src/diffusion_lm/model.py

"""GPT-sized bidirectional Transformer used as a masked-token denoiser."""

from __future__ import annotations

import argparse
import math
from pathlib import Path

import torch
from torch import Tensor, nn
from torch.utils.checkpoint import checkpoint

from diffusion_lm.config import ModelConfig, load_config


class DiffusionTransformer(nn.Module):
    """A GPT-like Transformer with the causal mask deliberately removed.

    The network predicts clean tokens from an input containing absorbing mask
    tokens. Passing ``output_positions`` avoids materializing vocabulary logits
    for already-visible tokens during training.
    """

    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.config = config
        self.tokenizer_sha256: str | None = None
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.position_embedding = nn.Embedding(config.max_seq_len, config.d_model)
        self.embedding_dropout = nn.Dropout(config.dropout)

        if config.use_flex_attention:
            from diffusion_lm.flexattn import FlexEncoder

            self.transformer = FlexEncoder(
                d_model=config.d_model,
                n_heads=config.n_heads,
                d_ff=config.d_ff,
                dropout=config.dropout,
                n_layers=config.n_layers,
                activation_checkpointing=config.activation_checkpointing,
            )
        else:
            layer = nn.TransformerEncoderLayer(
                d_model=config.d_model,
                nhead=config.n_heads,
                dim_feedforward=config.d_ff,
                dropout=config.dropout,
                activation="gelu",
                batch_first=True,
                norm_first=True,
            )
            self.transformer = nn.TransformerEncoder(
                layer,
                num_layers=config.n_layers,
                norm=nn.LayerNorm(config.d_model),
                enable_nested_tensor=False,
            )
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        self.apply(self._init_weights)
        self._init_residual_outputs()
        if config.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight
        self.register_buffer(
            "_forbidden_output_token_ids",
            torch.tensor(config.forbidden_output_token_ids, dtype=torch.long),
            persistent=False,
        )

    @staticmethod
    def _init_weights(module: nn.Module) -> None:
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.LayerNorm):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

    def _init_residual_outputs(self) -> None:
        """Scale residual branch outputs as in GPT-2 for stable deep training."""

        if self.config.use_flex_attention:
            self.transformer.init_residual_outputs(self.config.n_layers)
            return
        residual_std = 0.02 / math.sqrt(2 * self.config.n_layers)
        for layer in self.transformer.layers:
            nn.init.normal_(layer.self_attn.out_proj.weight, mean=0.0, std=residual_std)
            nn.init.normal_(layer.linear2.weight, mean=0.0, std=residual_std)

    def _checkpointed_transformer(
        self,
        hidden: Tensor,
        padding_mask: Tensor | None,
        attn_mask: Tensor | None = None,
    ) -> Tensor:
        for layer in self.transformer.layers:

            def run_layer(layer_input: Tensor, *, current_layer: nn.Module = layer) -> Tensor:
                return current_layer(
                    layer_input, src_mask=attn_mask, src_key_padding_mask=padding_mask
                )

            hidden = checkpoint(run_layer, hidden, use_reentrant=False)

        if self.transformer.norm is not None:
            hidden = self.transformer.norm(hidden)
        return hidden

    def _expand_attn_mask(self, attn_mask: Tensor | None, input_ids: Tensor) -> Tensor | None:
        """Broadcast a per-sample boolean blocking mask across attention heads.

        Accepts ``[L, L]`` shared masks or ``[B, L, L]`` per-sample masks with
        ``True`` marking blocked key positions, matching the src_mask convention.
        """

        if attn_mask is None:
            return None
        batch_size, sequence_length = input_ids.shape
        if attn_mask.dtype != torch.bool:
            raise ValueError("attn_mask must be boolean with True marking blocked positions")
        if attn_mask.shape == (sequence_length, sequence_length):
            return attn_mask
        if attn_mask.shape != (batch_size, sequence_length, sequence_length):
            raise ValueError("attn_mask must have shape [L, L] or [batch, L, L]")
        return attn_mask.repeat_interleave(self.config.n_heads, dim=0)

    def encode(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        attn_mask: Tensor | None = None,
    ) -> Tensor:
        """Return contextual token states; ``attn_mask`` restricts attention topology."""

        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, sequence]")
        batch_size, sequence_length = input_ids.shape
        if sequence_length > self.config.max_seq_len:
            raise ValueError(
                f"sequence length {sequence_length} exceeds max_seq_len "
                f"{self.config.max_seq_len}"
            )
        if attention_mask is not None and attention_mask.shape != input_ids.shape:
            raise ValueError("attention_mask must match input_ids")

        positions = torch.arange(sequence_length, device=input_ids.device)
        hidden = self.token_embedding(input_ids) + self.position_embedding(positions)[None, :, :]
        hidden = self.embedding_dropout(hidden)

        # TransformerEncoder expects True for padding, the inverse of the common
        # attention-mask convention. src_mask is only supplied by region-aware callers.
        padding_mask = None if attention_mask is None else ~attention_mask.bool()

        if self.config.use_flex_attention:
            from diffusion_lm.flexattn import build_block_mask

            if attn_mask is not None and attn_mask.dtype != torch.bool:
                raise ValueError("attn_mask must be boolean with True marking blocked positions")
            block_mask = build_block_mask(
                attn_mask, padding_mask, batch_size, sequence_length, hidden.device
            )
            return self.transformer(hidden, block_mask)

        expanded_attn_mask = self._expand_attn_mask(attn_mask, input_ids)
        if (
            self.config.activation_checkpointing
            and self.training
            and torch.is_grad_enabled()
        ):
            return self._checkpointed_transformer(hidden, padding_mask, expanded_attn_mask)
        return self.transformer(
            hidden, mask=expanded_attn_mask, src_key_padding_mask=padding_mask
        )

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor | None = None,
        output_positions: Tensor | None = None,
        attn_mask: Tensor | None = None,
    ) -> Tensor:
        """Predict vocabulary logits for all tokens or selected positions only."""

        hidden = self.encode(input_ids, attention_mask=attention_mask, attn_mask=attn_mask)
        if output_positions is not None:
            if output_positions.shape != input_ids.shape:
                raise ValueError("output_positions must match input_ids")
            hidden = hidden[output_positions.bool()]

        logits = self.lm_head(hidden)
        # Corruption/control tokens are never valid clean-token predictions. EOS
        # deliberately remains available so generation can terminate naturally.
        if self._forbidden_output_token_ids.numel():
            logits.index_fill_(
                -1,
                self._forbidden_output_token_ids,
                torch.finfo(logits.dtype).min,
            )
        return logits

    @property
    def num_parameters(self) -> int:
        """Count unique trainable parameters (shared embeddings count once)."""

        return sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)


def build_denoiser(
    config: ModelConfig,
    *,
    load_pretrained: bool = True,
    dtype: torch.dtype | None = None,
) -> nn.Module:
    """Construct the denoiser a config describes: project transformer or pretrained backbone.

    ``load_pretrained=False`` builds the architecture only, for callers that immediately
    restore weights from a project checkpoint.
    """

    if config.backbone == "hf-qwen3":
        from diffusion_lm.hf_bridge import Qwen3Denoiser

        return Qwen3Denoiser(config, load_pretrained=load_pretrained, dtype=dtype)
    return DiffusionTransformer(config)


def format_parameter_count(count: int) -> str:
    if count >= 1_000_000:
        return f"{count / 1_000_000:.2f}M"
    if count >= 1_000:
        return f"{count / 1_000:.2f}K"
    return str(count)


def main() -> None:
    parser = argparse.ArgumentParser(description="Report the exact model parameter count")
    parser.add_argument("--config", type=Path, required=True, help="experiment YAML")
    args = parser.parse_args()

    config = load_config(args.config)
    # Parameter inspection should not allocate four gigabytes for the 1B preset.
    with torch.device("meta"):
        model = DiffusionTransformer(config.model)
    print(f"parameters: {model.num_parameters:,} ({format_parameter_count(model.num_parameters)})")


if __name__ == "__main__":
    main()