wutaiqiang commented 11 months ago

I try to distill gpt2-1.5B -> gpt2-120M As I use 4 A100, so I change the GPUS_PER_NODE to ${3-4}

Batch size remains the same

wutaiqiang commented 11 months ago

Then I get the log:

wutaiqiang commented 11 months ago

max roughL is 24.353, and when I evaluate the model on dolly, I get 22.40, which is far less than 24.6

wutaiqiang commented 11 months ago

@t1101675 could you please help to figure out where I am wrong? thanks

the cmd is:

! /bin/bash

MASTER_ADDR=localhost MASTER_PORT=${2-2012} NNODES=1 NODE_RANK=0 GPUS_PER_NODE=${3-4}

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE \ --nnodes $NNODES \ --node_rank $NODE_RANK \ --master_addr $MASTER_ADDR \ --master_port $MASTER_PORT"

model

BASE_PATH=${1-"xxxx"} CKPT_NAME="base-init" CKPT="${BASE_PATH}/results/gpt2/train/minillm_init/gpt2-base" TEACHER_CKPT_NAME="xlarge-sft" TEACHER_CKPT="${BASE_PATH}/results/gpt2/train/sft/gpt2-xlarge/"

data

PROMPT_DATA_DIR="${BASE_PATH}/processed_data/dolly/prompt/gpt2/" LM_DATA_DIR="${BASE_PATH}/processed_data/openwebtext/gpt2/512/10M/"

runtime

SAVE_PATH="${BASE_PATH}/results/gpt2/train/minillm/"

hp

GRAD_ACC=1 BATCH_SIZE=4 CHUNK_SIZE=16

OPTS=""

model

OPTS+=" --base-path ${BASE_PATH}" OPTS+=" --model-path ${CKPT}" OPTS+=" --teacher-model-path ${TEACHER_CKPT}" OPTS+=" --ckpt-name ${CKPT_NAME}" OPTS+=" --teacher-ckpt-name ${TEACHER_CKPT_NAME}" OPTS+=" --n-gpu ${GPUS_PER_NODE}" OPTS+=" --teacher-model-fp16"

OPTS+=" --gradient-checkpointing"

data

OPTS+=" --prompt-data-dir ${PROMPT_DATA_DIR}" OPTS+=" --lm-data-dir ${LM_DATA_DIR}" OPTS+=" --dev-num 1000" OPTS+=" --num-workers 0"

hp

OPTS+=" --epochs 10" OPTS+=" --total-iters 5000" OPTS+=" --kd-ratio 0.5" OPTS+=" --batch-size ${BATCH_SIZE}" OPTS+=" --lr 5e-6" OPTS+=" --lr-min 5e-6" OPTS+=" --gradient-accumulation-steps ${GRAD_ACC}" OPTS+=" --max-length 512" OPTS+=" --max-prompt-length 256" OPTS+=" --warmup-iters 100"

runtime

OPTS+=" --save ${SAVE_PATH}" OPTS+=" --seed 10" OPTS+=" --seed-ppo 42" OPTS+=" --seed-lm 7" OPTS+=" --save-interval 100" OPTS+=" --eval-interval 20" OPTS+=" --log-interval 1" OPTS+=" --mid-log-num 1"

ppo

OPTS+=" --type minillm" OPTS+=" --ppo-epochs 4" OPTS+=" --num-rollouts 16" OPTS+=" --chunk-size ${CHUNK_SIZE}"

minillm

OPTS+=" --length-norm" OPTS+=" --single-step-reg" OPTS+=" --teacher-mixed-alpha 0.2"

reward

OPTS+=" --reward-scaling 0.5" OPTS+=" --cliprange-reward 100"

gen

OPTS+=" --do-sample" OPTS+=" --top-k 0" OPTS+=" --top-p 1.0" OPTS+=" --temperature 1.0"

deepspeed

OPTS+=" --deepspeed" OPTS+=" --deepspeed_config ${BASE_PATH}/configs/deepspeed/ds_config.json"

export NCCL_DEBUG="" export WANDB_DISABLED=True export TF_CPP_MIN_LOG_LEVEL=3 export PYTHONPATH=${BASE_PATH} CMD="torchrun ${DISTRIBUTED_ARGS} ${BASE_PATH}/train_minillm.py ${OPTS} $@"

echo ${CMD} echo "PYTHONPATH=${PYTHONPATH}" mkdir -p ${SAVE_PATH} ${CMD}

wutaiqiang commented 11 months ago

Change the batch size to 16, so that 4 GPU * 16= 64, then we get 19.82 on the dolly, a worse result.

t1101675 commented 11 months ago

The --epochs should be larger (> 300) and you will get the total global iters = 5000 in the log
The --num-rollouts should be larger (num-rollouts * num-gpus should be 256)
You can also increase --chunk-size for more efficient training.

We have updated our code for the convenience to set these hyper-parameters. You can ignore 1 & 2 if you use current scripts.

wutaiqiang commented 11 months ago

thanks Let me try the new hyper-parameter

wutaiqiang commented 11 months ago

thanks for your kind help. On dolly, mine ~= paper. On UnNI and selfInst and S-NI mine < paper. On Vicuna mine > paper. I think the seeds may cause it.

Thanks again. BTW, In pt_loss, it is

(1-self.args.kd_ratio) lm_loss + self.args.kd_ratio distil_loss

and the distil_loss is:

teacher_probs = F.softmax(teacher_logits, dim=-1, dtype=torch.float32) inf_mask = torch.isinf(logits) logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32) prod_probs = torch.masked_fill(teacher_probs logprobs, inf_mask, 0) x = torch.sum(prod_probs, dim=-1).view(-1) distil_loss = -torch.sum(x loss_mask.view(-1), dim=0) / torch.sum(loss_mask.view(-1), dim=0) #! div 0 风险

This distil_loss is forward KL. Meanwhile, I noticed that you use the get_rev_kl function in sampler.py, which is reverse KL. Could you please help figure out why two distill losses there?

thanks very much!

t1101675 commented 11 months ago

The forward KL loss works as a regularization to prevent the model from collapsing to a single mode when using reverse KL. The effect of this loss is controlled by args.kd_ratio and is optional because we find that the result does not change much with different args.kd_ratio (kd_ratio=0 means only using lm_loss).

AInkCode commented 11 months ago

@wutaiqiang @t1101675 hi, bro! In which file can I set the parameters of the distillation model, such as distill gpt2-1.5B -> gpt2-120M? Thanks？

wutaiqiang commented 11 months ago

https://github.com/microsoft/LMOps/blob/38f67def245e61cdbcf06c61ad98c030c6f2f505/minillm/scripts/gpt2/minillm/train_base_xl.sh

@AInkCode

microsoft / LMOps

[minillm] hard to reproduce the result #118

! /bin/bash

model

data

runtime

hp

model

OPTS+=" --gradient-checkpointing"

data

hp

runtime

ppo

minillm

reward

gen

deepspeed