bridge.diffusion.models.llada15.llada15_provider

LLaDA15ModelProvider: GPTModelProvider with bidirectional core attention.

Module Contents

Classes

LLaDA15ModelProvider

GPTModelProvider for LLaDA1.5 masked-diffusion dense models.

Functions

_llada15_layer_spec

Build a TE GPT layer spec and swap in LLaDA15TEDotProductAttention.

API

bridge.diffusion.models.llada15.llada15_provider._llada15_layer_spec(

config: LLaDA15ModelProvider,

) → megatron.core.transformer.ModuleSpec

Build a TE GPT layer spec and swap in LLaDA15TEDotProductAttention.

Injecting the custom core attention is the single change required versus a vanilla TE GPT layer: it forces bidirectional attention (no_mask) so the model matches LLaDA1.5’s reference implementation, which uses a zero attention bias at every forward pass.

class bridge.diffusion.models.llada15.llada15_provider.LLaDA15ModelProvider

Bases: megatron.bridge.models.gpt_provider.GPTModelProvider

GPTModelProvider for LLaDA1.5 masked-diffusion dense models.

Differences vs a vanilla Llama GPTModelProvider:

• transformer_layer_spec is overridden so every self-attention layer uses LLaDA15TEDotProductAttention, which forces bidirectional attention (the reference implementation uses zero attention bias everywhere — see modeling_llada.py:get_bidirectional_attention_bias).

• share_embeddings_and_output_weights defaults to False because LLaDA1.5 has weight_tying: false in its HF config — the LM head lives at model.transformer.ff_out.

Full RoPE is handled by Megatron’s standard rotary path (the HF reference rotates the full head_dim with a standard rotate-half pattern), so position_embedding_type stays at its default "rope" value.

transformer_layer_spec: Union[megatron.core.transformer.ModuleSpec, Callable]

None

share_embeddings_and_output_weights: bool

False

position_embedding_type: str

‘rope’