bridge.diffusion.conversion.llada15.llada15_bridge#

Megatron Bridge for LLaDA1.5 (GSAI-ML/LLaDA-*) masked-diffusion LLMs.

Converts between the trust_remote_code LLaDAModelLM HF class and a Megatron-Core GPTModel. The model is a dense LLaMA-style block with OLMo-style parameter naming and full RoPE.

Key mapping decisions, anchored to the reference implementation (modeling_llada.py):

  • QKV is separate (q_proj, k_proj, v_proj — see LLaDALlamaBlock.__init__), not fused. Use :class:QKVMapping, not

    class:

    ConcatenatedQKVMapping.

  • The SwiGLU MLP computes act(ff_proj(x)) * up_proj(x) (see LLaDALlamaBlock.forward), so ff_proj is the gate and up_proj is the up input for Megatron’s fused linear_fc1.

  • The LM head is model.transformer.ff_out (LLaDA1.5 has weight_tying: false).

  • Attention output projection is named attn_out, not o_proj.

  • Layer norms are attn_norm (pre-attention) and ff_norm (pre-MLP), fused into Megatron’s TE linear_qkv/linear_fc1 layer_norm_weight.

Module Contents#

Classes#

LLaDA15Bridge

HF LLaDAModelLM ↔ Megatron GPTModel bridge.

API#

class bridge.diffusion.conversion.llada15.llada15_bridge.LLaDA15Bridge#

Bases: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge

HF LLaDAModelLM ↔ Megatron GPTModel bridge.

provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.diffusion.models.llada15.llada15_provider.LLaDA15ModelProvider#
mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#