| """Padding and inference helpers for caller-batch-independent verification. |
| |
| Padding masks and no-drop routing settings address semantic batch dependence. |
| They do not guarantee bitwise equality: kernels can round differently for |
| changed shapes. Use ``strict_deterministic_inference`` when fixed per-example |
| shapes matter, with an eval model whose kernels and routing are deterministic. |
| """ |
| from __future__ import annotations |
|
|
| import math |
| from numbers import Real |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| def _validate_sequences(sequences: list[torch.Tensor], pad_token_id: int) -> None: |
| if not isinstance(sequences, list): |
| raise TypeError("sequences must be a list of tensors") |
| if not sequences: |
| raise ValueError("sequences must contain at least one nonempty tensor") |
| if isinstance(pad_token_id, bool) or not isinstance(pad_token_id, int): |
| raise TypeError("pad_token_id must be an integer") |
| for sequence in sequences: |
| if not isinstance(sequence, torch.Tensor): |
| raise TypeError("each sequence must be a tensor") |
| if sequence.ndim != 1 or sequence.numel() == 0: |
| raise ValueError("each sequence must be a nonempty 1D tensor") |
| if sequence.dtype not in (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64): |
| raise TypeError("token sequences must have an integer dtype") |
| if sequence.dtype != sequences[0].dtype or sequence.device != sequences[0].device: |
| raise ValueError("all sequences must have the same dtype and device") |
| limits = torch.iinfo(sequences[0].dtype) |
| if not limits.min <= pad_token_id <= limits.max: |
| raise ValueError("pad_token_id must be representable in the sequence dtype") |
|
|
|
|
| def prepare_batch_inputs( |
| sequences: list[torch.Tensor], pad_token_id: int = 0 |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Left-pad nonempty 1D integer sequences; return IDs and an int64 mask. |
| |
| Input dtype/device are preserved in IDs, and masks share their device. |
| Masks are length-derived, so genuine tokens equal to ``pad_token_id`` remain |
| visible. For positional models derive positions as |
| ``(attention_mask.cumsum(-1) - 1).clamp_min(0)`` and pass them explicitly. |
| """ |
| _validate_sequences(sequences, pad_token_id) |
| width = max(sequence.numel() for sequence in sequences) |
| ids = sequences[0].new_full((len(sequences), width), pad_token_id) |
| mask = torch.zeros(ids.shape, dtype=torch.long, device=ids.device) |
| for row, sequence in enumerate(sequences): |
| ids[row, -sequence.numel():] = sequence |
| mask[row, -sequence.numel():] = 1 |
| return ids, mask |
|
|
|
|
| def enforce_moe_deterministic_routing(model_or_config: object) -> tuple[str, ...]: |
| """Configure exposed no-drop controls, returning the touched attribute paths. |
| |
| Supports a config object or an ``nn.Module`` with ``drop_tokens`` on itself, |
| its submodules, or their ``config`` objects. Sets boolean ``drop_tokens`` to |
| False. On that *same object* only, raises numeric ``capacity_factor``, |
| ``train_capacity_factor``, and ``eval_capacity_factor`` to at least 100; |
| zero is deliberately never used as a purported no-limit sentinel. |
| |
| This is a best-effort configuration adapter, not proof of deterministic |
| routing or unlimited capacity. A large finite factor is not an infinite |
| capacity, and setting config attributes may not update instantiated routers. |
| Verify the actual model implementation honors these controls. Unrecognized |
| routers and exposed capacity controls without a local boolean drop flag |
| raise ValueError before any planned changes. Other architectures need an |
| architecture-specific adapter. No floating-point equality is promised. |
| """ |
| candidates: list[tuple[str, object]] = [("root", model_or_config)] |
| if isinstance(model_or_config, nn.Module): |
| candidates.extend((f"root.{name}", module) for name, module in model_or_config.named_modules() if name) |
| candidates += [(f"{path}.config", obj.config) for path, obj in candidates if hasattr(obj, "config")] |
| factors = ("capacity_factor", "train_capacity_factor", "eval_capacity_factor") |
| |
| absolute_capacities = ("expert_capacity", "token_capacity", "max_tokens_per_expert") |
| updates: list[tuple[str, object, str, object]] = [] |
| seen: set[int] = set() |
| for path, obj in candidates: |
| if id(obj) in seen: |
| continue |
| seen.add(id(obj)) |
| local_factors = [name for name in factors if hasattr(obj, name)] |
| if any(hasattr(obj, name) for name in absolute_capacities): |
| raise ValueError(f"{path}: absolute capacity requires an architecture-specific adapter") |
| has_drop = hasattr(obj, "drop_tokens") |
| if not has_drop and local_factors: |
| raise ValueError(f"{path}: capacity controls require a local boolean drop_tokens flag") |
| if not has_drop: |
| continue |
| if not isinstance(obj.drop_tokens, bool): |
| raise ValueError(f"{path}.drop_tokens must be boolean") |
| updates.append((path, obj, "drop_tokens", False)) |
| for name in local_factors: |
| value = getattr(obj, name) |
| if isinstance(value, bool) or not isinstance(value, Real) or not math.isfinite(value): |
| raise ValueError(f"{path}.{name} must be a finite numeric capacity factor") |
| updates.append((path, obj, name, max(100, value))) |
| if not updates: |
| raise ValueError("unsupported router: no boolean drop_tokens control found") |
| for _, obj, name, value in updates: |
| setattr(obj, name, value) |
| return tuple(f"{path}.{name}" for path, _, name, _ in updates) |
|
|
|
|
| def strict_deterministic_inference( |
| model: nn.Module, sequences: list[torch.Tensor], pad_token_id: int = 0 |
| ) -> list[torch.Tensor]: |
| """Return per-example logits using one unpadded forward per example. |
| |
| Removes the singleton batch axis: classifier logits ``(1, classes)`` become |
| ``(classes,)`` and token logits ``(1, length, classes)`` become |
| ``(length, classes)``. The model must accept |
| ``input_ids``, ``attention_mask``, and ``position_ids`` keyword arguments and |
| return a Tensor or an object with a Tensor ``logits`` attribute. All modules |
| must already be in eval mode. No training state or global flags are changed. |
| Models needing extra sequence-aligned inputs (such as segment masks) need |
| a wrapper that supplies those inputs for each unpadded example; this helper |
| intentionally does not infer them from token values. |
| |
| Serial execution trades batching throughput for identical per-example |
| shapes regardless of caller batch composition/order. Exact repeatability |
| still requires deterministic kernels/routing, fixed parameters, no mutable |
| cross-call state or stochastic eval behavior, and the same device/runtime. |
| This helper cannot make a nondeterministic model deterministic. |
| """ |
| _validate_sequences(sequences, pad_token_id) |
| if any(module.training for module in model.modules()): |
| raise ValueError("strict inference requires every module to be in eval mode") |
| results = [] |
| with torch.inference_mode(): |
| for sequence in sequences: |
| ids, mask = prepare_batch_inputs([sequence], pad_token_id) |
| positions = (mask.cumsum(-1) - 1).clamp_min(0) |
| output = model(input_ids=ids, attention_mask=mask, position_ids=positions) |
| logits = output if isinstance(output, torch.Tensor) else getattr(output, "logits", None) |
| if not isinstance(logits, torch.Tensor) or logits.ndim < 2 or logits.shape[0] != 1: |
| raise ValueError("model must return logits with a singleton leading batch dimension") |
| results.append(logits[0].clone()) |
| return results |
|
|
