"""Sparse whole-block routing with top-k training and causal threshold inference. Raposo et al., arXiv:2404.02258, sections 3.4/3.5 (same-router BCE option). Only selected tokens enter attention AND SwiGLU. Original RoPE positions and sequence order survive packing. Eval never uses sequence-wide top-k decisions. """ import copy import torch from torch import nn from torch.nn import functional as F from transformers.cache_utils import DynamicCache from transformers.modeling_outputs import BaseModelOutputWithPast from .modeling_q50m import ( Q50MDecoderLayer, Q50MForCausalLM, Q50MPreTrainedModel, Q50MRMSNorm, Q50MRotaryEmbedding, ) from .configuration_qmod import QMoDConfig class QMoDCache(DynamicCache): """Packed per-layer KV with source positions, validity and a global clock. Padding between batch members/chunks can remain in storage; masked entries never become keys. The clock counts source tokens, not sparse KV entries. """ def __init__(self, config): super().__init__(config=config) self.seen_tokens = 0 self.token_positions = [None] * config.num_hidden_layers self.valid_tokens = [None] * config.num_hidden_layers def get_seq_length(self, layer_idx=0): return self.seen_tokens def append_metadata(self, layer_idx, positions, valid): old = self.token_positions[layer_idx] self.token_positions[layer_idx] = positions if old is None else torch.cat((old, positions), -1) old_valid = self.valid_tokens[layer_idx] self.valid_tokens[layer_idx] = valid if old_valid is None else torch.cat((old_valid, valid), -1) return self.token_positions[layer_idx], self.valid_tokens[layer_idx] def reorder_cache(self, beam_idx): super().reorder_cache(beam_idx) for i, positions in enumerate(self.token_positions): if positions is not None: self.token_positions[i] = positions.index_select(0, beam_idx.to(positions.device)) self.valid_tokens[i] = self.valid_tokens[i].index_select(0, beam_idx.to(positions.device)) def batch_repeat_interleave(self, repeats): super().batch_repeat_interleave(repeats) for i, positions in enumerate(self.token_positions): if positions is not None: self.token_positions[i] = positions.repeat_interleave(repeats, dim=0) self.valid_tokens[i] = self.valid_tokens[i].repeat_interleave(repeats, dim=0) def batch_select_indices(self, indices): self.reorder_cache(indices) def crop(self, max_length): raise NotImplementedError("QMoD packed cache does not support assisted-decoding rollback") def reset(self): # DynamicLayer.reset retains storage, which would desynchronize metadata. for layer in self.layers: if layer.is_initialized: layer.keys = layer.keys[..., :0, :] layer.values = layer.values[..., :0, :] self.seen_tokens = 0 self.token_positions = [None] * len(self.layers) self.valid_tokens = [None] * len(self.layers) def gather_tokens(x, indices): return x.gather(1, indices[..., None].expand(-1, -1, x.shape[-1])) class QMoDDecoderLayer(Q50MDecoderLayer): def __init__(self, config, layer_idx): block_config = copy.deepcopy(config) self.routed = layer_idx in config.mod_layers if self.routed: block_config.intermediate_size = config.mod_intermediate_size super().__init__(block_config, layer_idx) self.config = config self.layer_idx = layer_idx self.router = nn.Linear(config.hidden_size, 1, bias=False) if self.routed else None def route(self, hidden_states, valid): # FP32 router and BCE, also when the block runs in BF16 autocast. with torch.autocast(device_type=hidden_states.device.type, enabled=False): logits = F.linear(hidden_states.float(), self.router.weight.float()).squeeze(-1) if self.training: # Deterministic tie breaking preserves checkpoint replay. Capacity is # per valid sequence, never per batch and never allocated to padding. counts = torch.ceil(valid.sum(-1) * self.config.mod_capacity).long() order = logits.detach().masked_fill(~valid, -torch.inf).argsort(dim=-1, descending=True, stable=True) ranks = torch.empty_like(order).scatter(1, order, torch.arange(order.shape[1], device=order.device).expand_as(order)) selected = (ranks < counts[:, None]) & valid bce = F.binary_cross_entropy_with_logits(logits, selected.float(), reduction="none") aux = (bce * valid).sum() / valid.sum().clamp_min(1) else: # No top-k, capacity clipping or sequence normalization in eval: # changing future tokens must not alter any prefix routing decision. selected = (logits > 0) & valid aux = logits.new_zeros(()) return logits, selected, aux def forward(self, hidden_states, valid, token_positions, position_embeddings, past_key_values=None): batch, seq, _ = hidden_states.shape zero = hidden_states.new_zeros((), dtype=torch.float32) if self.routed: logits, selected, aux = self.route(hidden_states, valid) count = selected.sum(-1) width = int(count.max().item()) stats = torch.stack((selected.sum() / valid.sum().clamp_min(1), (((logits > 0) != selected) & valid).sum() / valid.sum().clamp_min(1))).detach() if width == 0: return hidden_states, aux, stats # Sort by ORIGINAL time, not router score. Surplus entries are dummy # slots for shorter batch members, excluded from keys and scatter. order = torch.arange(seq, device=hidden_states.device).expand(batch, -1) indices = order.masked_fill(~selected, seq).sort(-1).values[:, :width].clamp_max(seq - 1) packed_valid = torch.arange(width, device=hidden_states.device)[None, :] < count[:, None] packed = gather_tokens(hidden_states, indices) positions = token_positions.gather(1, indices) rotary = tuple(gather_tokens(x, indices) for x in position_embeddings) else: aux, stats = zero, torch.stack((zero, zero)) packed, packed_valid, positions, rotary = hidden_states, valid, token_positions, position_embeddings key_positions, key_valid = positions, packed_valid if past_key_values is not None: key_positions, key_valid = past_key_values.append_metadata(self.layer_idx, positions, packed_valid) # The common unpadded training path can use fused SDPA's causal flag. mask = None if not (past_key_values is None and self.config._attn_implementation == "sdpa" and bool(packed_valid.all())): allowed = (key_positions[:, None, :] <= positions[:, :, None]) & key_valid[:, None, :] mask = packed.new_zeros((batch, 1, packed.shape[1], key_positions.shape[1])) mask.masked_fill_(~allowed[:, None], torch.finfo(packed.dtype).min) transformed = super().forward(packed, attention_mask=mask, position_embeddings=rotary, past_key_values=past_key_values) if not self.routed: return torch.where(valid[..., None], transformed, hidden_states), aux, stats # Weight the residual UPDATE, never the identity/skip path. Raw scalar # weights follow equation (1); sigmoid is only the BCE interpretation. update = (transformed - packed) * logits.gather(1, indices)[..., None].to(packed.dtype) update = update * packed_valid[..., None] result = hidden_states.scatter_add(1, indices[..., None].expand_as(update), update) return result, aux, stats class QMoDModel(Q50MPreTrainedModel): config_class = QMoDConfig _no_split_modules = ["QMoDDecoderLayer"] def __init__(self, config): super().__init__(config) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id) self.layers = nn.ModuleList([QMoDDecoderLayer(config, i) for i in range(config.num_hidden_layers)]) self.norm = Q50MRMSNorm(config.hidden_size, config.rms_norm_eps) self.rotary_emb = Q50MRotaryEmbedding(config) self.gradient_checkpointing = False self.router_loss = None self.last_router_metrics = {} self.post_init() def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs): if (input_ids is None) == (inputs_embeds is None): raise ValueError("Specify exactly one of input_ids or inputs_embeds") if self.training and (use_cache or past_key_values is not None): raise ValueError("MoD top-k training does not support a KV cache; use model.eval() for inference") if past_key_values is not None and not isinstance(past_key_values, QMoDCache): raise ValueError("QMoD requires its own packed QMoDCache") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) batch, seq, _ = inputs_embeds.shape if seq == 0: raise ValueError("Input sequence must not be empty") if use_cache and past_key_values is None: past_key_values = QMoDCache(self.config) past = 0 if past_key_values is None else past_key_values.get_seq_length() positions = torch.arange(past, past + seq, device=inputs_embeds.device).expand(batch, -1) valid = torch.ones((batch, seq), dtype=torch.bool, device=inputs_embeds.device) if attention_mask is not None: if attention_mask.ndim != 2 or attention_mask.shape != (batch, past + seq): raise ValueError("QMoD expects a 2D attention_mask covering the full cached + current sequence") valid = attention_mask[:, -seq:].to(device=inputs_embeds.device, dtype=torch.bool) if position_ids is None: position_ids = positions if attention_mask is None else attention_mask.long().cumsum(-1)[:, -seq:] - 1 position_ids = position_ids.clamp_min(0) position_ids = position_ids.to(inputs_embeds.device).expand(batch, seq) rotary = self.rotary_emb(inputs_embeds, position_ids) hidden = inputs_embeds losses = [] self.router_loss = None self.last_router_metrics = {} for i, layer in enumerate(self.layers): if self.gradient_checkpointing and self.training: hidden, aux, stats = self._gradient_checkpointing_func( layer.forward, hidden, valid, positions, rotary) else: hidden, aux, stats = layer(hidden, valid, positions, rotary, past_key_values=past_key_values) if layer.routed: losses.append(aux) self.last_router_metrics[f"mod/layer_{i}/selected_fraction"] = stats[0] self.last_router_metrics[f"mod/layer_{i}/threshold_disagreement"] = stats[1] if self.training: self.router_loss = torch.stack(losses).mean() if losses else hidden.sum() * 0 if past_key_values is not None: past_key_values.seen_tokens += seq return BaseModelOutputWithPast(last_hidden_state=self.norm(hidden), past_key_values=past_key_values) class QMoDForCausalLM(Q50MForCausalLM): config_class = QMoDConfig _no_split_modules = ["QMoDDecoderLayer"] def __init__(self, config): Q50MPreTrainedModel.__init__(self, config) self.model = QMoDModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() @classmethod def _supports_default_dynamic_cache(cls): return False def forward(self, *args, **kwargs): out = super().forward(*args, **kwargs) if out.loss is not None and self.training: out.loss = out.loss + self.consume_router_loss() return out def consume_router_loss(self, metrics=None): aux = self.model.router_loss if aux is None: raise RuntimeError("Consume router loss once after a training forward") self.model.router_loss = None if metrics is not None: metrics["train/mod_router_bce_raw"] = aux.detach() metrics["train/mod_router_bce_weighted"] = aux.detach() * self.config.router_aux_coef return aux * self.config.router_aux_coef def router_metrics(self): return self.model.last_router_metrics.copy() def parameter_budget(self): total = sum(p.numel() for p in self.parameters()) conditional = sum(p.numel() for layer in self.model.layers if layer.routed for name, p in layer.named_parameters() if not name.startswith("router.")) return {"total": total, "always_active": total - conditional, "expected_active_training": total - conditional + self.config.mod_capacity * conditional, "active_min": total - conditional, "active_max": total, "capacity": self.config.mod_capacity}