| """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") |
| |
| 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") |
| 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 |
| self.peft_model = peft_model |
| self.head = head |
| self.taps = list(taps) |
| self._captured = {} |
| self._capturing = False |
| self._hooks = [] |
| |
| 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: |
| |
| |
| 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] |
| 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: |
| |
| 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) |
| |
| 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: |
| 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) |
| 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() |
| 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() |
| 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() |
| 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 |
|
|
