training_code/nla/models.py

"""NLACriticModel: truncated transformer + vector value head.

Architecture:
  - First K transformer layers only (K = extraction layer, set via config override)
  - No final layernorm — raw residual stream goes to head
  - Linear(d_model, d_model) head, no bias (the paper says "affine"; we drop
    the bias deliberately — predictions are L2-normalized in the loss anyway,
    and identity-init is cleaner without an offset term)
  - Forward returns .values at every position; training extracts at the LAST
    real token of the critic prompt (suffix-anchored) for the MSE

Layer truncation: set config.num_hidden_layers BEFORE from_pretrained so the
weight loader only reads K layers. Do NOT slice nn.ModuleList post-hoc — breaks
FSDP sharding assumptions.
"""

import json
import os
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn as nn
from huggingface_hub import snapshot_download
from safetensors import safe_open
from safetensors.torch import load_file, save_file
from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel

from nla.arch_adapters import resolve_text_config, resolve_text_model


# Embedding weight key suffixes across architectures.
# Llama/Qwen/Mistral/Gemma: model.embed_tokens.weight
# GPT-2: transformer.wte.weight
# Falcon: transformer.word_embeddings.weight
_EMBED_KEY_SUFFIXES = ("embed_tokens.weight", "wte.weight", "word_embeddings.weight")


@dataclass
class NLACriticOutput:
    values: torch.Tensor  # [B, T, d_model] — value-head output at every position
    backbone_last_hidden: torch.Tensor  # [B, T, d_model] — pre-value-head (for norm tracking)


def _truncate_config_layers(config, num_layers: int) -> None:
    """Set num_hidden_layers AND truncate per-layer arrays to match.

    transformers >=4.50 validates len(layer_types) == num_hidden_layers at
    config init (configuration_utils.py:layer_type_validation). Qwen2/Llama3
    configs carry per-layer arrays that must be sliced consistently.
    """
    config.num_hidden_layers = num_layers
    for attr in ("layer_types", "sliding_window_pattern", "no_rope_layers"):
        v = getattr(config, attr, None)
        if isinstance(v, (list, tuple)) and len(v) > num_layers:
            setattr(config, attr, type(v)(v[:num_layers]))


def _inner_transformer(backbone: PreTrainedModel) -> nn.Module:
    """Get the inner transformer module (the part with .layers + .norm).

    Qwen/Llama/Mistral/Gemma (post-resolve_text_model → CausalLM wrapper): backbone.model
    GPT-2/Falcon: backbone.transformer
    """
    if hasattr(backbone, "model"):
        return backbone.model
    if hasattr(backbone, "transformer"):
        return backbone.transformer
    raise AssertionError(
        f"{type(backbone).__name__} has neither .model nor .transformer — "
        f"add the attribute name here if supporting a new arch"
    )


class NLACriticModel(PreTrainedModel):
    """Wraps an HF causal LM backbone with layer truncation + vector value head.

    Delegates everything structural (save/load/fsdp-wrapping) to the backbone's
    PreTrainedModel machinery. Only the forward path + head are NLA-specific.
    """

    # We never run attention ourselves; backbone does and validates its own
    # attn_implementation. Advertise support so HF doesn't refuse construction.
    _supports_sdpa = True
    _supports_flash_attn_2 = True
    _supports_flex_attn = True

    def __init__(self, config, backbone: PreTrainedModel):
        super().__init__(config)
        self.backbone = backbone
        self.value_head = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        # FSDP's apply_fsdp2 reads this to decide which modules to wrap.
        # Instance attr (not class attr) — two NLACriticModels with different
        # backbones in one process would clobber each other on class attr.
        self._no_split_modules = backbone._no_split_modules

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *, nla_num_layers: int | None = None, **kwargs):
        """Load an NLACriticModel from an HF checkpoint.

        Normal case: checkpoint was produced by a previous NLA training run →
        config.json already has the truncated num_hidden_layers. Just load.

        Bootstrapping case (fresh truncation from a full base model): pass
        nla_num_layers = the datagen extraction layer_index. Truncation keeps
        blocks 0..layer_index INCLUSIVE — we need the output OF block K, which
        means we need block K to exist, so num_hidden_layers = K+1.

        Indexing convention (matches datagen/extractors.py):
          - datagen `layer_index=K` hooks `model.model.layers[K]` →
            captures its output (= HF's `hidden_states[K+1]`, post-residual-add,
            the stream entering block K+1).
          - critic `last_hidden_state` (with final-LN → Identity) is the output
            of the LAST block.
          - So last block must be block K → num_hidden_layers = K+1.
        """
        # Callers may pass trust_remote_code + attn_implementation via kwargs.
        # Default trust_remote_code to True if caller doesn't specify.
        kwargs.setdefault("trust_remote_code", True)
        config = AutoConfig.from_pretrained(
            pretrained_model_name_or_path, trust_remote_code=kwargs["trust_remote_code"]
        )
        # For multimodal wrappers (Gemma-3), truncate the nested text config —
        # the wrapper config references it by identity so the truncation is
        # visible to from_pretrained. For plain text models, text_config IS
        # config (see arch_adapters).
        text_config = resolve_text_config(config)
        if nla_num_layers is not None:
            needed = nla_num_layers + 1
            assert needed <= text_config.num_hidden_layers, (
                f"nla_num_layers={nla_num_layers} needs blocks 0..{nla_num_layers} "
                f"inclusive (num_hidden_layers={needed}), but base model has "
                f"only {text_config.num_hidden_layers}."
            )
            _truncate_config_layers(text_config, needed)

        backbone = AutoModelForCausalLM.from_pretrained(
            pretrained_model_name_or_path,
            config=config,
            **kwargs,
        )
        # Discard the vision tower / projection for multimodal wrappers — keep
        # only the text transformer. No-op for plain causal LMs.
        backbone = resolve_text_model(backbone)
        # Strip lm_head — critic never produces logits. Frees memory and
        # prevents FSDP from sharding a weight matrix we never use.
        if hasattr(backbone, "lm_head"):
            backbone.lm_head = nn.Identity()

        # Strip final layernorm — the value head must see the raw layer-K
        # residual stream, not a normalized version of it.
        inner = _inner_transformer(backbone)
        for attr in ("norm", "final_layernorm", "ln_f"):
            if hasattr(inner, attr):
                setattr(inner, attr, nn.Identity())
                break
        else:
            raise AssertionError(
                f"could not find final layernorm on {type(inner).__name__} — "
                f"add the attribute name to the list above"
            )

        model = cls(text_config, backbone)

        # If an HF-format critic checkpoint exists (from save_pretrained), load
        # the value head from it. Fresh init otherwise. Skip under
        # init_empty_weights (rank≠0 in fsdp_utils/actor.py:202 when
        # tie_word_embeddings=False): nn.Linear was registered on meta;
        # load_state_dict(cpu_tensors) would crash. Rank 0 loads,
        # _fsdp2_load_full_state_dict broadcasts.
        head_path = Path(pretrained_model_name_or_path) / "value_head.safetensors"
        if head_path.exists() and not model.value_head.weight.is_meta:
            model.value_head.load_state_dict(load_file(str(head_path)))

        # Qwen3-8B stability: freeze value_head as identity. The trainable Linear
        # is what's causing 1/||p||² gradient explosions in bf16 during the
        # first few steps. With value_head frozen, only the backbone learns —
        # and the backbone has stable bf16 dynamics (we've verified AV-SFT trains
        # without divergence on this same setup). The fixed identity head means
        # pred = backbone_last_hidden, so the AR's prediction is literally the
        # truncated forward's layer-K residual — what we want anyway, no need
        # for a learnable affine.
        if os.environ.get("NLA_FREEZE_VALUE_HEAD") == "1":
            for p in model.value_head.parameters():
                p.requires_grad_(False)
            print("[NLACriticModel] value_head FROZEN (identity)")

        # nn.Linear in __init__ is CPU/fp32; load_state_dict upcasts the bf16
        # file tensors into that. Backbone got device+dtype via kwargs (incl.
        # device_map="auto" → accelerate sharding), but value_head is outside
        # the safetensors index so accelerate doesn't see it. Align to the last
        # layer's placement so forward()'s `h` → value_head doesn't bounce
        # device/dtype. Skip on meta (FSDP rank≠0 — broadcast handles it).
        if not model.value_head.weight.is_meta:
            last = next(inner.layers[-1].parameters())
            model.value_head.to(device=last.device, dtype=last.dtype)

        return model

    def forward(self, input_ids=None, position_ids=None, attention_mask=None, **kwargs):
        out = _inner_transformer(self.backbone)(
            input_ids=input_ids,
            position_ids=position_ids,
            attention_mask=attention_mask,
            **kwargs,
        )
        # out.last_hidden_state: [B, T, d_model] — post last-kept transformer block,
        # pre-final-LN (norm was replaced with Identity in from_pretrained).
        h = out.last_hidden_state
        # Align dtype before the value_head matmul: under 4-bit QLoRA,
        # prepare_model_for_kbit_training casts the backbone's final norm to
        # fp32, so h may be fp32 while value_head is bf16 (or vice-versa).
        return NLACriticOutput(
            values=self.value_head(h.to(self.value_head.weight.dtype)),
            backbone_last_hidden=h,
        )

    def get_input_embeddings(self):
        return self.backbone.get_input_embeddings()

    def save_pretrained(self, save_directory, state_dict=None, **kwargs):
        if state_dict is None:
            state_dict = self.state_dict()
        backbone_sd = {k.removeprefix("backbone."): v for k, v in state_dict.items() if k.startswith("backbone.")}
        head_sd = {k.removeprefix("value_head."): v for k, v in state_dict.items() if k.startswith("value_head.")}

        self.backbone.save_pretrained(save_directory, state_dict=backbone_sd, **kwargs)
        save_file(head_sd, str(Path(save_directory) / "value_head.safetensors"))
        # config.json written by backbone.save_pretrained includes the truncated num_hidden_layers

    def gradient_checkpointing_enable(self, **kwargs):
        self.backbone.gradient_checkpointing_enable(**kwargs)

    def gradient_checkpointing_disable(self):
        self.backbone.gradient_checkpointing_disable()


# Path the actor dumps a fresh embedding to after each SGLang weight sync.
# nla_generate reloads from here so rollout embeddings track the trained
# model (not the stale initial HF checkpoint). Default goes under args.save
# (disk); set NLA_EMBED_DUMP_DIR=/dev/shm/nla to use tmpfs — the 1GB reload
# drops from ~1.8s to ~0.2s per rollout.
ROLLOUT_EMBED_DUMP = "_nla_rollout_embed.pt"


def embed_dump_path(save_dir: str | None) -> Path | None:
    """Resolve the embedding dump path. NLA_EMBED_DUMP_DIR overrides save_dir."""
    override = os.environ.get("NLA_EMBED_DUMP_DIR")
    base = override or save_dir
    if base is None:
        return None
    return Path(base) / ROLLOUT_EMBED_DUMP


def _find_embed_key(keys: list[str], where: str) -> str:
    matches = [k for k in keys if k.endswith(_EMBED_KEY_SUFFIXES)]
    assert len(matches) == 1, (
        f"expected exactly one input-embedding key in {where} "
        f"(suffixes {_EMBED_KEY_SUFFIXES}), got {matches!r}"
    )
    return matches[0]


def load_embedding_only(hf_checkpoint: str, dtype: torch.dtype = torch.float32) -> nn.Embedding:
    """Load ONLY the input embedding layer from an HF checkpoint's safetensors.

    Returns a plain `nn.Embedding` wrapping the weight tensor. This is the RAW
    lookup — if the model's embedding forward applies a scale (Gemma ×√d, T5),
    the CALLER must multiply. See `arch_adapters.resolve_embed_scale()` and
    how `nla_generate.py` uses `_EMBED_SCALE` explicitly.

    Avoids materializing the full model — just reads the one weight tensor
    via safe_open's lazy loading. Handles HF Hub names via snapshot_download.
    """
    root = Path(hf_checkpoint)
    if not root.exists():
        root = Path(snapshot_download(hf_checkpoint))

    index_path = root / "model.safetensors.index.json"
    if index_path.exists():
        weight_map = json.loads(index_path.read_text())["weight_map"]
        key = _find_embed_key(list(weight_map), str(index_path))
        shard = root / weight_map[key]
    else:
        shard = root / "model.safetensors"
        assert shard.exists(), (
            f"no model.safetensors or .index.json at {root!r}"
        )
        with safe_open(str(shard), framework="pt") as f:
            key = _find_embed_key(list(f.keys()), str(shard))

    with safe_open(str(shard), framework="pt") as f:
        weight = f.get_tensor(key).to(dtype)

    vocab, d_model = weight.shape
    embed = nn.Embedding(vocab, d_model, _weight=weight)
    embed.requires_grad_(False)
    embed.eval()
    return embed