# MODIFIED FILE NOTICE (2026-10-17) # This file contains project modifications to upstream LOOM commit # a2ee5a7ecd38b083d9c049d97271ece38f4bce4c for the training/runtime implementation. # Those modifications are preserved in the original trained-model artifact. # This distribution adds this notice only; executable Python syntax is # unchanged from that artifact. from typing import Tuple import os import torch from torch import nn from torch import Tensor import torch.distributed as dist import torch.nn.functional as F from torch.utils.checkpoint import checkpoint from pydantic import BaseModel from models.layers import LinearInit, ScaledEmbeddingInit, Carry from models.common import IGNORE_LABEL_ID, packing_sequence_sum class LMHeadConfig(BaseModel): vocab_size: int class LMHead(nn.Module): def __init__(self, model: nn.Module, config_dict: dict) -> None: super().__init__() self.model = model # Create cache function self.create_cache = self.model.create_cache # Train extra args function self.compute_train_extra_args = self.model.compute_train_extra_args config = LMHeadConfig(**config_dict) head_hint: dict = self.model.head_hint # pyright: ignore[reportAssignmentType] # LMHead input and output self.embed_tokens = ScaledEmbeddingInit(config.vocab_size, head_hint["in"]["dim"], init_std=head_hint["in"]["init_std"]) # pyright: ignore[reportArgumentType] self.lm_head = LinearInit(head_hint["out"]["dim"], config.vocab_size, bias=False, init_std=head_hint["out"]["init_std"]) # pyright: ignore[reportArgumentType] def _linear_ce_chunk(self, hidden: Tensor, labels: Tensor) -> tuple[Tensor, Tensor]: logits = self.lm_head(hidden) loss = F.cross_entropy( logits.to(torch.float32), labels.to(torch.long), ignore_index=IGNORE_LABEL_ID, reduction="sum", ) return loss, logits.detach().argmax(dim=-1) def _chunked_linear_ce(self, hidden: Tensor, labels: Tensor, chunk_size: int) -> tuple[Tensor, Tensor]: """Sum checkpointed packed-token CE without retaining full-vocabulary logits.""" if hidden.shape[0] == 0: return self._linear_ce_chunk(hidden, labels) losses = [] predictions = [] for start in range(0, hidden.shape[0], chunk_size): chunk_hidden = hidden[start:start + chunk_size] chunk_labels = labels[start:start + chunk_size] if self.training and torch.is_grad_enabled(): loss, predicted = checkpoint( self._linear_ce_chunk, chunk_hidden, chunk_labels, use_reentrant=False, ) else: loss, predicted = self._linear_ce_chunk(chunk_hidden, chunk_labels) losses.append(loss) predictions.append(predicted) return torch.stack(losses).sum(), torch.cat(predictions) def ce_from_logits(self, logits: Tensor, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, tuple[Tensor, Tensor]]]: """CE + metrics from hidden-state logits (used by full forward and segmented BP).""" labels = batch["labels"] masks = labels != IGNORE_LABEL_ID loss = F.cross_entropy(logits.to(torch.float32), labels.to(torch.long), ignore_index=IGNORE_LABEL_ID, reduction="sum") loss_divisor = masks.sum().to(torch.float32) dist.all_reduce(loss_divisor, op=dist.ReduceOp.AVG) normalized_loss = loss / loss_divisor with torch.no_grad(): is_correct = torch.argmax(logits, dim=-1) == labels local_valid_counts = masks.sum() seq_num_tokens_correct = packing_sequence_sum(is_correct, batch["cu_seqlens"]) seq_num_valid_tokens = packing_sequence_sum(masks, batch["cu_seqlens"]) seq_is_valid = seq_num_valid_tokens > 0 metrics = { "loss": (normalized_loss.detach() * local_valid_counts, local_valid_counts), "accuracy": (is_correct.sum(), local_valid_counts), "exact_accuracy": (((seq_num_tokens_correct == seq_num_valid_tokens) & seq_is_valid).sum(), seq_is_valid.sum()), } return normalized_loss, metrics def forward(self, carry: Carry, batch: dict[str, Tensor], **kwargs) -> Tuple[Carry, Tensor] | Tuple[Carry, Tensor, dict[str, Tuple[Tensor, Tensor]]]: # Token embedding input_embedding = self.embed_tokens(batch["inputs"]) # Model forward new_carry, hidden = self.model(carry, input_embedding, **{k: v for k, v in batch.items() if k not in ("inputs", "labels")}, **kwargs) chunk_env = "LM_HEAD_CE_CHUNK_SIZE" if self.training else "LM_HEAD_EVAL_CE_CHUNK_SIZE" chunk_size = int(os.environ.get(chunk_env, "0")) chunked_ce = "labels" in batch and chunk_size > 0 if chunked_ce: loss, predicted = self._chunked_linear_ce(hidden, batch["labels"], chunk_size) else: logits = self.lm_head(hidden) # Loss & Metrics if "labels" in batch: # Masks & labels labels = batch["labels"] masks = labels != IGNORE_LABEL_ID # Loss (CE in F32) if not chunked_ce: loss = F.cross_entropy(logits.to(torch.float32), labels.to(torch.long), ignore_index=IGNORE_LABEL_ID, reduction="sum") # AllReduce loss divisor. Divide by mean of valid tokens across all processes, as gradient will be averaged. loss_divisor = masks.sum().to(torch.float32) dist.all_reduce(loss_divisor, op=dist.ReduceOp.AVG) router_z_coef = getattr(self.model, "router_z_loss_coef", 0.0) z_loss = None if router_z_coef > 0 and hasattr(self.model, "pop_router_z_loss"): z_loss = self.model.pop_router_z_loss() normalized_loss = loss / loss_divisor if z_loss is not None: normalized_loss = normalized_loss + router_z_coef * z_loss # Deep supervision: mean CE over intermediate H-cycle states. deep_sup_coef = float(getattr(self.model, "deep_sup_coef", 0.0) or 0.0) deep_sup_loss = None if ( deep_sup_coef > 0.0 and self.training and hasattr(self.model, "pop_deep_sup_states") ): inter = self.model.pop_deep_sup_states() if inter: aux_sum = None for h in inter: if chunked_ce: aux_i, _ = self._chunked_linear_ce(h, labels, chunk_size) else: aux_logits = self.lm_head(h) aux_i = F.cross_entropy( aux_logits.to(torch.float32), labels.to(torch.long), ignore_index=IGNORE_LABEL_ID, reduction="sum", ) aux_sum = aux_i if aux_sum is None else (aux_sum + aux_i) deep_sup_loss = (aux_sum / float(len(inter))) / loss_divisor normalized_loss = normalized_loss + deep_sup_coef * deep_sup_loss elif hasattr(self.model, "pop_deep_sup_states"): # Eval / coef=0: drop any stashed states so they don't leak. self.model.pop_deep_sup_states() # Scale-visible norm penalty (arXiv:2606.24898): λ · mean_k E[‖Hk‖²_rms]. # ‖v‖²_rms = mean_d(v²); keeps RMSNorm readout but makes scale visible to loss. norm_pen_coef = float(getattr(self.model, "state_norm_penalty_coef", 0.0) or 0.0) norm_pen_loss = None if ( norm_pen_coef > 0.0 and self.training and hasattr(self.model, "pop_norm_pen_states") ): states = self.model.pop_norm_pen_states() if states: rms2_sum = None for h in states: # E[‖h‖²_rms] over tokens; h is [N, D] packed. rms2 = h.to(torch.float32).pow(2).mean() rms2_sum = rms2 if rms2_sum is None else (rms2_sum + rms2) norm_pen_loss = rms2_sum / float(len(states)) normalized_loss = normalized_loss + norm_pen_coef * norm_pen_loss elif hasattr(self.model, "pop_norm_pen_states"): self.model.pop_norm_pen_states() # Loop-Δz RMS match (CoD / residual-scaling): unitless update-norm balance. # Final-z RMS is LN-flat (rms2 saw match=0); use RMS(z_t−z_{t−1}) instead. # Full-H: L = (rms_ΔH / rms_Δ2 − 1)². Plus stream-gain: # λ* = clamp(λ · ratio, λ_min, λ_max); L += (λ − stopgrad(λ*))² # Plus depth prior (fixed-λ BEST / transfer point ≈1.6), gated by |ratio−1| # so it only pulls when the stream is imbalanced — not pack-specific LR. rms_match = bool(getattr(self.model, "cycle_lambda_rms_match", False)) rms_match_loss = None stream_gain_loss = None depth_prior_loss = None if rms_match and self.training: ref, delta_last = (None, None) if hasattr(self.model, "pop_lambda_rms_full_h"): ref, delta_last = self.model.pop_lambda_rms_full_h() if ref is not None and delta_last is not None: rms_last = ( delta_last.to(torch.float32).pow(2).mean().clamp_min(1e-12).sqrt() ) ratio = rms_last / ref rms_match_loss = (ratio - 1.0).pow(2) normalized_loss = normalized_loss + rms_match_loss get_lam = getattr(self.model, "_cycle_lambda_res_live", None) if not callable(get_lam): get_lam = getattr(self.model, "_cycle_lambda_res", None) if callable(get_lam): lam = get_lam() if isinstance(lam, torch.Tensor): lam_f = lam.to(torch.float32).reshape(()) lam_min = float( getattr(self.model, "cycle_scale_lambda_min", 1.5) or 1.5 ) lam_max = float( getattr(self.model, "cycle_scale_lambda_max", 0.0) or 0.0 ) if lam_max <= 0.0: lam_max = lam_min + 2.0 # Operating band (residual-scaling BEST ≈1.6–1.7). lo = float( getattr(self.model, "cycle_lambda_depth_prior_lo", 1.58) or 1.58 ) hi = float( getattr(self.model, "cycle_lambda_depth_prior_hi", 1.75) or 1.75 ) # Match residual H-aware hi so prior/stream gates align with detach. H = float( max(int(getattr(self.model, "num_loops", 1) or 1), 1) ) hi = min(hi, 1.65 + 0.10 * (9.0 / H) ** 0.5) mid = 0.5 * (lo + hi) # ≈1.65 # Climb zone [lo, mid): keep stream-gain so init15 can rise # (rms7 froze at lo→λ≈1.58→eval 7.2). # rms7b: gate stream in [mid,hi]; rms8: also gate when λ>hi # (CE+stream yank caused soft15→1.80 / soft20→2.17 rebound). mode = str( getattr(self.model, "cycle_lambda_detach_mode", "rms7b") or "rms7b" ) if mode == "rms8": stream_off = lam_f >= mid elif mode == "rms9": latched = getattr( self.model, "_cycle_lambda_settled_latched", None ) latched_b = bool( latched.detach().item() if isinstance(latched, torch.Tensor) else bool(latched) ) stream_off = latched_b or ((lam_f >= mid) & (lam_f <= hi)) elif mode == "rms10": # Stream-gain off at/above hi (ceiling); keep below mid live. stream_off = lam_f >= hi elif mode == "rms11": # Settled band OR above hi (ceiling zone): no stream yank. stream_off = (lam_f >= mid) | (lam_f >= hi) elif mode == "rms12": # STE ceiling: keep stream/prior ON when λ>hi (pull down); # gate only in settled [mid,hi]. stream_off = (lam_f >= mid) & (lam_f <= hi) else: stream_off = (lam_f >= mid) & (lam_f <= hi) ratio_d = ratio.detach() lam_star = (lam_f.detach() * ratio_d).clamp(lam_min, lam_max) stream_raw = (lam_f - lam_star).pow(2) stream_gain_loss = torch.where( stream_off, lam_f.new_zeros(()), stream_raw ) normalized_loss = normalized_loss + stream_gain_loss lam_ref = float( getattr(self.model, "cycle_lambda_depth_prior", 1.6) or 0.0 ) if lam_ref > 0.0: # Asymmetric operating band: # λ > hi → (λ−ref)² # λ < mid → (λ−mid)² (climb through [lo,mid) too) # mid≤λ≤hi → 0 (settled; stream gated above) coef = float( getattr(self.model, "cycle_lambda_depth_prior_coef", 1.0) or 1.0 ) high_pen = torch.where( lam_f > hi, (lam_f - lam_ref).pow(2), lam_f.new_zeros(()) ) low_pen = torch.where( lam_f < mid, (lam_f - mid).pow(2), lam_f.new_zeros(()) ) depth_prior_loss = coef * (high_pen + low_pen) normalized_loss = normalized_loss + depth_prior_loss elif hasattr(self.model, "pop_lambda_rms_states"): # Legacy fallback (non-segmented / short H). states = self.model.pop_lambda_rms_states() if len(states) >= 2: rms_list = [ h.to(torch.float32).pow(2).mean().clamp_min(1e-12).sqrt() for h in states ] ref0 = rms_list[0].detach() acc = None for r in rms_list[1:]: term = (r / ref0 - 1.0).pow(2) acc = term if acc is None else (acc + term) rms_match_loss = acc / float(len(rms_list) - 1) normalized_loss = normalized_loss + rms_match_loss elif hasattr(self.model, "pop_lambda_rms_states"): self.model.pop_lambda_rms_states() # Accuracy with torch.no_grad(): is_correct = (predicted if chunked_ce else torch.argmax(logits, dim=-1)) == labels local_valid_counts = masks.sum() # Sequence-level statistics seq_num_tokens_correct = packing_sequence_sum(is_correct, batch["cu_seqlens"]) seq_num_valid_tokens = packing_sequence_sum(masks, batch["cu_seqlens"]) seq_is_valid = seq_num_valid_tokens > 0 # Metrics metrics = { "loss": (normalized_loss.detach() * local_valid_counts, local_valid_counts), "accuracy": (is_correct.sum(), local_valid_counts), "exact_accuracy": (((seq_num_tokens_correct == seq_num_valid_tokens) & seq_is_valid).sum(), seq_is_valid.sum()), } if z_loss is not None: metrics["router_z_loss"] = (z_loss.detach(), torch.ones((), device=z_loss.device, dtype=torch.float32)) if deep_sup_loss is not None: metrics["deep_sup_loss"] = ( deep_sup_loss.detach(), torch.ones((), device=deep_sup_loss.device, dtype=torch.float32), ) if norm_pen_loss is not None: metrics["state_norm_penalty"] = ( norm_pen_loss.detach(), torch.ones((), device=norm_pen_loss.device, dtype=torch.float32), ) if rms_match_loss is not None: metrics["lambda_rms_match"] = ( rms_match_loss.detach(), torch.ones((), device=rms_match_loss.device, dtype=torch.float32), ) if stream_gain_loss is not None: metrics["lambda_stream_gain"] = ( stream_gain_loss.detach(), torch.ones((), device=stream_gain_loss.device, dtype=torch.float32), ) if depth_prior_loss is not None: metrics["lambda_depth_prior"] = ( depth_prior_loss.detach(), torch.ones((), device=depth_prior_loss.device, dtype=torch.float32), ) return new_carry, normalized_loss, metrics return new_carry, logits