transformer_lens.model_bridge.supported_architectures.bd3lm module

BD3LM (Block Diffusion Language Model) architecture adapter.

class transformer_lens.model_bridge.supported_architectures.bd3lm.BD3LMArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for BD3LM (Block Diffusion LM, ICLR 2025).

BD3LM uses adaLN conditioning on diffusion timesteps, a custom Rotary embedding, joint QKV projections, and non-causal block-diffusion masking. Because adaLN modulation varies per-timestep, it cannot be folded into weights — the adapter uses DelegatedAttentionBlockBridge to delegate each DDiTBlock.forward() wholesale to the original HF module. Hooks fire at block boundaries and on mapped submodules.

applicable_phases: list[int] = [1, 2, 3]
component_mapping: ComponentMapping | None
convert_weights() → dict[str, Tensor]

Return empty dict — delegation means no weight rearrangement.

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

Patch BD3LM dynamic class before from_pretrained runs.

Modeling code has a custom __getattr__ that fails to delegate back to PreTrainedModel, raising AttributeError on all_tied_weights_keys.

prepare_model(hf_model: Any) → None

Patch BD3LM quirks that prevent standard bridge construction.

Three issues must be fixed before the bridge can wrap the model:

  1. vocab_embed is an nn.Parameter, not nn.Embedding, so it lacks a .weight attribute that EmbeddingBridge expects.

  2. The flex attention backend crashes on CPU; fall back to sdpa and regenerate block_diff_mask for the new backend.

  3. The HF forward() does not accept output_attentions and other kwargs the bridge unconditionally injects; patch at runtime because no other hook point allows filtering them before HF’s forward call.

supports_generation: bool = False
uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None