ermine-olmoe

ERMINE-OLMoE: 1B-active MoE Masked Diffusion Language Models

This repository hosts intermediate training checkpoints (step 170,000) for two ~1B-active-parameter masked diffusion language models (MDMs) built on an OLMoE-style Mixture-of-Experts transformer backbone. The two checkpoints share the same architecture and data but differ in the schedule used to sample the diffusion timestep t during training.

Code: https://github.com/chen-hao-chao/olmoe-mdm (fork of AI2's OLMo)

Model description

• Objective: masked (absorbing-state) discrete diffusion language modeling, trained with a Monte-Carlo ELBO objective rather than next-token autoregression. The model has bidirectional attention (bidirectional: true) and is trained to denoise randomly masked token subsets.
• Backbone: OLMo-style transformer with Mixture-of-Experts FFN blocks (OLMoE-style), RoPE position embeddings, RMSNorm, and SwiGLU activations.
• Architecture:
• d_model: 2240
• n_layers: 18
• n_heads: 35
• MoE: 8 experts, top-2 routing per token (dropless, grouped-GEMM MLP impl)
• MLP hidden size (per expert): 2816
• Context length: 2048 tokens
• Vocab size: 50,280 (+ 1 reserved [MASK] token at id 50280; embedding table padded to 50,304)
• ~1.3B active parameters / ~3.3B total parameters (approximate, dense-equivalent — MoE total includes all 8 experts' weights, active counts top-2)
• Tokenizer: AI2's GPT-NeoX-based Dolma v1.5 tokenizer (allenai_gpt-neox-olmo-dolma-v1_5).

Training

• Data: tokenized shards of AI2's olmoe-mix-0924 pretraining mix (the same mix used for AI2's official OLMoE-1B-7B-0924 model), packed to 2048-token sequences.
• Optimizer: AdamW, lr 4e-4, weight decay 0.1, betas (0.9, 0.95), cosine schedule with 1,000-step warmup, bf16 mixed precision, FSDP full-shard.
• Batch size: global batch size 512 sequences × 2048 tokens ≈ 1.05M tokens/step.
• Checkpoint step: 170,000 / 520,000 planned total steps (~33% through training), corresponding to roughly 178B tokens seen (of a ~545B-token full run).
• Both runs use identical hyperparameters except for the diffusion-time sampling schedule noted above (mdm.gpu_allocation_path, mdm.remove_factor); see configs/MDM-1B-rl.yml and configs/MDM-1B-rl-ermine.yml for the exact configs used.

Checkpoint format

Each .tar archive contains a raw PyTorch Distributed Checkpoint (DCP) directory (step170000/), i.e. a set of __<rank>_<shard>.distcp shard files plus a .metadata index — the native sharded-checkpoint format saved by this codebase's FSDP trainer (sharded_checkpointer: olmo_core). This is not a Hugging Face transformers-ready checkpoint and cannot be loaded with from_pretrained.

To use it:

1. Clone https://github.com/chen-hao-chao/olmoe-mdm and install its requirements.
2. Download and extract the .tar into a directory, e.g.: bash hf download chen-hao-chao/ermine-olmoe olmoe-mdm-1b-ermine-step170000.tar --local-dir . tar -xf olmoe-mdm-1b-ermine-step170000.tar
3. Reconstruct a full (unsharded) state dict with the repo's checkpoint utilities, e.g. scripts/unshard.py, pointing at the extracted step170000/ directory (which also contains the training config.yaml needed to rebuild the model).
4. For zero-shot downstream evaluation (SciQ, PIQA, Social IQa, HellaSwag, ARC-Easy, ARC-Challenge) with the same MC-ELBO scoring used during training, see eval_mdm_downstream.py.

Note: the MoE blocks use a grouped_gemm kernel build that must include your GPU's SM architecture (see the repo's Dockerfile) — it will error ("Grouped GEMM execution not possible with HW") on unsupported GPUs.

Intended use

These are research checkpoints from an in-progress training run, released for reproducibility and analysis of masked diffusion language model training dynamics (in particular, the effect of non-uniform diffusion-time sampling). They are not instruction-tuned or safety-aligned, and are not intended for production or end-user-facing applications.

License

Apache 2.0 (inherited from the AI2 OLMo codebase this project is built on).