Is your feature request related to a problem? Please describe.
Megatron Core supports autoregressive (GPT), masked language models (BERT), and state-space models (Mamba), but lacks support for discrete diffusion language models.
Recent models like MDLM, SEDD, and Meta's LLaDA (8B) demonstrate that discrete diffusion can match autoregressive quality while offering unique advantages:
• Parallel token generation (faster inference)
• Natural support for infilling and editing
• Better controllability and guidance
Existing open-source implementations lack Megatron's distributed training, making it difficult to scale.
Describe the solution you'd like
Support for training discrete diffusion language models at scale along with training recipes.
Describe alternatives you've considered
MDLM/SEDD repos exist but lack distributed training support for scaling
Additional context
• LLaDA https://arxiv.org/abs/2502.09992 — Meta's 8B diffusion LM
Is your feature request related to a problem? Please describe.
Megatron Core supports autoregressive (GPT), masked language models (BERT), and state-space models (Mamba), but lacks support for discrete diffusion language models.
Recent models like MDLM, SEDD, and Meta's LLaDA (8B) demonstrate that discrete diffusion can match autoregressive quality while offering unique advantages:
• Parallel token generation (faster inference)
• Natural support for infilling and editing
• Better controllability and guidance
Existing open-source implementations lack Megatron's distributed training, making it difficult to scale.
Describe the solution you'd like
Support for training discrete diffusion language models at scale along with training recipes.
Describe alternatives you've considered
MDLM/SEDD repos exist but lack distributed training support for scaling
Additional context
• LLaDA https://arxiv.org/abs/2502.09992 — Meta's 8B diffusion LM