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 \