File size: 7,859 Bytes
a82e638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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 lack a portable no-limit value, even with drop_tokens.
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
|