Training

February 1, 2026 ยท View on GitHub

Train on OWT

HF_DATASETS_CACHE=".cache/datasets";
METHODS="algo=xdlm algo.k1=0.1 wandb.name=xdlm-k01-1M-owt";
OTHERS="";
# OTHERS="trainer.log_every_n_steps=1 training.finetune_path=tmp.ckpst";
python -u -m main loader.batch_size=16 loader.eval_batch_size=16 model=small data=openwebtext-split model.length=1024  +wandb.offline=True data.cache_dir=$HF_DATASETS_CACHE ${METHODS} ${OTHERS};

Methods for models

  • MDLM: algo=mdlm wandb.name=mdlm-1M-owt
  • GIDD: algo=gidd wandb.name=gidd-1M-owt
  • XDLM-k0001: algo=xdlm algo.k1=0.001 wandb.name=xdlm-k0001-1M-owt
  • XDLM-k01: algo=xdlm algo.k1=0.1 wandb.name=xdlm-k01-1M-owt
  • XDLM-k05: algo=xdlm algo.k1=0.5 wandb.name=xdlm-k05-1M-owt
  • XDLM-k09: algo=xdlm algo.k1=0.9 wandb.name=xdlm-k09-1M-owt
  • UDLM: algo=udlm wandb.name=udlm-1M-owt

Train on LM1B

HF_DATASETS_CACHE=".cache/datasets";
METHODS="algo=xdlm algo.k1=0.1 wandb.name=xdlm-k01-1M-lm1b-wrap";
OTHERS="";
# OTHERS="trainer.log_every_n_steps=1 checkpointing.resume_ckpt_path=tmp.ckpt";
python -u -m main loader.batch_size=64 loader.eval_batch_size=32 model=small data=lm1b-wrap model.length=128 +wandb.offline=True data.cache_dir=$HF_DATASETS_CACHE ${METHODS} ${OTHERS};

Train on ImageNet-1K

IMNET="/htmp/ImageNet/ILSVRC/Data/CLS-LOC";
TOKENIZER_CKPT="$HF_HOME/others/LLamaGen/vq_ds16_c2i.pt";
METHODS="algo=xdlm algo.k1=0.1 wandb.name=xdlm-k01-imnet-16k-cond";
OTHERS="";
python -u -m main is_vision=True loader.batch_size=64 loader.eval_batch_size=64 model=small model.length=256 data=imagenet data.train=$IMNET data.valid=$IMNET data.tokenizer_name_or_path=imagenet_16384 data.tokenizer_checkpoint=$TOKENIZER_CKPT trainer.max_steps=500000 +wandb.offline=True +trainer.check_val_every_n_epoch=null training.guidance=True ${METHODS} ${OTHERS};

Train on CIFAR10

export CIFAR10="${HF_HOME}/datasets/cifar10";
METHODS="algo=xdlm algo.k1=0.1 wandb.name=xdlm-k01-cifar10-cond";
OTHERS="";
python -u -m main is_vision=True loader.batch_size=64 loader.eval_batch_size=64 model=unet model.length=3072 data=cifar10 data.train=$CIFAR10 data.valid=$CIFAR10 algo.backbone=unet optim.lr=2e-4 lr_scheduler.num_warmup_steps=5000 trainer.max_steps=300000 +wandb.offline=True +trainer.check_val_every_n_epoch=null training.guidance=True ${METHODS} ${OTHERS};