transformer_lens.model_bridge.supported_architectures.dream module

Dream diffusion LM architecture adapter.

HKU-NLP’s Dream 7B (DreamModel, remote code; also Apple’s DiffuCoder): a discrete-diffusion text model initialized from Qwen2.5, so the module tree is exactly Qwen2 (biased q/k/v, gated SiLU MLP, RMS norms, shared model.rotary_emb, untied lm_head) — but attention is fully bidirectional (is_causal = False) and generation is iterative denoising via diffusion_generate, not autoregressive decoding.

Attention is therefore delegated to HF wholesale: the bridge’s reimplemented attention assumes causal masking. Q/K/V/O hooks fire on the wrapped projections; there is no reconstructed pattern hook.

The remote code targets transformers 4.46; v5 removed the “default” key from ROPE_INIT_FUNCTIONS, so prepare_loading re-registers it with the v4 semantics (plain inverse-frequency rope, attention factor 1.0).

class transformer_lens.model_bridge.supported_architectures.dream.DreamArchitectureAdapter(cfg: Any)

Bases: Qwen2ArchitectureAdapter

Architecture adapter for DreamModel diffusion LMs.

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

component_mapping: ComponentMapping | None

native_sampler: str = 'diffusion_generate'

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

Dream denoises a fixed-length canvas; one step per token is its default ratio.

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

Shim the remote code’s two transformers-v4 dependencies.

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

Delegated attention computes rotary inside HF; nothing to wire.

supports_generation: bool = False

uses_split_attention: bool

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