"""Parallel candidate scoring over shared text/image computation.""" import hashlib import json from pathlib import Path import torch import torch.nn.functional as F from torch import Tensor, nn from jolt.adapters import Qwen35Adapter from jolt.recipe import IMAGE_PIXELS, MAX_LENGTH, normalize_architecture def model_revision(directory): """Identify a downloaded Hub snapshot or a locally pinned training base.""" directory = Path(directory) revision = directory / "revision.txt" if revision.is_file(): return revision.read_text().strip() if directory.parent.name == "snapshots" and len(directory.name) == 40: return directory.name raise ValueError("The base model must be a pinned Hub snapshot or contain revision.txt") class ResidualBottleneck(nn.Module): def __init__(self, hidden_size, width, activation): super().__init__() self.norm = nn.LayerNorm(hidden_size, eps=1e-5) self.down = nn.Linear(hidden_size, width, bias=False) self.up = nn.Linear(width, hidden_size, bias=False) self.activation = F.silu if activation == "silu" else F.gelu nn.init.xavier_uniform_(self.down.weight) nn.init.zeros_(self.up.weight) def forward(self, hidden): residual = hidden block_input = hidden.to(dtype=self.norm.weight.dtype) update = self.up(self.activation(self.down(self.norm(block_input)))) return residual + update.to(dtype=residual.dtype) class CandidateInteraction(nn.Module): def __init__(self, hidden_size): super().__init__() width = 128 self.heads = 4 self.head_width = width // self.heads self.norm_attention = nn.LayerNorm(hidden_size, eps=1e-5) self.query = nn.Linear(hidden_size, width) self.key = nn.Linear(hidden_size, width) self.value = nn.Linear(hidden_size, width) self.output = nn.Linear(width, hidden_size) self.norm_mlp = nn.LayerNorm(hidden_size, eps=1e-5) self.mlp_in = nn.Linear(hidden_size, 256) self.mlp_out = nn.Linear(256, hidden_size) for layer in (self.query, self.key, self.value, self.mlp_in): nn.init.xavier_uniform_(layer.weight) nn.init.zeros_(layer.bias) for layer in (self.output, self.mlp_out): nn.init.zeros_(layer.weight) nn.init.zeros_(layer.bias) def forward(self, hidden, valid): if valid is None: raise ValueError("Candidate attention requires a validity mask") output_dtype = hidden.dtype hidden = hidden.to(dtype=self.norm_attention.weight.dtype) normalized = self.norm_attention(hidden) batch, candidates, _ = hidden.shape def split_heads(tensor): return tensor.view(batch, candidates, self.heads, self.head_width).transpose(1, 2) attended = F.scaled_dot_product_attention( split_heads(self.query(normalized)), split_heads(self.key(normalized)), split_heads(self.value(normalized)), attn_mask=valid[:, None, None, :], dropout_p=0.0, ) attended = attended.transpose(1, 2).reshape(batch, candidates, -1) hidden = hidden + self.output(attended) hidden = hidden + self.mlp_out(F.gelu(self.mlp_in(self.norm_mlp(hidden)))) return hidden.masked_fill(~valid[..., None], 0).to(dtype=output_dtype) class GeluScoringHead(nn.Module): def __init__(self, hidden_size): super().__init__() self.in_proj = nn.Linear(hidden_size, 256, bias=False) self.out_proj = nn.Linear(256, 1, bias=False) nn.init.xavier_uniform_(self.in_proj.weight) nn.init.normal_(self.out_proj.weight, std=0.01) def forward(self, hidden): return self.out_proj(F.gelu(self.in_proj(hidden.float()))).squeeze(-1) class InterpolatedActivation(nn.Module): def __init__(self, initial_mix): super().__init__() self.mix = nn.Parameter(torch.tensor(float(initial_mix), device="cuda")) def forward(self, value): mix = self.mix.to(value.dtype) return (1 - mix) * F.silu(value) + mix * F.gelu(value) class DecisionModel(nn.Module): def __init__(self, checkpoint: str | Path, *, adapter=None, architecture=None): super().__init__() self.adapter = adapter or Qwen35Adapter() self.base_path = Path(checkpoint) self.architecture = normalize_architecture(architecture) self.backbone = self.adapter.load(checkpoint, architecture=self.architecture) hidden_size = self.adapter.hidden_size(self.backbone) self.processor = self.adapter.processor(checkpoint) self.head = nn.Linear(hidden_size, 1, bias=False, device="cuda", dtype=torch.float32) nn.init.normal_(self.head.weight, std=0.01) self.lora_rank = 8 self.lora_alpha = 16 self.lora_dropout = 0.0 self.marker_id = self.processor.tokenizer.convert_tokens_to_ids(self.adapter.marker) if self.processor.tokenizer.encode(self.adapter.marker, add_special_tokens=False) != [ self.marker_id ]: raise ValueError("The candidate marker must be a single reserved token") # Added modules use a forked RNG so they do not shift the baseline LoRA # initialization that follows enable_lora(). with torch.random.fork_rng(devices=[torch.cuda.current_device()]): if self.architecture["head"] == "gelu-mlp": self.head = GeluScoringHead(hidden_size).to("cuda") if self.architecture["candidate_attention"]: self.candidate_block = CandidateInteraction(hidden_size).to("cuda") else: self.candidate_block = None adapter_config = self.architecture["adapters"] self.layer_adapters = nn.ModuleDict( { str(index): ResidualBottleneck( hidden_size, adapter_config["width"], adapter_config["activation"] ).to("cuda") for index in adapter_config["layers"] } ) self.layer_mixture_weights = ( nn.Parameter(torch.zeros(3, device="cuda")) if self.architecture["readout"] == "layer-mixture" else None ) activation_config = self.architecture["backbone_activation"] self.activation_mix_layers = tuple(activation_config["layers"]) text_layers = self.adapter.language_layers(self.backbone) for index in self.activation_mix_layers: if not 0 <= index < len(text_layers): raise ValueError(f"Backbone activation layer is out of range: {index}") text_layers[index].mlp.act_fn = InterpolatedActivation( activation_config["initial_mix"] ).to("cuda") layer_count = len(self.adapter.language_layers(self.backbone)) self.capture_layers = (20, 21, 22, 23) if self.layer_mixture_weights is not None else () for index in self.capture_layers: if index >= layer_count: raise ValueError("Layer mixture requires decoder layers 21 through 24") for index in self.layer_adapters: if not 0 <= int(index) < layer_count: raise ValueError(f"Residual adapter layer is out of range: {index}") self._adapter_hooks = [] def forward( self, inputs: dict[str, Tensor], positions: Tensor, *, candidate_mask: Tensor | None = None ) -> Tensor: hidden, offset, intermediate = self.adapter.forward( self.backbone, inputs, capture_layers=self.capture_layers ) local_positions = (positions - offset).clamp_min(0) rows = torch.arange(hidden.shape[0], device=hidden.device)[:, None] if self.layer_mixture_weights is None: candidate_hidden = hidden[rows, local_positions] else: norm = self.adapter.final_norm(self.backbone) layer_features = torch.stack( [ norm(intermediate[index][rows, local_positions]) for index in self.capture_layers ], dim=0, ) layer_features = layer_features.float() final = layer_features[-1] deltas = layer_features[:-1] - final candidate_hidden = final + ( self.layer_mixture_weights.float()[:, None, None, None] * deltas ).sum(0) if self.candidate_block is not None: candidate_hidden = self.candidate_block(candidate_hidden, candidate_mask) if isinstance(self.head, GeluScoringHead): return self.head(candidate_hidden) return self.head(candidate_hidden.float()).squeeze(-1) def enable_lora(self, *, checkpointing=True, rank=8, alpha=16, dropout=0.0): self.lora_rank = rank self.lora_alpha = alpha self.lora_dropout = dropout self.adapter.adapt_language( self.backbone, training=checkpointing, rank=rank, alpha=alpha, dropout=dropout, ) layers = self.adapter.language_layers(self.backbone) for index in self.activation_mix_layers: layers[index].mlp.act_fn.mix.requires_grad_(True) for handle in self._adapter_hooks: handle.remove() self._adapter_hooks.clear() for index, module in self.layer_adapters.items(): layer = layers[int(index)] def apply_adapter(_layer, _inputs, output, adapter_module=module): if isinstance(output, tuple): return (adapter_module(output[0]), *output[1:]) return adapter_module(output) self._adapter_hooks.append(layer.register_forward_hook(apply_adapter)) def load_adapter(self, directory): """Load the same verified unmerged weights for learning or deployment.""" from safetensors.torch import load_file directory = Path(directory) metadata = json.loads((directory / "dohnuts.json").read_text()) weights = directory / "adapter.safetensors" if ( metadata.get("format_version") not in {1, 2, 3} or metadata.get("adapter", "qwen3.5") != self.adapter.name or metadata.get("base_revision") != model_revision(self.base_path) or (metadata["image_pixels"], metadata["max_length"], metadata["lora_rank"]) != (IMAGE_PIXELS, MAX_LENGTH, self.lora_rank) or metadata.get("lora_alpha", 16) != self.lora_alpha or metadata.get("lora_dropout", 0.0) != self.lora_dropout or normalize_architecture(metadata.get("architecture")) != self.architecture or hashlib.sha256(weights.read_bytes()).hexdigest() != metadata["weights_sha256"] ): raise ValueError( "Checkpoint differs from the pinned model, training recipe, or checksum" ) state = load_file(str(weights)) parameters = {n: p for n, p in self.named_parameters() if p.requires_grad} if state.keys() != parameters.keys() or any( state[name].shape != parameter.shape for name, parameter in parameters.items() ): raise ValueError("Checkpoint trainable parameters differ") with torch.no_grad(): for name, parameter in parameters.items(): parameter.copy_(state[name]) return metadata def merge(self): self.adapter.merge(self.backbone) self.requires_grad_(False).eval() def train(self, mode: bool = True): super().train(mode) self.adapter.freeze_vision(self.backbone) return self def marker_positions(input_ids: Tensor, marker_id: int, candidates: int) -> Tensor: positions = [] for row in input_ids: found = (row == marker_id).nonzero(as_tuple=True)[0] if found.numel() != candidates: raise ValueError(f"Expected {candidates} candidate markers, found {found.numel()}") positions.append(found) return torch.stack(positions)