Blame for code/model.py

"""Backbone (Qwen3.5/3.6-MoE, the default in this repository via SH_BACKBONE=qwen35moe, or Ministral-3 Mistral3ForConditionalGeneration)
+ LoRA + joint schema head (DESIGN.md 2.1, 3.1).

- The inner model (`.model`, Mistral3Model) is called; its `last_hidden_state` is the post-norm final state.
  `lm_head` is never called and is deleted after load (131,072 x 4,096 bf16 = 1.07e9 B).
- Extra taps: outputs of decoder layers by index, captured with forward hooks.
- LoRA: language-model linear layers only (q/k/v/o/gate/up/down), vision tower + projector frozen, alpha = 2r.
  Same leaf discovery and vision-exclusion rule as the Standard One LoRA trainer.
"""
import json
import os
from pathlib import Path

import torch
import torch.nn as nn

from head import HeadConfig, JointSchemaHead
from render import token_roles

VISION_EXCLUDE_REGEX = r".*\.(vision_tower|multi_modal_projector|visual)\..*"
QWEN = os.environ.get("SH_BACKBONE", "qwen35moe") in ("qwen35moe", "qwen35")   # Qwen3.5/3.6 (MoE or dense; Gated DeltaNet + full attention)
# Qwen3.5 linear-attention projections that get LoRA (in_proj_a/b are 32-wide gate projections: left frozen)
LINEAR_ATTN_LEAVES = {"in_proj_qkv", "in_proj_z", "out_proj"}


def lora_leaves(model):
    """Leaf names of the language-model attention + MLP linears (as the trainer's `attention_mlp` targets)."""
    attention, mlp = set(), set()
    for name, module in model.named_modules():
        if not isinstance(module, nn.Linear) or ".language_model." not in name:
            continue
        leaf = name.rsplit(".", 1)[-1]
        if ".self_attn." in name:
            attention.add(leaf)
        elif ".linear_attn." in name and leaf in LINEAR_ATTN_LEAVES:
            attention.add(leaf)
        elif ".mlp." in name and leaf in {"gate_proj", "up_proj", "down_proj"}:
            mlp.add(leaf)
    if not attention:
        raise ValueError("no attention projections found")
    return sorted(attention | mlp)


ENCODER = os.environ.get("SH_BACKBONE", "qwen35moe") == "encoder"


def decoder_layers(model):
    if ENCODER:
        return model.layers
    lm = model.model.language_model
    return lm.layers


def load_encoder(path, attn="sdpa"):
    """Text encoder: ModernBERT/mmBERT (AutoModel) or the text tower of EmbeddingGemma 2 (vision/audio dropped)."""
    import transformers
    cfg = transformers.AutoConfig.from_pretrained(str(path), local_files_only=True)
    kw = dict(local_files_only=True, dtype=torch.bfloat16, attn_implementation=attn)
    if cfg.model_type == "embedding_gemma2":
        kw.pop("attn_implementation")   # the audio tower has no SDPA path; the default picks SDPA for the text tower
        full = transformers.AutoModel.from_pretrained(str(path), **kw)
        enc = full.language_model
        del full
        return enc
    return transformers.AutoModel.from_pretrained(str(path), **kw)


def tiny_config(vocab_size=131072, layers=3, hidden=64, heads=4, kv_heads=2, intermediate=128):
    """A tiny, randomly initialised Mistral3 (Ministral3 text + Pixtral vision) config for CPU smoke tests."""
    from transformers import Mistral3Config
    text = {"model_type": "ministral3", "hidden_size": hidden, "intermediate_size": intermediate,
            "num_hidden_layers": layers, "num_attention_heads": heads, "num_key_value_heads": kv_heads,
            "head_dim": hidden // heads, "vocab_size": vocab_size, "max_position_embeddings": 262144,
            "rms_norm_eps": 1e-5, "hidden_act": "silu", "tie_word_embeddings": False, "sliding_window": None,
            "rope_parameters": {"beta_fast": 32.0, "beta_slow": 1.0, "factor": 16.0, "llama_4_scaling_beta": 0.1,
                                "mscale": 1.0, "mscale_all_dim": 1.0, "original_max_position_embeddings": 16384,
                                "rope_theta": 1000000.0, "rope_type": "yarn", "type": "yarn"}}
    vision = {"model_type": "pixtral", "hidden_size": 32, "intermediate_size": 64, "num_hidden_layers": 1,
              "num_attention_heads": 2, "head_dim": 16, "image_size": 1540, "patch_size": 14, "num_channels": 3,
              "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, "rope_theta": 10000.0}
    return Mistral3Config(text_config=text, vision_config=vision, image_token_index=10, spatial_merge_size=2,
                          multimodal_projector_bias=False, projector_hidden_act="gelu", vision_feature_layer=-1,
                          tie_word_embeddings=False)


class SchemaModel(nn.Module):
    """Backbone (optionally LoRA-wrapped) + head. `forward_micro` runs one micro-batch of encoded requests."""

    def __init__(self, backbone, head, taps, peft_model=None):
        super().__init__()
        self.backbone = backbone          # Mistral3ForConditionalGeneration (LoRA layers injected in place)
        self.peft_model = peft_model      # the PeftModel wrapper (for save_pretrained) or None
        self.head = head
        self.taps = list(taps)
        self._captured = {}
        self._capturing = False   # hooks record only during our own forward, not during checkpoint recompute
        self._hooks = []
        # EmbeddingGemma 2 projects its final state 512 -> 768 after the norm; the head reads the 512-d norm output
        self._final_hook = ENCODER and hasattr(backbone, "embedding_projection")
        if self._final_hook:
            self._hooks.append(backbone.norm.register_forward_hook(self._make_hook("final")))
        layers = decoder_layers(backbone)
        for t in self.taps:
            if t == "final":
                continue
            idx = int(t)
            if not 0 <= idx < len(layers):
                raise ValueError(f"tap layer {idx} out of range (0..{len(layers) - 1})")
            self._hooks.append(layers[idx].register_forward_hook(self._make_hook(idx)))

    def _make_hook(self, idx):
        def hook(module, inputs, output):
            if self._capturing:
                self._captured[idx] = output[0] if isinstance(output, (tuple, list)) else output
        return hook

    def device(self):
        return next(self.head.parameters()).device

    def backbone_states(self, batch):
        """Run the inner model on a padded batch; returns list of [B, T, H] tensors in tap order."""
        self._captured = {}
        dev = batch["input_ids"].device
        kwargs = {k: v for k, v in batch.items() if k in ("input_ids", "attention_mask", "position_ids",
                                                           "pixel_values", "image_sizes", "image_grid_thw",
                                                           "mm_token_type_ids")}
        if QWEN:
            # Qwen3.5 uses 3-axis (M-RoPE) positions computed from input_ids (+ image grids); rows are right-padded
            # so the model's own positions equal arange per row.
            kwargs.pop("position_ids", None)
        use_bf16 = dev.type == "cuda"
        self._capturing = True
        try:
            with torch.autocast(device_type=dev.type, dtype=torch.bfloat16, enabled=use_bf16):
                if ENCODER:
                    out = self.backbone(input_ids=kwargs["input_ids"], attention_mask=kwargs["attention_mask"],
                                        return_dict=True)
                else:
                    out = self.backbone.model(**kwargs, use_cache=False, return_dict=True)
        finally:
            self._capturing = False
        states = []
        for t in self.taps:
            if t == "final":
                states.append(self._captured["final"] if self._final_hook else out.last_hidden_state)
            else:
                states.append(self._captured[int(t)])
        self._captured = {}
        return states

    def forward_micro(self, encs, batch):
        """encs: list of encoded requests (render.Encoder.encode outputs, right-padded in `batch`).
        Returns list (per request) of lists (per question) of fp32 logit vectors (canonical option order)."""
        states = self.backbone_states(batch)
        outputs = []
        dev = states[0].device
        with torch.autocast(device_type=dev.type, enabled=False):
            for i, enc in enumerate(encs):
                n = enc["n_tokens"]
                taps = [s[i, :n] for s in states]   # bf16 on GPU; the head casts inside its checkpointed projection
                roles = torch.tensor(enc.get("_roles") or token_roles(enc), device=dev)
                outputs.append(self.head.forward_one(taps, enc, roles))
        return outputs

    def trainable_named(self):
        return [(n, p) for n, p in self.named_parameters() if p.requires_grad]


def pad_pixel_values(chunks):
    """Join per-row pixel tensors [n_i, C, H_i, W_i] into one [sum n_i, C, maxH, maxW] tensor, zero-padded at the
    bottom/right. Every row went through the processor on its own, so H/W differ between rows; this is exactly what
    PixtralImageProcessor._pad_for_batching does for a multi-image call, and the vision tower crops each image back
    to its `image_sizes` entry (patch_embeds[..., :h // patch, :w // patch]) so the padding never reaches the model."""
    h = max(int(c.shape[-2]) for c in chunks)
    w = max(int(c.shape[-1]) for c in chunks)
    return torch.cat([torch.nn.functional.pad(c, (0, w - int(c.shape[-1]), 0, h - int(c.shape[-2])))
                      for c in chunks], 0)


def make_batch(encs, pad_id, device):
    """Right-padded text batch (position ids restart per row); image micro-batches carry pixel values."""
    longest = max(e["n_tokens"] for e in encs)
    ids = torch.full((len(encs), longest), pad_id, dtype=torch.long)
    mask = torch.zeros((len(encs), longest), dtype=torch.long)
    for i, e in enumerate(encs):
        ids[i, :e["n_tokens"]] = torch.tensor(e["input_ids"])
        mask[i, :e["n_tokens"]] = 1
    pos = torch.arange(longest).unsqueeze(0).expand(len(encs), -1).contiguous()
    batch = {"input_ids": ids.to(device), "attention_mask": mask.to(device), "position_ids": pos.to(device)}
    pix = [e.get("_pixel") for e in encs]
    if any(p is not None for p in pix):
        if not all(p is not None for p in pix):
            raise ValueError("image micro-batches must be homogeneous")
        if QWEN:
            # Qwen: flattened patches [sum patches, dim] + one (t, h, w) grid row per image, joined in row order
            batch["pixel_values"] = torch.cat([torch.as_tensor(p["pixel_values"]) for p in pix], 0).to(device)
            batch["image_grid_thw"] = torch.cat([torch.as_tensor(p["image_grid_thw"]) for p in pix], 0).to(device)
            # M-RoPE needs to know which tokens are image tokens (processor output, right-padded like input_ids)
            mm = torch.zeros((len(encs), longest), dtype=torch.long)
            for i, p in enumerate(pix):
                mm[i, :len(p["mm_token_type_ids"])] = torch.tensor(p["mm_token_type_ids"])
            batch["mm_token_type_ids"] = mm.to(device)
            if device.type == "cuda":
                batch["pixel_values"] = batch["pixel_values"].to(torch.bfloat16)
            return batch
        batch["pixel_values"] = pad_pixel_values([torch.as_tensor(p["pixel_values"]) for p in pix]).to(device)
        batch["image_sizes"] = torch.cat([torch.as_tensor(p["image_sizes"]) for p in pix], 0).to(device)
        if device.type == "cuda":
            batch["pixel_values"] = batch["pixel_values"].to(torch.bfloat16)
    return batch


def backbone_class(snapshot=None):
    import transformers
    if not QWEN:
        return transformers.Mistral3ForConditionalGeneration
    if snapshot is not None:   # dense Qwen3_5ForConditionalGeneration or MoE Qwen3_5MoeForConditionalGeneration
        arch = json.load(open(Path(snapshot) / "config.json"))["architectures"][0]
        return getattr(transformers, arch)
    return transformers.Qwen3_5MoeForConditionalGeneration


def tiny_qwen_config(vocab_size=248320, layers=4, hidden=64):
    """A tiny random Qwen3.5-MoE (3 linear-attention layers + 1 full-attention layer, 4 experts) for CPU tests."""
    from transformers import Qwen3_5MoeConfig
    text = {"hidden_size": hidden, "num_hidden_layers": layers, "num_attention_heads": 4, "num_key_value_heads": 2,
            "head_dim": 16, "vocab_size": vocab_size, "num_experts": 4, "num_experts_per_tok": 2,
            "moe_intermediate_size": 32, "shared_expert_intermediate_size": 32, "linear_num_value_heads": 4,
            "linear_num_key_heads": 2, "linear_key_head_dim": 16, "linear_value_head_dim": 16,
            "linear_conv_kernel_dim": 4, "full_attention_interval": 4, "max_position_embeddings": 262144,
            "layer_types": ["linear_attention"] * (layers - 1) + ["full_attention"]}
    vision = {"depth": 1, "hidden_size": 32, "intermediate_size": 64, "num_heads": 2, "out_hidden_size": hidden,
              "patch_size": 16, "spatial_merge_size": 2, "temporal_patch_size": 2}
    return Qwen3_5MoeConfig(text_config=text, vision_config=vision)


def build_model(snapshot=None, tiny=None, lora_rank=64, lora_alpha=None, lora_dropout=0.0, taps=("final", 25),
                head_overrides=None, device="cpu", grad_checkpointing=True, init_adapter=None, attn="sdpa",
                seed=0):
    """Load (or create tiny) backbone, delete lm_head, freeze, add LoRA (rank 0 = none), attach the head.

    Returns (SchemaModel, info dict)."""
    import peft
    import transformers
    torch.manual_seed(seed)
    if ENCODER:
        return _build_encoder_model(snapshot, taps, head_overrides, device, grad_checkpointing, attn, seed,
                                    full_ft=lora_rank == 0 and os.environ.get("SH_FULL_FT") == "1")
    if tiny is not None:
        torch.manual_seed(1234)   # fixed: the tiny random "base" must be identical in training and serving
        cfg = tiny if not isinstance(tiny, dict) else (tiny_qwen_config if QWEN else tiny_config)(**tiny)
        backbone = backbone_class()._from_config(cfg, attn_implementation=attn)
        backbone = backbone.to(device)
    else:
        backbone = backbone_class(snapshot).from_pretrained(
            str(snapshot), local_files_only=True, trust_remote_code=False, dtype=torch.bfloat16,
            attn_implementation=attn, device_map={"": device} if str(device) != "cpu" else None)
        if str(device) == "cpu":
            backbone = backbone.to(device)
    torch.manual_seed(seed)
    backbone.config.use_cache = False
    lm_head_bytes = backbone.lm_head.weight.numel() * backbone.lm_head.weight.element_size()
    del backbone.lm_head
    backbone.lm_head = nn.Identity()  # never called; keeps attribute access working
    for p in backbone.parameters():
        p.requires_grad_(False)
    if grad_checkpointing:
        backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
        vt = getattr(backbone.model, "visual" if QWEN else "vision_tower", None)
        if vt is not None and hasattr(vt, "gradient_checkpointing_disable"):
            vt.gradient_checkpointing_disable()
    leaves = lora_leaves(backbone)
    peft_model = None
    if init_adapter:
        peft_model = peft.PeftModel.from_pretrained(backbone, str(init_adapter), is_trainable=True)
    elif lora_rank:
        alpha = lora_alpha if lora_alpha is not None else 2 * lora_rank
        peft_model = peft.get_peft_model(backbone, peft.LoraConfig(
            r=lora_rank, lora_alpha=alpha, lora_dropout=lora_dropout, target_modules=leaves,
            exclude_modules=VISION_EXCLUDE_REGEX, bias="none"))
    leaked = 0
    for name, p in backbone.named_parameters():
        if "vision_tower" in name or "multi_modal_projector" in name or ".visual." in name:
            if p.requires_grad:
                leaked += 1
            p.requires_grad_(False)
        elif p.requires_grad:
            p.data = p.data.float()  # LoRA weights fp32 (as the trainer)
    text_cfg = backbone.config.text_config
    hcfg = HeadConfig(hidden_size=text_cfg.hidden_size, n_taps=len(taps), taps=list(taps),
                      **(head_overrides or {}))
    head = JointSchemaHead(hcfg).to(device=device, dtype=torch.float32)
    model = SchemaModel(backbone, head, taps, peft_model)
    info = {"lora_leaves": leaves, "lora_rank": lora_rank,
            "lora_alpha": (lora_alpha if lora_alpha is not None else 2 * lora_rank) if lora_rank else None,
            "lora_params": sum(p.numel() for n, p in backbone.named_parameters() if p.requires_grad),
            "head_params": sum(p.numel() for p in head.parameters()), "lm_head_bytes_freed": lm_head_bytes,
            "vision_lora_leaked_frozen": leaked, "taps": list(taps), "head_config": hcfg.to_dict(),
            "layers": len(decoder_layers(backbone)), "hidden_size": text_cfg.hidden_size}
    return model, info


def _build_encoder_model(snapshot, taps, head_overrides, device, grad_checkpointing, attn, seed, full_ft):
    backbone = load_encoder(snapshot, attn).to(device)
    torch.manual_seed(seed)
    for p in backbone.parameters():
        p.requires_grad_(False)
    if full_ft:
        for p in backbone.parameters():
            p.data = p.data.float()     # fp32 master weights, bf16 autocast in forward
            p.requires_grad_(True)
    if grad_checkpointing:
        backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    hidden = backbone.config.hidden_size
    hcfg = HeadConfig(hidden_size=hidden, n_taps=len(taps), taps=list(taps), **(head_overrides or {}))
    head = JointSchemaHead(hcfg).to(device=device, dtype=torch.float32)
    model = SchemaModel(backbone, head, taps, None)
    info = {"lora_leaves": [], "lora_rank": 0, "lora_alpha": None, "full_ft": full_ft,
            "lora_params": sum(p.numel() for p in backbone.parameters() if p.requires_grad),
            "backbone_params": sum(p.numel() for p in backbone.parameters()),
            "head_params": sum(p.numel() for p in head.parameters()), "lm_head_bytes_freed": 0,
            "vision_lora_leaked_frozen": 0, "taps": list(taps), "head_config": hcfg.to_dict(),
            "layers": len(decoder_layers(backbone)), "hidden_size": hidden, "encoder": type(backbone).__name__}
    return model, info


def save_head(head, folder, extra=None):
    from safetensors.torch import save_file
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    save_file({k: v.detach().contiguous().cpu() for k, v in head.state_dict().items()},
              str(folder / "schema_head.safetensors"))
    cfg = {"head": head.cfg.to_dict(), **(extra or {})}
    tmp = folder / "schema_head_config.json.tmp"
    tmp.write_text(json.dumps(cfg, indent=2) + "\n")
    os.replace(tmp, folder / "schema_head_config.json")


def load_head_weights(head, folder):
    from safetensors.torch import load_file
    state = load_file(str(Path(folder) / "schema_head.safetensors"))
    missing, unexpected = head.load_state_dict(state, strict=True), None
    return missing