transformer_lens.model_bridge.supported_architectures.llada module

LLaDA architecture support for one masked-token transformer forward pass.

LLaDA is a masked discrete-diffusion model. Its transformer uses bidirectional self-attention and returns logits for every input position; the iterative denoising sampler lives outside the model. This adapter intentionally supports only the transformer forward pass. TransformerBridge’s autoregressive generation APIs are disabled for this architecture.

The remote LLaDALlamaBlock owns its Q/K/V, output, and gated-MLP projections directly instead of grouping them in attention and MLP modules. The adapter preserves the reviewed remote block forward and replaces only its attention(...) method with a hook-aware reconstruction. This keeps the native pre-norm residual order, dropout, RoPE implementation, and MLP math while exposing the standard attention scores and pattern hooks.

class transformer_lens.model_bridge.supported_architectures.llada.LLaDAArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Adapter for the dense LLaDAModelLM architecture.

Support is deliberately limited to the released dense LLaDA block contract: Llama-style blocks, RMSNorm, separate bias-free projections, RoPE, bidirectional attention, an untied LM head, and no KV cache. The external iterative denoising/remasking loop is not a TransformerBridge generation API. Loading the Hugging Face checkpoint requires the caller to opt in with trust_remote_code=True.

applicable_phases: list[int] = []

component_mapping: ComponentMapping | None

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

Disable remote branches incompatible with single-pass hook support.

prepare_model(hf_model: Any) → None

Keep the wrapper and underlying model on the no-cache path.

supports_causal_loss: bool = False

supports_generation: bool = False

supports_hf_output_attentions: bool = False

uses_split_attention: bool

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