chen-hao-chao/ermine-megatron

Ermine: Pretrained Megatron-LM Diffusion Language Models

This repository hosts pretrained checkpoints for a suite of diffusion language models and autoregressive baselines across six model families, trained using Megatron-LM and used to study neural differential scaling laws.

Model Families

Family	Directory	Description
ARM	arm/	Autoregressive language model baseline
MDM	mdm/	Masked diffusion language model
MDM-MoE	mdm-moe/	MDM with Mixture-of-Experts FFN
MDM-MoE (Ermine)	mdm-moe-ermine/	MDM-MoE trained with the Ermine compute-optimal scaling recipe
UDM	udm/	Uniform-state diffusion language model
UDM-MoE	udm-moe/	UDM with Mixture-of-Experts FFN

Checkpoints

ARM (`arm/`)

File	Params
arm_param_14M_model_only.tar	14M
arm_param_25M_model_only.tar	25M
arm_param_36M_model_only.tar	36M
arm_param_48M_model_only.tar	48M
arm_param_63M_model_only.tar	63M
arm_param_77M_model_only.tar	77M
arm_param_106M_model_only.tar	106M
arm_param_153M_model_only.tar	153M
arm_param_199M_model_only.tar	199M
arm_param_249M_model_only.tar	249M
arm_param_352M_model_only.tar	352M
arm_param_411M_model_only.tar	411M
arm_param_919M_model_only.tar	919M

MDM (`mdm/`)

File	Params
mdm_param_48M_model_only.tar	48M
mdm_param_63M_model_only.tar	63M
mdm_param_77M_model_only.tar	77M
mdm_param_106M_model_only.tar	106M
mdm_param_153M_model_only.tar	153M
mdm_param_199M_model_only.tar	199M
mdm_param_249M_model_only.tar	249M
mdm_param_352M_model_only.tar	352M
mdm_param_566M_model_only.tar	566M
mdm_param_766M_model_only.tar	766M
mdm_param_919M_model_only.tar	919M
mdm_param_1104M_model_only.tar	1.1B
mdm_param_1529M_model_only.tar	1.5B
mdm_param_2350M_model_only.tar	2.4B

MDM-MoE (`mdm-moe/`)

File	Params
mdm_moe_param_14M_model_only.tar	14M
mdm_moe_param_25M_model_only.tar	25M
mdm_moe_param_36M_model_only.tar	36M
mdm_moe_param_48M_model_only.tar	48M
mdm_moe_param_63M_model_only.tar	63M
mdm_moe_param_77M_model_only.tar	77M
mdm_moe_param_106M_model_only.tar	106M
mdm_moe_param_153M_model_only.tar	153M
mdm_moe_param_199M_model_only.tar	199M
mdm_moe_param_249M_model_only.tar	249M
mdm_moe_param_411M_model_only.tar	411M
mdm_moe_param_566M_model_only.tar	566M
mdm_moe_param_766M_model_only.tar	766M

MDM-MoE (Ermine) (`mdm-moe-ermine/`)

File	Params
mdm_moe_param_14M_model_only_ermine.tar	14M
mdm_moe_param_25M_model_only_ermine.tar	25M
mdm_moe_param_36M_model_only_ermine.tar	36M
mdm_moe_param_48M_model_only_ermine.tar	48M
mdm_moe_param_63M_model_only_ermine.tar	63M
mdm_moe_param_77M_model_only_ermine.tar	77M
mdm_moe_param_106M_model_only_ermine.tar	106M
mdm_moe_param_153M_model_only_ermine.tar	153M
mdm_moe_param_199M_model_only_ermine.tar	199M
mdm_moe_param_249M_model_only_ermine.tar	249M
mdm_moe_param_352M_model_only_ermine.tar.tar	352M
mdm_moe_param_411M_model_only_ermine.tar.tar	411M
mdm_moe_param_566M_model_only_ermine.tar.tar	566M
mdm_moe_param_766M_model_only_ermine.tar.tar	766M

UDM (`udm/`)

File	Params
udm_param_7M_model_only.tar	7M
udm_param_10M_model_only.tar	10M
udm_param_14M_model_only.tar	14M
udm_param_25M_model_only.tar	25M
udm_param_36M_model_only.tar	36M
udm_param_48M_model_only.tar	48M
udm_param_63M_model_only.tar	63M
udm_param_77M_model_only.tar	77M
udm_param_106M_model_only.tar	106M
udm_param_153M_model_only.tar	153M
udm_param_199M_model_only.tar	199M
udm_param_249M_model_only.tar	249M
udm_param_352M_model_only.tar	352M
udm_param_411M_model_only.tar	411M
udm_param_566M_model_only.tar	566M
udm_param_766M_model_only.tar	766M

UDM-MoE (`udm-moe/`)

File	Params
udm_moe_param_5M_model_only.tar	5M
udm_moe_param_7M_model_only.tar	7M
udm_moe_param_10M_model_only.tar	10M
udm_moe_param_14M_model_only.tar	14M
udm_moe_param_25M_model_only.tar	25M
udm_moe_param_36M_model_only.tar	36M
udm_moe_param_48M_model_only.tar	48M
udm_moe_param_63M_model_only.tar	63M
udm_moe_param_77M_model_only.tar	77M
udm_moe_param_106M_model_only.tar	106M
udm_moe_param_153M_model_only.tar	153M
udm_moe_param_199M_model_only.tar	199M
udm_moe_param_249M_model_only.tar	249M
udm_moe_param_352M_model_only.tar	352M
udm_moe_param_411M_model_only.tar	411M
udm_moe_param_566M_model_only.tar	566M
udm_moe_param_766M_model_only.tar	766M

Training Curves

Full training curves and evaluation logs are available on Weights & Biases: lance_chao/ermine-all-runs

Model Description

Architecture: Megatron-LM transformer with optional Mixture-of-Experts FFN blocks, RoPE position embeddings, RMSNorm, and SwiGLU activations.

Tokenizer: GPT-2 BPE

Data: chen-hao-chao/ermine-c4-dataset

Checkpoint Format

Each .tar file contains a torch_dist Megatron-LM distributed checkpoint directory.

Note: These are not Hugging Face transformers-ready checkpoints and cannot be loaded with from_pretrained.

License

Apache 2.0 (inherited from the NVIDIA/Megatron-LM codebase this project is built on)