"""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