News

[2025-10-27] We release the codebase, all training checkpoints, and logs. The codebase is highly optimized and is industry-level in terms scalability and efficiency.

[2025-10-03] The blog is out! Check it out here!

Code

The codebase is released here. It is a highly-optimized codebase for any-scale DLMs training backend with Megatron-LM.

The full MoE implementation is not yet released. We plan to release it after the main training is done.

Resources

We opensource all model checkpoints and training logs mentioned in the paper. All of them can be downloaded at https://huggingface.co/collections/jinjieni/mdga.

The easiest way to download a folder is using this script (setup the variables properly):

python utils/hf_download_folder.py

Alternatively, you can also use wget to directly download individual files from the folder, e.g.:

wget https://huggingface.co/datasets/MDGA-1/openmoe2_logs/blob/main/dense_vs_moe/dense_100b_1e_1b7_difflm/tensorboard/events.out.tfevents.1755443508.0648415733

We link the related resources below:

• Parameter–Compute Trade-off
  • [ckpt][log] dense 1.7b
  • [ckpt][log] dense 8b
  • [ckpt][log] moe-8b1.7a-tc
  • [ckpt][log] moe-8b1.7a-ec
• Diffusion + MoE is a Double Win
  • [ckpt][log] AR dense 1b
  • [ckpt][log] DLM dense 1b
  • [ckpt][log] AR dense 8b
  • [ckpt][log] DLM dense 8b
  • [ckpt][log] AR MoE 8b1a
  • [ckpt][log] DLM MoE 8b1a
• Token-choice vs. Expert-choice
  • [ckpt][log] tc
  • [ckpt][log] ec
• Token-Wise Load-Balancing
  • [ckpt][log] with
  • [ckpt][log] without
• with and without shared experts
  • [ckpt][log] with
  • [ckpt][log] without
• scaling the expert granularities
  • [ckpt][log] 16
  • [ckpt][log] 64
  • [ckpt][log] 256
• scratch vs. upcycling
  • [ckpt][log] scratch
  • [ckpt][log] upcycling
• Skip the First 2 MoE Layers
  • [ckpt][log] skip
  • [ckpt][log] keep
• with and without scaling factors
  • [ckpt][log] tc with
  • [ckpt][log] tc without
  • [ckpt][log] ec with
  • [ckpt][log] ec without
• Batch vs. Sequence Level Expert Choice
  • [ckpt][log] batch
  • [ckpt][log] seq
• Softmax vs. Sigmoid
  • [ckpt][log] softmax
  • [ckpt][log] sigmoid

You can refer to this script to inference with the huggingface checkpoints. Due to the large amount, most small checkpoints above are still in megatron formats. You may refer to this script to convert them (need to tweak the conversion scripts).

Todo List

This is an on-going project! We will tick the below todo list one-by-one.

• Architectural Design Choices
• Scaled-up Pre-training
• Post-training
• Routing & Other Analysis
• Full Paper
• Code & Checkpoint Open-sourcing

Citation

@misc{ni2025openmoe2,
  title={OpenMoE 2: Sparse Diffusion Language Models},
  author={Ni, Jinjie and team},
  year={2025}
  howpublished={\url{https://github.com/JinjieNi/OpenMoE2}},
}