File size: 854 Bytes
3cd1076 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | DATA_DIR="YOUR_DATA_DIR"
TEACHER_PATH="YOUR_FLM_CHECKPOINT_PATH"
python -u -m main \
loader.global_batch_size=128 \
loader.batch_size=16 \
loader.eval_batch_size=16 \
data=openwebtext-split \
data.cache_dir=$DATA_DIR \
wandb.project=owt_distill \
wandb.name=flm_distill \
model=small \
algo=fmlm_twomodel \
algo.teacher_path=$TEACHER_PATH \
trainer.max_steps=1000000 \
trainer.precision=bf16 \
trainer.val_check_interval=10000 \
model.length=1024 \
sampling.steps=[1,2,4,32] \
sampling.solver=euler \
optim.lr=3e-4 \
algo.double_temb=True \
algo.add_boundary=True \
+algo.boundary_prob=64 \
algo.bootstrap_ema=False \
algo.learnable_loss_weighting=True \
callbacks.checkpoint_every_n_steps.every_n_train_steps=20000 \
checkpointing.resume_from_ckpt=False \
|