transformer_lens.model_bridge.supported_architectures.llada2_moe module

LLaDA 2.0 MoE architecture adapter.

Ant Group’s LLaDA 2.x (LLaDA2MoeModelLM, remote code): masked block-diffusion language models on a DeepSeek-V3-style MoE decoder — fused query_key_value attention with full-width query/key layernorms, per-expert routed MLPs behind a bias-corrected router plus shared experts, and dense MLPs on the first first_k_dense_replace layers.

Attention is bidirectional (is_causal = False) and generation is block-diffusion sampling via the model’s own generate, reached through bridge.diffusion_generate; attention delegates to HF and the bridge’s autoregressive generation stays off. The fused QKV ships no HookedTransformer-format weight conversions, so LN folding is disabled.

The remote forward validates attention_mask strictly: it must be the 4D block-diffusion form (batch, 1, seq, seq) — full-ones for full bidirectional visibility — not a 2D padding mask. A forward pre-hook drops all-ones 2D masks (informationless) and rejects padded ones with a clear error instead of the remote validator’s opaque failure.

class transformer_lens.model_bridge.supported_architectures.llada2_moe.LLaDA2MoeArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for LLaDA2MoeModelLM models.

__init__(cfg: Any) → None

Initialize the LLaDA 2.0 MoE architecture adapter.

applicable_phases: list[int] = [1, 2, 3, 4]

component_mapping: ComponentMapping | None

native_sampler: str = 'generate'

native_sampler_kwargs(max_new_tokens: int, prompt_len: int) → dict

gen_length must cover whole blocks; block_length caps at the budget.

prepare_loading(model_name: str, model_kwargs: dict) → None

Restore the v4 ‘default’ rope init the remote code looks up (Dream shim).

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Delegated attention computes rotary inside HF; nothing to wire.

setup_hook_compatibility(bridge: Any) → None

Guard the remote forward against auto-passed 2D padding masks.

supports_fold_ln = False

supports_generation: bool = False

uses_split_attention: bool

weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None