michaelchenkj/mdlm-dclm-530m-53b

MDLM-530M on DCLM 53B

Masked diffusion language model with the DataDecide OLMo-1 530M block (bidirectional attention + mask token), trained on the same 53B unique DCLM-baseline tokens as the autoregressive companion michaelchenkj/datadecide-dclm-530m-53b.

This repo stores the raw training checkpoints (model + optimizer in each .pt), EasyBench evals (metrics + per-example generations), logs, and training code.

Recipe

• Architecture: d=1344, 16L, 16H, SwiGLU, RoPE, RMSNorm, seq 2048, vocab 50280 / embedding 50304, mask id 50280
• Loss: absorbing-state mask (SMDM/MDLM): t ~ U[0,1], p=(1-ε)t+ε, CE on masked tokens / p, mean over B×T
• Data: one pass over DataDecide DCLM-baseline shards
• Steps: 57,786; global batch 448 (112 seq/GPU × 4); LR 2.77e-3 cosine, 578 warmup, α=0.01; AdamW 0.9/0.95, wd 0.05
• Hardware: 4×H100
• Seed: 6198

Layout

• checkpoints/stepN.pt — every 1,000 steps; checkpoints/step57786.pt is the final (includes optimizer). checkpoints/config.json has run hyperparameters
• evals/ — per-step metrics.json / original.json and per-example .jsonl generations
• logs/ — training, nohup, and eval-worker logs
• scripts/ — train.py, launch.sh, and eval launchers
• src/mdlm/ — model, diffusion loss, data loader

Each stepN.pt (except the 6.7G final) is a torch.save dict with step, model, optim, and RNG state.

Load a checkpoint

from huggingface_hub import hf_hub_download
import torch

path = hf_hub_download("michaelchenkj/mdlm-dclm-530m-53b", "checkpoints/step57786.pt")
ckpt = torch.load(path, map_location="cpu")
state = ckpt["model"]  # or the whole dict to resume

Training data is not stored here. Use the same DCLM shards as the DataDecide 530M run (allenai/DataDecide-data-recipes).

Train

CUDA_VISIBLE_DEVICES=4,5,6,7 ./scripts/launch.sh --max-steps 2 --skip-eval   # smoke
CUDA_VISIBLE_DEVICES=4,5,6,7 ./scripts/launch.sh