File size: 946 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 30 31 32 33 34 35 | #!/bin/bash
TEACHER_PATH="YOUR_FLM_CHECKPOINT_PATH"
DATA_CACHE_DIR="YOUR_DATA_DIR"
python -u -m main \
loader.global_batch_size=512 \
loader.batch_size=128 \
loader.eval_batch_size=128 \
data=lm1b-wrap \
data.cache_dir=${DATA_CACHE_DIR} \
model=small \
model.length=128 \
algo=fmlm \
algo.double_temb=True \
algo.learnable_loss_weighting=False \
algo.distillation_method=PSD \
algo.use_mse_loss_psd=False \
algo.diagonal_fraction=0.5 \
algo.add_boundary=fixed \
algo.boundary_prob=32 \
algo.offdiagonal_sampling=uniform_diff \
algo.use_ema_for_psd_target=False \
algo.teacher_path=${TEACHER_PATH} \
algo.initialize_student_from_teacher=True \
sampling.steps=[1,2,4,8,16,32,64,128] \
trainer.max_steps=1000000 \
trainer.precision=bf16 \
trainer.val_check_interval=10000 \
trainer.limit_val_batches=10 \
optim.lr=3e-4 \
optim.beta2=0.95 \
wandb.project=lm1b_full \
wandb.name=lm1b_fmlm_PSD
|