Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _layer_census.py: 97%
717 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Census of what a layer of each kind keeps for its backward, on a fake layer.
17The census builds one decoder layer of a Transformers model under
18``FakeTensorMode``, so it allocates no memory, runs its forward and its
19backward at micro-batch 1, and counts bytes: those the autograd graph saves
20between the two passes, and the most the backward holds live at once, less
21the parameters and their gradients, which the memory model counts on its
22own. It runs the kernels the runtime runs: flash attention, which saves its
23inputs, its output and two fp32 softmax statistics per head and token, and
24HyperParallel's chunked gated-delta-rule kernel, which saves the queries,
25keys, values, the cumulated gates, beta and one chunk matrix. A layer built
26with half the heads and half the feed-forward widths, as a tensor-parallel
27rank of 2 holds it, tells the bytes tensor parallelism splits from those it
28does not. Run again under HyperParallel's selective activation checkpointing,
29its own checkpoint and the policy the installed trainer runs, a layer keeps
30its input and the outputs that policy saves: every other matmul's, the
31convolution's, and the attention kernel's where the policy saves the
32operator the census's kernel stands for (:func:`_policy_saves`); the matmuls
33it does not save its backward runs again, a share of the attention's
34projections and of the rest of the layer.
36What a layer saves is also told apart by the op of ND's op vector whose
37backward takes it (:class:`_SavedOps`), the records of shared decision S1 as
38the census fills them: the delta rule's, the attention kernel's (its
39inputs and output for the batched matmuls, its fp32 statistics for the
40softmax), a dropout's mask, what a norm saves, what the feed-forward's
41activation function saves, and the rest by the part of the layer saving it,
42the attention's projections or the feed-forward's.
44The output layer's census runs the final norm, the output projection and
45Transformers' causal-LM loss, which casts the logits to fp32, as
46HyperParallel's trainer runs them by default, dropping the model's logits
47before the backward as the trainer does; half the vocabulary tells the
48bytes a vocabulary-parallel loss splits.
50Where a train.yaml replaces Transformers' modules with HyperParallel's
51fused ones (``plan_overrides`` with ``replace.kind: module``: its RMSNorm, its
52grouped-query attention, its grouped experts), the census runs those
53instead, so what it measures is what the trainer saves
54(:func:`replacement_specs`). Their kernels are ``torch_npu``'s, which a
55host cannot call, so they run under the kernels' shape contracts
56(:mod:`hyper_parallel.auto_parallel._npu_contracts`). A factory that needs
57a library the host lacks leaves the census with a layer the run does not
58build, which it refuses (:exc:`CensusUnavailable`): Transformers' own
59modules keep other bytes, and under the selective policy a DeepSeek-V3.2
60layer's eager indexer keeps 2.3 times what the layer keeps plain. The FLOP census
61(:func:`census_flops`) and the shares of matmul FLOPs a selective layer
62recomputes (:func:`census_recomputed`) stay on Transformers' own modules:
63both count the arithmetic the time model prices, which fusing a module
64does not change, and a fused kernel runs matmuls no dispatch sees.
66A model whose layers cannot run alone is measured over its whole stack
67(:func:`census_stack`): DeepSeek-V4.1's layers read what earlier layers
68wrote in the same forward, and carry a hyper-connection dimension and a
69mixing state between them. The model the trainer builds runs its own forward
70and backward once, each save counting toward the layer that makes it, and a
71kind's record is the mean of its layers'.
73A layer's parameters are counted by part, as ND prices its parts: the
74attention's, the norms', the dense feed-forward's, the routed experts', the
75shared expert's and the router's (:func:`census_parameters`), for verify
76mode to set beside what ND prices. So are its forward FLOPs
77(:func:`census_flops`): each matmul's toward the part of the layer that
78runs it, flash attention's toward the scores, at every pair of tokens, as
79the runtime's kernel computes them under an explicit causal mask, and the
80gated delta rule's through Transformers' own chunked implementation, whose
81matmuls the runtime's kernel runs alike.
82"""
83from __future__ import annotations
85import contextlib
86import copy
87import functools
88import importlib
89import inspect
90import itertools
91import logging
92import weakref
93from typing import (
94 Any,
95 Callable,
96 Dict,
97 FrozenSet,
98 Iterable,
99 Iterator,
100 List,
101 Mapping,
102 NamedTuple,
103 Optional,
104 Sequence,
105 Tuple,
106)
108import torch # pylint: disable=forbidden-backend-import
109from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode # pylint: disable=forbidden-backend-import
110from torch.utils._python_dispatch import TorchDispatchMode # pylint: disable=forbidden-backend-import
111from torch.utils._pytree import tree_flatten # pylint: disable=forbidden-backend-import
112from torch.utils.checkpoint import DefaultDeviceType # pylint: disable=forbidden-backend-import
114from hyper_parallel.auto_parallel._model_spec import KindActivations
115from hyper_parallel.auto_parallel._npu_contracts import npu_contracts
116from hyper_parallel.core.activation_memory import api as activation_memory
117from hyper_parallel.core.activation_memory.policy import CheckpointPolicy
119logger = logging.getLogger(__name__)
122class CensusUnavailable(RuntimeError):
123 """A census cannot run here: this host cannot build the modules the run installs."""
126class LayerTraffic(NamedTuple):
127 """What one forward of a layer moves, by part: its bytes, the parameters among them, and its ops.
129 Bytes of the whole layer on a micro-batch of one sequence, as
130 :func:`census_flops` counts its FLOPs, so that the two quantities can be
131 set beside each other part by part. ``parameter`` is the share of
132 ``moved`` that is the layer's own weights, which a strategy shards where
133 it leaves the activations alone, and ``launches`` is how many ops the
134 forward dispatches, views excluded.
135 """
137 moved: Dict[str, float]
138 parameter: Dict[str, float]
139 launches: Dict[str, int]
141# The Transformers attention implementation the census registers its flash
142# attention under: a name Transformers does not take for a flash
143# attention's, as it does not take HyperParallel's default, sdpa, so a model
144# runs sdpa's path, and DeepSeek-V3's values are not padded to its queries'
145# width.
146_FLASH = "nd_census_attention"
148# The config fields a tensor-parallel rank holds a share of.
149_TP_FIELDS = (
150 "num_attention_heads", "num_key_value_heads", "intermediate_size", "moe_intermediate_size",
151 "shared_expert_intermediate_size", "linear_num_key_heads", "linear_num_value_heads",
152)
154# The softmax statistics flash attention keeps per head and token.
155_FLASH_STATS = 8
157# The gated delta rules a run can take, named as the runtime names them
158# (``components/modules/gated_delta_net.py``): Transformers' own chunked
159# implementation, and HyperParallel's kernel, whose saved set the census holds
160# as a contract. The runtime defaults to eager and so does the census.
161_EAGER_GDN = "eager"
162_HYPER_GDN = "triton"
163_GDN_BACKENDS = (_EAGER_GDN, _HYPER_GDN)
165# The census's attention kernels, each beside the runtime operator it stands
166# in for, named as HyperParallel's selective policy names the operators it
167# saves: its own flash attention for PyTorch's scaled-dot-product attention,
168# and the contract of the kernel HyperParallel's fused attention calls,
169# ``torch_npu.npu_fusion_attention``. A selective census saves what a kernel
170# returns exactly where the installed policy saves what its operator returns
171# (:func:`_policy_saves`).
172_ATTENTION_KERNELS = {
173 "nd_census::flash_attention": "aten.scaled_dot_product_attention",
174 "nd_census_npu::fusion_attention": "npu.npu_fusion_attention",
175}
177# The per-head query and key norms an attention holds, which ND prices with
178# the layer's norms.
179_QK_NORMS = ("q_norm", "k_norm")
182def _flash_outputs(query: torch.Tensor, key: torch.Tensor,
183 value: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
184 """Flash attention's outputs, shapes only: each head's, as wide as its values, and two fp32 statistics."""
185 del key
186 batch, heads, seq, _ = query.shape
187 stats = [torch.empty(batch, heads, seq, _FLASH_STATS, dtype=torch.float32, device=query.device)
188 for _ in range(2)]
189 return query.new_empty(batch, heads, seq, value.shape[-1]), stats[0], stats[1]
192# Flash attention as one op, as the runtime's kernel is: HyperParallel's
193# selective policy saves what the attention kernels it lists return, and
194# saves this one's too.
195_FLASH_KERNEL = torch.library.custom_op(
196 "nd_census::flash_attention", mutates_args=(),
197 schema="(Tensor query, Tensor key, Tensor value) -> (Tensor, Tensor, Tensor)")(_flash_outputs)
198_FLASH_KERNEL.register_fake(_flash_outputs)
201def _flash_saves(ctx: Any, inputs: Tuple[torch.Tensor, ...], output: Tuple[torch.Tensor, ...]) -> None:
202 """Keep the inputs, the output and the statistics, which no gradient reaches."""
203 ctx.mark_non_differentiable(*output[1:])
204 ctx.save_for_backward(*inputs, *output)
207def _flash_backward(ctx: Any, *grads: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
208 """Gradients of the inputs' shapes."""
209 del grads
210 query, key, value = ctx.saved_tensors[:3]
211 return torch.empty_like(query), torch.empty_like(key), torch.empty_like(value)
214_FLASH_KERNEL.register_autograd(_flash_backward, setup_context=_flash_saves)
217def _flash_attention(module, query, key, value, attention_mask, **kwargs): # pylint: disable=unused-argument
218 """A Transformers attention function running the census's flash attention, grouped K and V unexpanded."""
219 return _FLASH_KERNEL(query, key, value)[0].transpose(1, 2).contiguous(), None
222class _GatedDeltaRule(torch.autograd.Function):
223 """HyperParallel's chunked gated-delta-rule kernel's saved set, shapes only."""
225 @staticmethod
226 def forward(ctx: Any, query: torch.Tensor, key: torch.Tensor, # pylint: disable=arguments-differ
227 value: torch.Tensor, gate: torch.Tensor, beta: torch.Tensor, chunk: int) -> torch.Tensor:
228 """Save the inputs, the cumulated gates, beta and a chunk matrix; return an output of the values' shape."""
229 batch, seq, heads, _ = key.shape
230 gates = torch.empty(batch, seq, heads, dtype=torch.float32, device=query.device)
231 chunks = torch.empty(batch, seq, heads, chunk, dtype=key.dtype, device=query.device)
232 ctx.save_for_backward(query, key, value, gates, beta, chunks)
233 del gate
234 return torch.empty_like(value)
236 @staticmethod
237 def backward(ctx: Any, # pylint: disable=arguments-differ
238 grad: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]:
239 """Return gradients of the inputs' shapes, and none of the chunk size."""
240 del grad
241 query, key, value, gates, beta, _ = ctx.saved_tensors
242 return (torch.empty_like(query), torch.empty_like(key), torch.empty_like(value),
243 torch.empty_like(gates), torch.empty_like(beta), None)
246def _gated_delta_rule(modeling: Any) -> Callable[..., Tuple[torch.Tensor, None]]:
247 """The drop-in for a Transformers ``chunk_gated_delta_rule`` running :class:`_GatedDeltaRule`."""
248 l2norm = getattr(modeling, "l2norm", None)
250 def run(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, g: torch.Tensor, beta: torch.Tensor,
251 chunk_size: int = 64, use_qk_l2norm_in_kernel: bool = False,
252 **_kwargs: Any) -> Tuple[torch.Tensor, None]:
253 """Run the kernel's contract as Transformers calls its kernel, with no final state."""
254 if use_qk_l2norm_in_kernel and l2norm is not None:
255 query, key = l2norm(query, dim=-1, eps=1e-6), l2norm(key, dim=-1, eps=1e-6)
256 return _GatedDeltaRule.apply(query, key, value, g, beta, chunk_size), None
258 return run
261# Reading a scalar off a fake tensor raises, since it holds no data. Transformers'
262# grouped_mm fallback reads its group offsets so in its backward, and it takes that
263# fallback on CPU with torch 2.8 or older, the 910C's 2.7.1 among them.
264_SCALAR_READ = torch.ops.aten._local_scalar_dense.default # pylint: disable=protected-access
267class _LiveBytes(TorchDispatchMode):
268 """The storages the ops it sees allocate, and the bytes they hold while they live."""
270 def __init__(self) -> None:
271 """Start with no storage counted."""
272 super().__init__()
273 self.sizes: Dict[Tuple[int, int], int] = {}
274 self.generation: Dict[int, int] = {}
275 self.events: list = []
277 def key(self, tensor: torch.Tensor) -> Tuple[int, int]:
278 """The storage behind *tensor*, in the lifetime it has now."""
279 address = tensor.untyped_storage()._cdata # pylint: disable=protected-access
280 return address, self.generation.get(address, 0)
282 def track(self, tensor: torch.Tensor) -> None:
283 """Count *tensor*'s storage from now until it is freed."""
284 storage = tensor.untyped_storage()
285 key = self.key(tensor)
286 if key in self.sizes:
287 return
288 self.sizes[key] = storage.nbytes()
289 self.events.append((key, self.sizes[key]))
290 weakref.finalize(storage, self._free, key)
292 def _free(self, key: Tuple[int, int]) -> None:
293 size = self.sizes.pop(key, None)
294 self.generation[key[0]] = key[1] + 1
295 if size:
296 self.events.append((key, -size))
298 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
299 if func is _SCALAR_READ and isinstance(args[0], FakeTensor):
300 return 0 # shapes alone price the memory, whatever the scalar
301 out = func(*args, **(kwargs or {}))
302 for tensor in tree_flatten(out)[0]:
303 if isinstance(tensor, torch.Tensor):
304 self.track(tensor)
305 return out
307 def peak(self, start: int, excluded: Iterable[Tuple[int, int]], apart: Iterable[Tuple[int, int]] = ()) -> int:
308 """The bytes live when the most are, from event *start* on, leaving out the *excluded* storages.
310 The storages *apart* count toward finding that moment, not toward
311 the bytes returned.
312 """
313 excluded, apart = set(excluded), set(apart)
314 live = held = 0
315 for key, delta in self.events[:start]:
316 if key not in excluded:
317 live += delta
318 held += delta if key in apart else 0
319 top, bytes_at_top = live, live - held
320 for key, delta in self.events[start:]:
321 if key in excluded:
322 continue
323 live += delta
324 held += delta if key in apart else 0
325 if live > top:
326 top, bytes_at_top = live, live - held
327 return bytes_at_top
330@contextlib.contextmanager
331def _checkpoints_on_cpu() -> Iterator[None]:
332 """Run activation checkpointing as a CPU's while a census runs.
334 A census's tensors are fake and on the CPU, so a checkpoint over them
335 takes the default device type, which ``torch_npu`` sets to its own when
336 it imports. Replaying the forward would then open the NPU's autocast
337 and start a device the census never uses, and that a host whose NPU
338 cannot start fails on. The default is put back on the way out.
339 """
340 previous = DefaultDeviceType.get_device_type()
341 DefaultDeviceType.set_device_type("cpu")
342 try:
343 yield
344 finally:
345 DefaultDeviceType.set_device_type(previous)
348def _measure(params: Sequence[torch.Tensor], inputs: Iterable[torch.Tensor],
349 forward: Callable[[], torch.Tensor], backward: Callable[[torch.Tensor], None],
350 grad_inputs: Iterable[torch.Tensor] = (), checkpointed: bool = False,
351 shared: Iterable[torch.Tensor] = (), ops: Optional["_SavedOps"] = None) -> Tuple[int, int]:
352 """Bytes a forward saves for its backward, and the activations the backward holds when it holds the most.
354 Args:
355 params: The parameters of what runs.
356 inputs: The tensors the forward takes, which count from its start.
357 forward: Runs the forward and returns its output.
358 backward: Runs the backward from that output.
359 grad_inputs: The tensors the backward takes, which count from its
360 start.
361 checkpointed: Whether the forward runs under activation
362 checkpointing, whose own hooks hold what it keeps.
363 shared: Tensors the model shares between its layers, which a
364 checkpointed forward's count leaves out.
365 ops: Where given, tells each storage the forward saves apart by
366 the op that first saves it, counting its bytes into
367 ``ops.saved``, which then sums to the bytes returned.
369 Returns:
370 The bytes the forward saves and still holds once it has run, its
371 inputs included, and the bytes of activations live when the
372 backward peaks, its output left out. A graph the forward drops
373 before it ends, as a selection that keeps only the indices it picks
374 drops its scores, is freed then in a run and is no backward's. A
375 checkpointed forward keeps what is live once it has run, but for
376 the *shared* tensors: its inputs and the outputs its policy saves.
377 The parameters, which the forward's views of them bring into the
378 tracker, are the memory model's to count, and so are their
379 gradients, which hooks on the parameters tell apart from the
380 activations' however they are shaped: the gradients the backward
381 hands each parameter and the one it keeps. They settle when the
382 backward peaks, and the bytes returned are the activations held
383 then.
384 """
385 stored = {p.untyped_storage()._cdata for p in params} # pylint: disable=protected-access
386 # Counted by the storage in the lifetime it has now, as the tracker counts
387 # them: a fake tensor's storage is cheap and short lived, so one the
388 # forward has freed can hand its address to a later save, and counting by
389 # address alone lost whichever of the two it saw first. Which pair
390 # collided depended on what had run before, so a layer did not measure the
391 # same twice.
392 saved: Dict[Tuple[int, int], int] = {}
393 saved_op: Dict[Tuple[int, int], str] = {}
395 def pack(tensor: torch.Tensor) -> torch.Tensor:
396 """Count a tensor autograd saves, unless it is a parameter's, toward the op that first saves it."""
397 # Named before the storage is read, which dispatches ops of its own.
398 op = ops.op(tensor) if ops is not None else None
399 address = tensor.untyped_storage()._cdata # pylint: disable=protected-access
400 if address not in stored:
401 key = live.key(tensor)
402 if op is not None:
403 saved_op.setdefault(key, op)
404 saved[key] = tensor.untyped_storage().nbytes()
405 # Detached, so a graph holds the bytes it saves and not the tensor: an
406 # op saving its own output would otherwise hold itself through that
407 # output's grad_fn, a cycle no collector sees, which kept every graph
408 # a forward drops alive into the backward. On the DeepSeek-V4.1 crop
409 # at 4096 tokens that was the Indexer selection's scores, 8.9 GiB a
410 # rank a run frees as soon as the indices are taken.
411 return tensor.detach()
413 live = _LiveBytes()
414 grads: List[Tuple[int, int]] = []
415 handles = [p.register_hook(lambda grad: grads.append(live.key(grad))) for p in params]
416 try:
417 with _checkpoints_on_cpu(), live, ops or contextlib.nullcontext(), \
418 torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor):
419 for tensor in inputs:
420 live.track(tensor)
421 out = forward()
422 start = len(live.events)
423 left_out = {live.key(out), *(live.key(p) for p in params)}
424 not_kept = left_out | {live.key(tensor) for tensor in shared}
425 kept = sum(size for key, size in live.sizes.items() if key not in not_kept)
426 held = {key: size for key, size in saved.items() if key in live.sizes}
427 for tensor in grad_inputs:
428 live.track(tensor)
429 backward(out)
430 finally:
431 for handle in handles:
432 handle.remove()
433 # A storage's bytes count toward the op that first saved it, read from the
434 # same record as the total, so the two agree whatever a fake tensor does
435 # with a storage address it has freed.
436 for address, size in held.items():
437 op = saved_op.get(address)
438 if op is not None:
439 ops.saved[op] = ops.saved.get(op, 0) + size
440 grads += [live.key(p.grad) for p in params if p.grad is not None]
441 return kept if checkpointed else sum(held.values()), live.peak(start, left_out, grads)
444# The ops a matmul runs as.
445_MATMULS = ("mm", "addmm", "bmm", "_grouped_mm", "baddbmm")
448class _SavedOps(TorchDispatchMode):
449 """What each tensor a layer's forward saves is saved for: the op of ND's op vector whose backward takes it.
451 The delta rule's saves are its own (``linrec``); the attention kernel's
452 are its batched matmuls' (``attBMM``), its fp32 statistics the
453 softmax's; a dropout's mask is its own; what a norm saves, the norm's
454 (``normOp``), and what the feed-forward's activation function saves,
455 the function's (``ffAct``). The rest counts toward the part of the
456 layer saving it: the attention's projections (``attMM``), the
457 feed-forward's (``ffMM``), or ``other`` for the layer's own code.
458 """
460 def __init__(self, layer: Any) -> None:
461 """Follow which of *layer*'s modules runs, and when its delta rule does."""
462 super().__init__()
463 self.saved: Dict[str, int] = {}
464 self.running: List[str] = []
465 self.func: Any = None
466 self.delta_rule = False
467 self.modules = dict(layer.named_modules())
468 for name, module in self.modules.items():
469 if name:
470 module.register_forward_pre_hook(functools.partial(self._enter, name))
471 module.register_forward_hook(self._leave)
472 if hasattr(module, "chunk_gated_delta_rule"):
473 module.chunk_gated_delta_rule = self._rule(module.chunk_gated_delta_rule)
475 def _rule(self, rule: Callable[..., Any]) -> Callable[..., Any]:
476 """*rule*, noting that it runs."""
477 def run(*args: Any, **kwargs: Any) -> Any:
478 """The delta rule, its saves told apart as its own while it runs."""
479 self.delta_rule = True
480 try:
481 return rule(*args, **kwargs)
482 finally:
483 self.delta_rule = False
484 return run
486 def _enter(self, name: str, *_: Any) -> None:
487 """Note that the module at *name* runs."""
488 self.running.append(name)
490 def _leave(self, *_: Any) -> None:
491 """Note that the innermost module running has returned."""
492 self.running.pop()
494 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
495 # A tensor's metadata is read through prims, and _measure detaches
496 # each save, both between an op and its saves.
497 if func.namespace != "prim" and func is not torch.ops.aten.detach.default:
498 self.func = func
499 return func(*args, **(kwargs or {}))
501 def op(self, tensor: torch.Tensor) -> str:
502 """The op that saves *tensor*, which the forward saves now."""
503 if self.delta_rule:
504 return "linrec"
505 if self.func in _attention_ops():
506 return "softmax" if tensor.dtype == torch.float32 else "attBMM"
507 if self.func is torch.ops.aten.native_dropout.default:
508 return "dropout"
509 if self.func is torch.ops.nd_census_npu.rms_norm.default:
510 return "normOp"
511 if self.func is torch.ops.nd_census_npu.swiglu.default:
512 return "ffAct"
513 name = self.running[-1] if self.running else ""
514 module = self.modules.get(name)
515 if name and "norm" in type(module).__name__.lower():
516 return "normOp"
517 part = name.split(".", maxsplit=1)[0]
518 if not part:
519 return "other"
520 attention = "attn" in part or "attention" in part
521 if not attention and type(module).__module__ in ("transformers.activations", "torch.nn.modules.activation"):
522 return "ffAct"
523 return "attMM" if attention else "ffMM"
526def _matmul_flops(func: Any, args: Sequence[Any]) -> int:
527 """The FLOPs of *func* on *args* where it is a matmul, and 0 otherwise."""
528 matmuls = (torch.ops.aten.mm.default, torch.ops.aten.addmm.default, torch.ops.aten.bmm.default)
529 grouped = getattr(torch.ops.aten, "_grouped_mm", None) # first shipped in torch 2.8
530 if func not in (matmuls if grouped is None else (*matmuls, grouped.default)):
531 return 0
532 first, second = args[1:3] if func is torch.ops.aten.addmm.default else args[:2]
533 return 2 * first.numel() * second.shape[-1]
536class _RecomputedMatmuls:
537 """The matmul FLOPs a layer's selective forward saves and leaves to run again, its attention's and the rest's."""
539 def __init__(self, layer: Any) -> None:
540 """Tell the attention's parameters, those of a child named for it, from the rest of *layer*'s."""
541 self.attention = {param.untyped_storage()._cdata # pylint: disable=protected-access
542 for name, param in layer.named_parameters()
543 if "attn" in name.split(".")[0] or "attention" in name.split(".")[0]}
544 self.flops: Dict[Tuple[bool, bool], int] = {}
546 def note(self, func: Any, args: Sequence[Any], policy: CheckpointPolicy) -> None:
547 """Count the FLOPs of a matmul the policy decided on, by part and by whether it runs again."""
548 flops = _matmul_flops(func, args)
549 if flops:
550 attention = any(isinstance(arg, torch.Tensor)
551 and arg.untyped_storage()._cdata in self.attention # pylint: disable=protected-access
552 for arg in args)
553 key = (attention, policy == CheckpointPolicy.MUST_RECOMPUTE)
554 self.flops[key] = self.flops.get(key, 0) + flops
556 def share(self, attention: bool) -> float:
557 """The share of the attention's matmul FLOPs, or of the rest's, that runs again."""
558 again = self.flops.get((attention, True), 0)
559 total = again + self.flops.get((attention, False), 0)
560 return again / total if total else 0.0
563class _ForwardContext(NamedTuple):
564 """What a selective policy reads of its checkpoint context, in the forward: no recompute is running."""
566 is_recompute: bool = False
569def _checkpointing() -> Any:
570 """HyperParallel's activation checkpointing module, which holds the policy its trainer runs.
572 Its distributed package takes as long to import as the rest of the cost
573 model, so it loads only when a census runs.
574 """
575 return importlib.import_module("hyper_parallel.distributed.activation_checkpoint")
578def _policy_saves(checkpointing: Any, name: str) -> bool:
579 """Whether the installed selective policy saves what the runtime operator *name* returns.
581 Where this host registers the operator, the policy decides on it as it
582 does in a run. An NPU operator is registered only where ``torch_npu``
583 is installed: elsewhere the policy's own list of the operators it saves
584 is read for its name, which is what a run on the NPU resolves.
585 """
586 # pylint: disable=protected-access
587 operator = checkpointing._resolve_torch_op(name)
588 if operator is None:
589 return name in checkpointing._SELECTIVE_AC_COMPUTE_OP_NAMES
590 decide = checkpointing._make_selective_checkpoint_policy_fn()
591 return decide(_ForwardContext(), operator) == CheckpointPolicy.MUST_SAVE
594def _census_policy(ledger: Optional[_RecomputedMatmuls] = None) -> Callable[..., CheckpointPolicy]:
595 """The policy the installed trainer runs, deciding on the census's attention kernels as on their operators.
597 A census kernel is saved where the policy saves the operator it stands
598 for (:data:`_ATTENTION_KERNELS`), and recomputed where it does not, as
599 the policy recomputes any operator it does not name. Its forward
600 decisions on matmuls go to *ledger*, if given.
601 """
602 checkpointing = _checkpointing()
603 policy = checkpointing._make_selective_checkpoint_policy_fn() # pylint: disable=protected-access
604 saved = frozenset(_kernel_op(kernel) for kernel, name in _ATTENTION_KERNELS.items()
605 if _policy_saves(checkpointing, name))
607 def census_policy(ctx: Any, func: Any, *args: Any, **kwargs: Any) -> CheckpointPolicy:
608 """The trainer's policy, which saves a census attention kernel where it saves the kernel's operator."""
609 decision = CheckpointPolicy.MUST_SAVE if func in saved else policy(ctx, func, *args, **kwargs)
610 if ledger is not None and not ctx.is_recompute:
611 ledger.note(func, args, decision)
612 return decision
614 return census_policy
617def _selective_contexts(ledger: Optional[_RecomputedMatmuls] = None) -> Tuple[Any, Any]:
618 """HyperParallel's selective checkpointing contexts under :func:`_census_policy`."""
619 return activation_memory.create_selective_checkpoint_contexts(_census_policy(ledger))
622def _kernel_op(name: str) -> Any:
623 """The census kernel *name*, ``namespace::kernel``, as an op."""
624 namespace, kernel = name.split("::")
625 return getattr(getattr(torch.ops, namespace), kernel).default
628@functools.lru_cache(maxsize=None)
629def _attention_ops() -> FrozenSet[Any]:
630 """The census's attention kernels, as ops."""
631 return frozenset(_kernel_op(name) for name in _ATTENTION_KERNELS)
634def _import(path: str, where: str) -> Any:
635 """The object a dotted *path* names.
637 Raises:
638 ValueError: If the path names no module or no attribute of one.
639 """
640 module, _, name = str(path).rpartition(".")
641 try:
642 return getattr(importlib.import_module(module), name)
643 except (ImportError, AttributeError, ValueError) as exc:
644 raise ValueError(f"{where}: cannot import {path!r}: {exc}") from exc
647def replacement_specs(plan_overrides: Any) -> Tuple[Any, ...]:
648 """The module replacements a train.yaml's ``plan_overrides`` install, as replacement rules.
650 Args:
651 plan_overrides: The ``plan_overrides`` list, as plain mappings.
653 Returns:
654 One :class:`~hyper_parallel.models.replacement.ModuleReplacementSpec`
655 per entry that states ``replace.kind: module``, in the order stated; the
656 entries that state a sharding action alone, or a replacement
657 conditioned on a run the census does not price (``when``), are left
658 out.
660 Raises:
661 ValueError: If a replacement names a module or factory that cannot be
662 imported, or uses a removed replacement field.
663 """
664 from hyper_parallel.components.modules.family import resolve_module_family_target # pylint: disable=C0415
665 from hyper_parallel.models.replacement import ModuleReplacementSpec, module_replacement # pylint: disable=C0415
667 specs = []
668 for index, entry in enumerate(plan_overrides or ()):
669 where = f"plan_overrides[{index}]"
670 if not isinstance(entry, Mapping):
671 continue
672 removed = {
673 "replace_module", "module_type", "exact_type", "local_compute_fn",
674 "inner_wrapper", "inner_target", "inner_out_src",
675 }
676 if removed.intersection(entry):
677 raise ValueError(f"{where}: old replacement fields are no longer supported")
678 action = entry.get("replace")
679 if not isinstance(action, Mapping) or action.get("kind") != "module" or entry.get("when") is not None:
680 continue
681 if "implementation" not in action:
682 raise ValueError(f"{where}.replace requires implementation")
683 match = entry.get("match")
684 patterns = (match,) if isinstance(match, str) else tuple(match or ())
685 implementation = action["implementation"]
686 target_path = implementation.get("_target_") if isinstance(implementation, Mapping) else implementation
687 requested = _import(target_path, where)
688 args = {key: value for key, value in implementation.items() if key != "_target_"} if isinstance(
689 implementation, Mapping) else {}
690 selected, args = resolve_module_family_target(requested, args)
691 if args:
692 @module_replacement
693 def factory(
694 *, module: torch.nn.Module, module_fqn: str, context: Mapping[str, Any],
695 _selected: Any = selected, _args: tuple[tuple[str, Any], ...] = tuple(args.items()),
696 ) -> torch.nn.Module:
697 """Construct the selected family variant with ordinary YAML arguments."""
698 return _selected(module=module, module_fqn=module_fqn, context=context, **dict(_args))
699 factory.__name__ = selected.__name__
700 factory.__qualname__ = selected.__qualname__
701 factory.__module__ = selected.__module__
702 setattr(factory, "_hp_config", repr(sorted(args.items())))
703 else:
704 factory = selected
705 source_type = action.get("source_type")
706 specs.append(ModuleReplacementSpec(
707 match=patterns,
708 factory=factory,
709 module_type=_import(source_type, where) if source_type is not None else None,
710 exact_type=bool(action.get("exact_type", False)),
711 ))
712 return tuple(specs)
715class _Holder(torch.nn.Module):
716 """A module tree holding one part where a model holds it, so a rule's patterns match its path."""
718 def __init__(self, **parts: Any) -> None:
719 """Hold each part of *parts* under ``model``, a layer in a list as a model's stack holds it."""
720 super().__init__()
721 self.model = torch.nn.Module()
722 for name, part in parts.items():
723 setattr(self.model, name, torch.nn.ModuleList([part]) if name == "layers" else part)
726def _matching(spec: Any, names: Sequence[str]) -> Any:
727 """*spec* with the patterns that match a module of *names*, or ``None`` where none does."""
728 import fnmatch # pylint: disable=C0415
730 import dataclasses # pylint: disable=C0415
732 kept = tuple(pattern for pattern in spec.match
733 if any(fnmatch.fnmatchcase(name, pattern) for name in names))
734 return None if not kept else dataclasses.replace(spec, match=kept)
737def _replaced(holder: _Holder, specs: Sequence[Any]) -> None:
738 """Install on *holder* the replacements of *specs* whose patterns match one of its modules.
740 A rule's other patterns name modules no fake part holds, such as a
741 model's final norm beside a layer, and are left out rather than refused.
743 Raises:
744 CensusUnavailable: If a factory needs a library this host lacks, such
745 as an attention built on a native extension. Transformers' own
746 modules are not what the run trains, and measuring them prices a
747 layer nobody runs: a DeepSeek-V3.2 layer whose indexer runs eager
748 keeps, under the selective policy, 1192192 bytes a token against
749 the 516368 it keeps plain, its scores an fp32 tensor of one
750 sequence by another that the policy holds and no backward takes.
751 """
752 from hyper_parallel.models.replacement import ( # pylint: disable=C0415
753 apply_module_replacements,
754 compile_module_replacements,
755 )
757 names = [name for name, _ in holder.named_modules()]
758 kept = [found for found in (_matching(spec, names) for spec in specs) if found is not None]
759 if not kept:
760 return
761 try:
762 apply_module_replacements(holder, compile_module_replacements(holder, kept),
763 weights_mapping=[], context={}, capture_checkpoint_metadata=False)
764 except (ImportError, OSError) as exc:
765 raise CensusUnavailable(
766 f"this host cannot build the modules the run installs "
767 f"({[spec.factory.__name__ for spec in kept]}: {exc}); "
768 f"run the cost model where they build, or drop context.census and let the formulas price them"
769 ) from exc
772# Model families whose attention a census cannot price here, by model type,
773# with what is missing. Their layers build and run on this host, through
774# Transformers' own eager path, and that path is not the one the run takes:
775# measuring it prices a layer nobody runs, as a replacement this host cannot
776# build does (:func:`_replaced`, F53). A family belongs here only while no
777# contract in :mod:`~hyper_parallel.auto_parallel._npu_contracts` states what
778# its kernels save; stating one is what takes it out again.
779_NO_ACTIVATION_CONTRACT: Dict[str, str] = {
780 # Empty while every family with kernels of its own has its contracts
781 # stated. DeepSeek-V4 was the entry that forced this guard and now has
782 # them (:data:`_FAMILY_CONTRACTS`); the next family to arrive with a
783 # kernel and no contract belongs here rather than being priced from the
784 # path its run does not take.
785}
787# The families whose own modules call kernels this host has not got, with the
788# contracts their census runs under. A family here is measured as its run
789# behaves; a family in :data:`_NO_ACTIVATION_CONTRACT` is refused instead.
790_FAMILY_CONTRACTS = ("deepseek_v4",)
792# The families whose layers a census cannot run one at a time, by model type,
793# with why. Their census runs the model's own forward over the whole stack
794# (:func:`census_stack`), which calls each layer as the model calls it.
795_STACK_FAMILIES: Dict[str, str] = {
796 "deepseek_v4": (
797 "31 of the DeepSeek-V4.1 crop's 40 layers consume compressed keys and values, or an Indexer "
798 "key, that a source layer wrote earlier in the same forward and refuse to run without it; the "
799 "hidden state carries a hyper-connection dimension beside a mixing state; and the rotary "
800 "embedding is a pair of flavours the model picks between"
801 ),
802}
805@contextlib.contextmanager
806def _family_contracts(config: Any) -> Iterator[None]:
807 """The contracts the modules of *config*'s family run under, for what runs inside.
809 DeepSeek-V4.1's sparse attention is chosen on the device the tensors are
810 on, so a census, whose tensors are fake and on none, would take the
811 reference beside the kernel and measure a layer no run builds: the
812 attention's contract stands in. The Indexer's selection keeps nothing a
813 backward reads on either path, fused or reference, since a run frees its
814 scores' graph as soon as the indices are taken; its contract runs
815 the selection without recording that graph, and the census measures the
816 same bytes with it as without.
818 Whether this tree builds the model at all is the stack census's to
819 refuse (:func:`_model_class`): a layer's parameters do not depend on
820 which kernel runs it and are counted from the same fake layer.
821 """
822 if str(getattr(config, "model_type", "")) not in _FAMILY_CONTRACTS:
823 yield
824 return
825 from hyper_parallel.auto_parallel._npu_contracts import ( # pylint: disable=C0415
826 v41_attention_contract,
827 v41_indexer_contract,
828 )
829 rope = int(getattr(config, "qk_rope_head_dim", 0) or 0)
830 with v41_attention_contract(rope), v41_indexer_contract():
831 yield
834def _no_contract(config: Any) -> None:
835 """Refuse a census of one layer of a model whose kernels no contract states, or whose layers cannot run alone.
837 Raises:
838 CensusUnavailable: If *config*'s family is one of
839 :data:`_NO_ACTIVATION_CONTRACT`, or one of
840 :data:`_STACK_FAMILIES`, whose census is the whole stack's
841 (:func:`census_stack`). A layer's parameters do not depend on
842 which kernel runs it and are counted as ever
843 (:func:`census_parameters`); what it keeps for its backward
844 does, so this refuses rather than answer with the bytes of a
845 layer the run does not build.
846 """
847 family = str(getattr(config, "model_type", ""))
848 missing = _NO_ACTIVATION_CONTRACT.get(family)
849 if missing is not None:
850 raise CensusUnavailable(
851 f"no census of a {config.model_type} layer: {missing}; state the kernels' contracts, "
852 f"or drop context.census and read the formulas' memory as an estimate, not as a bound"
853 )
854 alone = _STACK_FAMILIES.get(family)
855 if alone is not None:
856 raise CensusUnavailable(
857 f"no census of one {family} layer, whose layers cannot be priced one at a time: {alone}. "
858 f"census_activations measures the whole stack instead (census_stack)"
859 )
862def _modeling(config: Any) -> Any:
863 """The Transformers modeling module of *config*'s model."""
864 model_type = str(config.model_type)
865 names = [model_type]
866 if model_type.endswith("_text"):
867 names.append(model_type[:-len("_text")])
868 for name in names:
869 try:
870 return importlib.import_module(f"transformers.models.{name}.modeling_{name}")
871 except ImportError:
872 continue
873 raise ValueError(f"no Transformers modeling module for model type {model_type!r}")
876def _classes(modeling: Any) -> Tuple[type, type]:
877 """The decoder layer and the language model's rotary embedding of *modeling*."""
878 classes = {name: cls for name, cls in vars(modeling).items() if isinstance(cls, type)}
879 layers = [cls for name, cls in classes.items() if name.endswith("DecoderLayer")]
880 rotary = [cls for name, cls in classes.items()
881 if name.endswith("RotaryEmbedding") and "Vision" not in name]
882 rotary.sort(key=lambda cls: "Text" not in cls.__name__)
883 if not layers or not rotary:
884 raise ValueError(f"{modeling.__name__} states no decoder layer or rotary embedding")
885 return layers[0], rotary[0]
888def tp_config(config: Any, tp: int) -> Any:
889 """*config* as one of *tp* tensor-parallel ranks holds its layers: heads and widths divided.
891 Two of the fields divided can be one field. A Transformers config class
892 may declare an ``attribute_map``, and DeepSeek-V4's maps
893 ``intermediate_size`` onto ``moe_intermediate_size``, so a model whose
894 feed-forward is all experts states one width under two names. Dividing
895 each name in turn then divides that width by *tp* twice, and the layer
896 built as a rank's share is a quarter as wide at TP 2 rather than half,
897 which makes the split the census derives from the pair too large and ND
898 over-reward tensor parallelism. Qwen3-MoE keeps the two apart, which is
899 why this stayed hidden. Each underlying attribute is divided once.
900 """
901 config = copy.deepcopy(config)
902 config.head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
903 aliases = getattr(type(config), "attribute_map", None) or {}
904 divided = set()
905 for name in _TP_FIELDS:
906 target = aliases.get(name, name)
907 if target in divided:
908 continue
909 value = getattr(config, name, None)
910 if isinstance(value, int) and not isinstance(value, bool) and value > 0:
911 setattr(config, name, max(1, value // tp))
912 divided.add(target)
913 return config
916def _grid_rotary(rotary: Any) -> bool:
917 """Whether *rotary* embeds a position on three grids, time, height and width, as a multimodal model's does.
919 Qwen3.5's and Qwen3-VL's hold the grids' ``mrope_section``; Qwen2-VL's
920 config states it.
921 """
922 if getattr(rotary, "mrope_section", None) is not None:
923 return True
924 config = getattr(rotary, "config", None)
925 return any(isinstance(getattr(config, name, None), Mapping) and "mrope_section" in getattr(config, name)
926 for name in ("rope_parameters", "rope_scaling"))
929def _positions(rotary: Any, hidden: torch.Tensor, seq_length: int) -> Tuple[torch.Tensor, Any]:
930 """The position ids of a sequence of *seq_length* tokens, and the embeddings *rotary* gives them.
932 A rotary on three grids (:func:`_grid_rotary`) is handed a row of ids
933 for each grid, as its model hands them, a text token's index on all
934 three; the layers keep the one row. Transformers 5.17 and 5.18 refuse
935 the single row the versions before them expanded themselves.
936 """
937 position_ids = torch.arange(seq_length).unsqueeze(0)
938 rotary_ids = position_ids[None].expand(3, *position_ids.shape) if _grid_rotary(rotary) else position_ids
939 with torch.no_grad():
940 return position_ids, rotary(hidden, rotary_ids)
943def _run(layer: Any, hidden: torch.Tensor, positions: Tuple[torch.Tensor, Any],
944 call: Optional[Callable[..., Any]] = None) -> torch.Tensor:
945 """One forward of *layer*, as its model calls it at *positions*, through *call* if given."""
946 position_ids, position_embeddings = positions
947 accepted = inspect.signature(layer.forward).parameters
948 kwargs = {name: value for name, value in (
949 ("position_embeddings", position_embeddings), ("position_ids", position_ids),
950 ("attention_mask", None), ("use_cache", False)) if name in accepted}
951 out = (call or layer)(hidden, **kwargs)
952 return out[0] if isinstance(out, tuple) else out
955# The function a gated-delta-rule attention of Transformers 5.15 and later
956# calls by name, which its modeling module defines; earlier versions hold the
957# rule in the module's ``chunk_gated_delta_rule``.
958_DELTA_RULE_FUNCTION = "torch_chunk_gated_delta_rule"
961def delta_rule_module(module: Any) -> bool:
962 """Whether *module* runs a gated delta rule, whichever way its Transformers version calls it.
964 Up to 5.14 the module holds the rule in ``chunk_gated_delta_rule``; from
965 5.15 its forward calls its modeling module's ``torch_chunk_gated_delta_rule``
966 by name, and the module holds nothing.
967 """
968 if hasattr(module, "chunk_gated_delta_rule"):
969 return True
970 forward = inspect.unwrap(getattr(type(module), "forward", None) or (lambda: None))
971 return _DELTA_RULE_FUNCTION in getattr(getattr(forward, "__code__", None), "co_names", ())
974@contextlib.contextmanager
975def _delta_rule_held(modeling: Any, layer: Any, rule: Callable[..., Any]) -> Iterator[None]:
976 """Have *layer*'s delta rule module run *rule* through its ``chunk_gated_delta_rule``, while the context lasts.
978 The census swaps and follows the rule through that attribute. A module
979 of Transformers 5.15 or later, which calls its modeling module's function
980 by name, gets the attribute too, and the name runs it.
982 Raises:
983 ValueError: If the layer holds more than one such module, which the
984 name could not tell apart.
985 """
986 modules = [module for module in layer.modules() if delta_rule_module(module)]
987 for module in modules:
988 module.chunk_gated_delta_rule = rule
989 by_name = getattr(modeling, _DELTA_RULE_FUNCTION, None)
990 if not modules or by_name is None:
991 yield
992 return
993 if len(modules) > 1:
994 raise ValueError(f"{type(layer).__name__} holds {len(modules)} gated delta rule modules, and "
995 f"{_DELTA_RULE_FUNCTION} cannot tell which one calls it")
996 module = modules[0]
998 def by_attribute(*args: Any, **kwargs: Any) -> Any:
999 """The module's rule as the attribute holds it when called: the census wraps it afterwards."""
1000 return module.chunk_gated_delta_rule(*args, **kwargs)
1002 setattr(modeling, _DELTA_RULE_FUNCTION, by_attribute)
1003 try:
1004 yield
1005 finally:
1006 setattr(modeling, _DELTA_RULE_FUNCTION, by_name)
1009@contextlib.contextmanager
1010def _fake_layer(config: Any, layer_index: int, gdn_backend: str = _EAGER_GDN,
1011 replacements: Sequence[Any] = ()) -> Iterator[Tuple[Any, Any]]:
1012 """Layer *layer_index* of *config* and its model's rotary embedding, on fake tensors and the runtime's kernels.
1014 *gdn_backend* names the gated delta rule the layer runs, as the runtime
1015 names it: ``"eager"`` for Transformers' own chunked implementation,
1016 ``"triton"`` for HyperParallel's kernel's contract, which saves the rule's
1017 inputs, the cumulated gates, beta and one chunk matrix and nothing more.
1018 Eager is the default because it is the rule a run reaches: the runtime's
1019 dispatcher defaults to it (``components/modules/gated_delta_net.py``,
1020 ``chunk_gated_delta_rule``), the only caller that can select the kernel is
1021 the context-parallel wrapper, which defaults to eager as well, and no
1022 recipe states a backend. The two are not interchangeable: a layer keeping
1023 the kernel's saved set keeps less than half as much.
1025 Where *replacements* are given, the modules they name are HyperParallel's
1026 fused ones, running under the kernels' contracts for as long as the layer
1027 does. Eager is Transformers' own Python rule, never a kernel a hub or the
1028 ``fla`` package would put in its place.
1030 Raises:
1031 ValueError: If *gdn_backend* names no backend the runtime has.
1032 """
1033 if gdn_backend not in _GDN_BACKENDS:
1034 raise ValueError(
1035 f"unsupported GDN backend {gdn_backend!r}; "
1036 f"expected one of {sorted(_GDN_BACKENDS)}."
1037 )
1038 modeling = _modeling(config)
1039 layer_cls, rotary_cls = _classes(modeling)
1040 config = copy.deepcopy(config)
1041 config._attn_implementation = _FLASH # pylint: disable=protected-access
1042 config._experts_implementation = "grouped_mm" # pylint: disable=protected-access
1043 modeling.ALL_ATTENTION_FUNCTIONS.register(_FLASH, _flash_attention)
1044 with npu_contracts() if replacements else contextlib.nullcontext(), \
1045 _family_contracts(config), \
1046 FakeTensorMode(allow_non_fake_inputs=True):
1047 default = torch.get_default_dtype()
1048 torch.set_default_dtype(torch.bfloat16)
1049 try:
1050 layer = layer_cls(config, layer_index)
1051 rotary = rotary_cls(config=config)
1052 finally:
1053 torch.set_default_dtype(default)
1054 rule = (_gated_delta_rule(modeling) if gdn_backend == _HYPER_GDN
1055 else inspect.unwrap(getattr(modeling, _DELTA_RULE_FUNCTION, None) or (lambda: None)))
1056 with _delta_rule_held(modeling, layer, rule):
1057 if replacements:
1058 holder = _Holder(layers=layer)
1059 _replaced(holder, replacements)
1060 layer = holder.model.layers[0]
1061 layer.train()
1062 yield layer, rotary
1065class _ByPart(TorchDispatchMode):
1066 """A walk over a layer's forward that attributes each op to the part of the layer running it."""
1068 def __init__(self, layer: Any, recurrent: bool = True) -> None:
1069 """Follow which of *layer*'s modules runs, and which tensors are its parameters.
1071 *recurrent* is whether an attention of this layer has a recurrence to
1072 name: where it has none, an op of the attention that no weight takes
1073 part in is the attention's own, its mask or its rotary. A walk that
1074 counts only matmuls leaves it at the default, because a weightless
1075 matmul under an attention is a recurrence or a score, and a score is
1076 counted on its own.
1077 """
1078 super().__init__()
1079 self.recurrent = recurrent
1080 self.params = {param.untyped_storage()._cdata # pylint: disable=protected-access
1081 for param in layer.parameters()}
1082 self.running: List[str] = []
1083 for name, module in layer.named_modules():
1084 if name:
1085 module.register_forward_pre_hook(functools.partial(self._enter, name))
1086 module.register_forward_hook(self._leave)
1088 def _enter(self, name: str, *_: Any) -> None:
1089 """Note that the module at *name* runs."""
1090 self.running.append(name)
1092 def _leave(self, *_: Any) -> None:
1093 """Note that the innermost module running has returned."""
1094 self.running.pop()
1096 def _part(self, args: Sequence[Any]) -> str:
1097 """The part an op on *args* runs for: its module's, or an attention's recurrence where no weight takes part."""
1098 part = _parameter_part(f"{self.running[-1]}.weight") if self.running else "ffn"
1099 weighted = any(isinstance(arg, torch.Tensor)
1100 and arg.untyped_storage()._cdata in self.params # pylint: disable=protected-access
1101 for arg in args)
1102 if part == "attention" and not weighted and self.recurrent:
1103 return "linrec"
1104 return part
1107class _PartFlops(_ByPart):
1108 """A layer's forward FLOPs by part: its matmuls' toward the part running them, flash attention's the scores'."""
1110 def __init__(self, layer: Any) -> None:
1111 """Count no FLOPs yet, and follow *layer*'s modules."""
1112 super().__init__(layer)
1113 self.flops: Dict[str, int] = {}
1115 def _add(self, part: str, flops: int) -> None:
1116 """Count *flops* toward *part*."""
1117 self.flops[part] = self.flops.get(part, 0) + flops
1119 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
1120 out = func(*args, **(kwargs or {}))
1121 if func is torch.ops.nd_census.flash_attention.default:
1122 query, key, value = args[:3]
1123 batch, heads, seq, width = query.shape
1124 self._add("scores", 2 * batch * heads * seq * key.shape[2] * (width + value.shape[-1]))
1125 elif flops := _matmul_flops(func, args):
1126 self._add(self._part(args), flops)
1127 return out
1130class _PartTraffic(_ByPart):
1131 """The bytes a layer's forward moves by part, the parameters among them, and the ops it dispatches.
1133 An op reads the tensors it is given and writes the tensors it returns.
1134 A view writes nothing and reads nothing: it renames bytes its producer
1135 already wrote, and no kernel runs for it, so it counts for neither. A
1136 tensor given to an op twice is read twice, as the count is of the op's
1137 arguments rather than of the storages behind them.
1139 The parameters are counted again on their own, because a strategy
1140 divides their traffic where it leaves the activations' alone: the degree
1141 sharding a parameter divides the bytes read of it, and the routed
1142 experts' part is the one the expert-parallel degree divides.
1143 """
1145 def __init__(self, layer: Any, recurrent: bool = True) -> None:
1146 """Move nothing yet, and follow *layer*'s modules."""
1147 super().__init__(layer, recurrent=recurrent)
1148 self.moved: Dict[str, float] = {}
1149 self.parameter: Dict[str, float] = {}
1150 self.launches: Dict[str, int] = {}
1152 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
1153 kwargs = kwargs or {}
1154 out = func(*args, **kwargs)
1155 if func.namespace == "prim" or getattr(func, "is_view", False):
1156 return out
1157 part = ("scores" if func is torch.ops.nd_census.flash_attention.default
1158 else self._part(args))
1159 moved = parameter = 0
1160 for tensor in tree_flatten((args, kwargs))[0]:
1161 if not isinstance(tensor, torch.Tensor):
1162 continue
1163 size = tensor.numel() * tensor.element_size()
1164 moved += size
1165 if tensor.untyped_storage()._cdata in self.params: # pylint: disable=protected-access
1166 parameter += size
1167 for tensor in tree_flatten(out)[0]:
1168 if isinstance(tensor, torch.Tensor):
1169 moved += tensor.numel() * tensor.element_size()
1170 self.moved[part] = self.moved.get(part, 0.0) + moved
1171 self.parameter[part] = self.parameter.get(part, 0.0) + parameter
1172 self.launches[part] = self.launches.get(part, 0) + 1
1173 return out
1176def census_flops(config: Any, layer_index: int, seq_length: int) -> Dict[str, int]:
1177 """The forward FLOPs of layer *layer_index* of *config* on one sequence, by part.
1179 Args:
1180 config: The language model's Transformers config.
1181 layer_index: The layer to build, which settles its kind.
1182 seq_length: Tokens of the sequence it runs.
1184 Returns:
1185 The FLOPs of each part that runs a matmul: the attention's
1186 projections (``attention``), its scores and values at every pair of
1187 tokens (``scores``), a linear attention's recurrence (``linrec``),
1188 and the parts of the feed-forward :func:`_parameter_part` names.
1189 """
1190 with _fake_layer(config, layer_index) as (layer, rotary):
1191 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)
1192 positions = _positions(rotary, hidden, seq_length)
1193 counter = _PartFlops(layer)
1194 with counter:
1195 _run(layer, hidden, positions)
1196 return counter.flops
1199def census_traffic(config: Any, layer_index: int, seq_length: int,
1200 gdn_backend: str = _EAGER_GDN) -> LayerTraffic:
1201 """What one forward of layer *layer_index* of *config* moves, by the part of the layer moving it.
1203 The quantity the time model counts is matmul FLOPs, and the cluster's
1204 kernel tables put those at 6.3 to 7.1% of a profiled step of this model.
1205 This counts the other quantity a step could be priced by, on the same
1206 forward and with no device, so that a measured round can set each
1207 beside the time it took.
1209 Args:
1210 config: The language model's Transformers config.
1211 layer_index: The layer to build, which settles its kind.
1212 seq_length: Tokens of the micro-batch of one sequence it runs.
1213 gdn_backend: The gated delta rule the run takes (:func:`_fake_layer`).
1214 It is the largest single choice here: Transformers' own chunked
1215 rule runs a matmul a chunk and the kernel's contract is one op.
1217 Returns:
1218 A :class:`LayerTraffic` of the whole layer at tensor and context
1219 parallelism 1, by the parts :func:`census_flops` prices
1220 (:class:`_PartTraffic`).
1221 """
1222 with _fake_layer(config, layer_index, gdn_backend=gdn_backend) as (layer, rotary):
1223 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)
1224 positions = _positions(rotary, hidden, seq_length)
1225 # The rule's own module is what a layer has a recurrence by, as
1226 # _fake_layer finds it to settle which rule the layer runs.
1227 recurrent = any(delta_rule_module(module) for module in layer.modules())
1228 counter = _PartTraffic(layer, recurrent=recurrent)
1229 with counter:
1230 _run(layer, hidden, positions)
1231 return LayerTraffic(counter.moved, counter.parameter, counter.launches)
1234def census_layer(config: Any, layer_index: int, seq_length: int, selective: bool = False,
1235 ops: Optional[Dict[str, int]] = None,
1236 replacements: Sequence[Any] = (),
1237 gdn_backend: str = _EAGER_GDN) -> Tuple[int, int]:
1238 """Bytes layer *layer_index* of *config* keeps for its backward, and the most its backward holds.
1240 Args:
1241 config: The language model's Transformers config.
1242 layer_index: The layer to build, which settles its kind.
1243 seq_length: Tokens of the micro-batch of one sequence it runs.
1244 selective: Whether the layer runs under HyperParallel's selective
1245 activation checkpointing, as its trainer wraps a layer.
1246 ops: Where given, and the layer runs without checkpointing, takes
1247 the bytes it saves by the op saving them (:class:`_SavedOps`).
1248 replacements: The module replacements a run installs
1249 (:func:`replacement_specs`).
1250 gdn_backend: The gated delta rule the run takes (:func:`_fake_layer`).
1252 Returns:
1253 The bytes the forward keeps for the backward, its input included,
1254 and the bytes of activations the backward holds when it holds the
1255 most, its own gradients included, less the parameters and those
1256 gradients.
1258 Raises:
1259 CensusUnavailable: If no contract states what this model's kernels
1260 save (:func:`_no_contract`).
1261 """
1262 _no_contract(config)
1263 with _fake_layer(config, layer_index, gdn_backend=gdn_backend,
1264 replacements=replacements) as (layer, rotary):
1265 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16, requires_grad=True)
1266 grad = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)
1267 if replacements:
1268 # A fused kernel keeps a constant per device, such as the mask
1269 # the attention runs its causal sparse mode with, which the
1270 # runtime allocates once for every layer and micro-batch: one
1271 # forward allocates them before the layer's own bytes are
1272 # counted, so they count toward no layer.
1273 with torch.no_grad():
1274 _run(layer, hidden, _positions(rotary, hidden, seq_length))
1275 if not selective:
1276 saved_ops = None if ops is None else _SavedOps(layer)
1277 measured = _measure(list(layer.parameters()), (hidden,),
1278 lambda: _run(layer, hidden, _positions(rotary, hidden, seq_length)),
1279 lambda out: out.backward(grad), grad_inputs=(grad,), ops=saved_ops)
1280 if saved_ops is not None:
1281 ops.update(saved_ops.saved)
1282 return measured
1283 # The model embeds the positions once for all its layers, outside
1284 # their checkpoints, and the trainer checkpoints each layer's call.
1285 positions = _positions(rotary, hidden, seq_length)
1286 call = functools.partial(activation_memory.checkpoint, layer, swap_inputs=False,
1287 context_fn=_selective_contexts)
1288 return _measure(list(layer.parameters()), (hidden,), lambda: _run(layer, hidden, positions, call),
1289 lambda out: out.backward(grad), grad_inputs=(grad,), checkpointed=True,
1290 shared=tree_flatten(positions)[0])
1293def census_saved_ops(config: Any, layer_index: int, seq_length: int,
1294 replacements: Sequence[Any] = ()) -> Dict[str, float]:
1295 """What layer *layer_index* of *config* keeps for its backward for each op, per token, the whole layer.
1297 Args:
1298 config: The language model's Transformers config.
1299 layer_index: The layer to build, which settles its kind.
1300 seq_length: Tokens of the micro-batch of one sequence it runs.
1301 replacements: The module replacements a run installs
1302 (:func:`replacement_specs`).
1304 Returns:
1305 ``{op: bytes per token}``, by the op saving them
1306 (:class:`_SavedOps`), which sum to what the layer keeps.
1307 """
1308 ops: Dict[str, int] = {}
1309 census_layer(config, layer_index, seq_length, ops=ops, replacements=replacements)
1310 return {op: size / seq_length for op, size in ops.items()}
1313def census_recomputed(config: Any, layer_index: int, seq_length: int) -> Tuple[float, float]:
1314 """The shares of layer *layer_index*'s matmul FLOPs its backward runs again under selective checkpointing.
1316 Args:
1317 config: The language model's Transformers config.
1318 layer_index: The layer to build, which settles its kind.
1319 seq_length: Tokens of the micro-batch of one sequence it runs.
1321 Returns:
1322 The share of its attention's projections' FLOPs, and of the rest of
1323 the layer's, that HyperParallel's selective policy recomputes: it
1324 saves every other matmul's output, in the order the layer runs
1325 them, and a biased projection's always.
1326 """
1327 with _fake_layer(config, layer_index) as (layer, rotary):
1328 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16, requires_grad=True)
1329 ledger = _RecomputedMatmuls(layer)
1330 call = functools.partial(activation_memory.checkpoint, layer, swap_inputs=False,
1331 context_fn=functools.partial(_selective_contexts, ledger))
1332 with _checkpoints_on_cpu():
1333 _run(layer, hidden, _positions(rotary, hidden, seq_length), call)
1334 return ledger.share(attention=True), ledger.share(attention=False)
1337def _parameter_part(name: str) -> str:
1338 """The part of a decoder layer the parameter at path *name* belongs to, as ND prices the parts.
1340 The layer's norms and its attention's per-head query and key norms are
1341 its norms; the rest of a child named for attention is its attention.
1342 Of the feed-forward, a shared expert's parameters, the routed experts'
1343 and the router's are parts of their own, and the rest is the dense
1344 feed-forward's.
1345 """
1346 segments = name.split(".")
1347 module = segments[-2] if len(segments) > 1 else ""
1348 if "norm" in segments[0] or module in _QK_NORMS:
1349 return "norm"
1350 if "attn" in segments[0] or "attention" in segments[0]:
1351 return "attention"
1352 if any("shared" in segment for segment in segments):
1353 return "shared"
1354 if "experts" in segments:
1355 return "routed"
1356 if module in ("gate", "router"):
1357 return "router"
1358 return "ffn"
1361def census_parameters(config: Any, layer_index: int, replacements: Sequence[Any] = ()) -> Dict[str, int]:
1362 """The parameters of layer *layer_index* of *config*, by part.
1364 Args:
1365 config: The language model's Transformers config.
1366 layer_index: The layer to build, which settles its kind.
1367 replacements: The module replacements a run installs
1368 (:func:`replacement_specs`). A replacement holds the
1369 parameters its source held, fused or renamed, so it counts the
1370 same.
1372 Returns:
1373 The parameter count of each part the layer has
1374 (:func:`_parameter_part`).
1375 """
1376 parts: Dict[str, int] = {}
1377 with _fake_layer(config, layer_index, replacements=replacements) as (layer, _):
1378 for name, param in layer.named_parameters():
1379 part = _parameter_part(name)
1380 parts[part] = parts.get(part, 0) + param.numel()
1381 return parts
1384def census_final_norm(config: Any, replacements: Sequence[Any] = ()) -> int:
1385 """The parameters of *config*'s final norm, of the class of a layer's input norm."""
1386 with _fake_layer(config, 0, replacements=replacements) as (layer, _):
1387 return sum(param.numel() for param in layer.input_layernorm.parameters())
1390def census_output(config: Any, seq_length: int, replacements: Sequence[Any] = ()) -> Tuple[int, int]:
1391 """Bytes the output layer of *config* keeps for its backward, and the most its backward holds.
1393 Args:
1394 config: The language model's Transformers config.
1395 seq_length: Tokens of the micro-batch of one sequence it runs.
1396 replacements: The module replacements a run installs
1397 (:func:`replacement_specs`), of which the final norm's is the
1398 one that reaches the output layer.
1400 Returns:
1401 As :func:`census_layer`: the final norm's input included, the
1402 output table and its gradient left out.
1403 """
1404 modeling = _modeling(config)
1405 layer_cls, _ = _classes(modeling)
1406 loss_function = importlib.import_module("transformers.loss.loss_utils").ForCausalLMLoss
1407 with npu_contracts() if replacements else contextlib.nullcontext(), \
1408 FakeTensorMode(allow_non_fake_inputs=True):
1409 default = torch.get_default_dtype()
1410 torch.set_default_dtype(torch.bfloat16)
1411 try:
1412 # The final norm is of the class of a layer's input norm.
1413 norm = layer_cls(config, 0).input_layernorm
1414 head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1415 finally:
1416 torch.set_default_dtype(default)
1417 if replacements:
1418 holder = _Holder(norm=norm)
1419 _replaced(holder, replacements)
1420 norm = holder.model.norm
1421 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16, requires_grad=True)
1422 labels = torch.randint(0, config.vocab_size, (1, seq_length))
1423 # The loss alone holds the logits: the trainer drops the model's
1424 # output before the backward.
1425 return _measure([*norm.parameters(), *head.parameters()], (hidden,),
1426 lambda: loss_function(head(norm(hidden)), labels, config.vocab_size),
1427 lambda loss: loss.backward())
1430def census_output_activations(config: Any, seq_length: int = 4096,
1431 replacements: Sequence[Any] = ()) -> KindActivations:
1432 """What the output layer of *config* keeps and holds, per token, with its vocabulary and half of it.
1434 Args:
1435 config: The language model's Transformers config.
1436 seq_length: The tokens the census runs the layer at.
1437 replacements: The module replacements a run installs
1438 (:func:`replacement_specs`).
1440 Returns:
1441 The layer's :class:`KindActivations`, whose TP part is the part a
1442 vocabulary-parallel loss over two ranks halves.
1443 """
1444 half = copy.deepcopy(config)
1445 half.vocab_size = max(1, config.vocab_size // 2)
1446 (saved_1, working_1), (saved_2, working_2) = (
1447 census_output(each, seq_length, replacements) for each in (config, half))
1448 return KindActivations(
1449 saved=max(0.0, 2 * saved_2 - saved_1) / seq_length,
1450 saved_tp=max(0.0, 2 * (saved_1 - saved_2)) / seq_length,
1451 working=max(0.0, 2 * working_2 - working_1) / seq_length,
1452 working_tp=max(0.0, 2 * (working_1 - working_2)) / seq_length,
1453 seq_length=int(seq_length),
1454 )
1457class _FakeScalars(TorchDispatchMode):
1458 """A scalar read of a fake tensor answered with zero, as :class:`_LiveBytes` answers it.
1460 A model can read a value out of a tensor while it builds, as
1461 DeepSeek-V4.1's Engram reads its token map's pad entry to build its hash
1462 mapping, which a fake tensor refuses. The value settles a padding id,
1463 never a shape the census counts.
1464 """
1466 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
1467 if func is _SCALAR_READ and isinstance(args[0], FakeTensor):
1468 return 0
1469 return func(*args, **(kwargs or {}))
1472def _model_class(config: Any) -> Tuple[Any, bool]:
1473 """What the trainer builds *config* as, by its ``from_config``: its family's own class, else Transformers'.
1475 The trainer's builder discovers the family's registration, asks the
1476 model registry for a class registered by the architecture the config
1477 names, and builds Transformers' own model where none is.
1479 Returns:
1480 The class, and whether it is Transformers' own.
1482 Raises:
1483 CensusUnavailable: If the config names an architecture neither the
1484 registry nor Transformers builds, as a DeepSeek-V4.1 config does
1485 on a tree that does not carry that model: Transformers'
1486 DeepSeek-V4 in its place would price a different model.
1487 """
1488 import transformers # pylint: disable=C0415
1489 from hyper_parallel.models._transformers.config_resolver import get_is_hf_model # pylint: disable=C0415
1490 from hyper_parallel.models.registry import _resolve_custom_model_cls # pylint: disable=C0415
1492 architecture = (getattr(config, "architectures", None) or [""])[0] or ""
1493 if not get_is_hf_model(config):
1494 return _resolve_custom_model_cls(architecture), False
1495 if architecture and not hasattr(transformers, architecture):
1496 raise CensusUnavailable(
1497 f"the config names {architecture}, which neither this tree's model registry nor Transformers "
1498 f"builds; price it from a tree carrying the model, or drop context.census and read the "
1499 f"formulas' memory as an estimate, not as a bound"
1500 )
1501 return transformers.AutoModelForCausalLM, True
1504class _Stack:
1505 """Which layer of a stack runs, forward and backward, and where in a tracker's log each backward runs."""
1507 def __init__(self, layers: Sequence[Any], clock: Callable[[], int]) -> None:
1508 """Follow *layers* through both passes, placing their backwards by *clock*."""
1509 self.running: List[int] = []
1510 self.windows: Dict[int, List[int]] = {}
1511 self.clock = clock
1512 self.handles: List[Any] = []
1513 for index, layer in enumerate(layers):
1514 self.handles += [
1515 layer.register_forward_pre_hook(functools.partial(self._enter, index)),
1516 layer.register_forward_hook(self._leave),
1517 layer.register_full_backward_pre_hook(functools.partial(self._enter_backward, index)),
1518 layer.register_full_backward_hook(self._leave_backward),
1519 ]
1520 self.backward: List[int] = []
1522 # A hook's return replaces what the module takes or gives, so each of
1523 # these returns None.
1525 def _enter(self, index: int, *_: Any) -> None:
1526 """Note that the layer at *index* runs its forward."""
1527 self.running.append(index)
1529 def _leave(self, *_: Any) -> None:
1530 """Note that the innermost layer running its forward has returned."""
1531 if self.running:
1532 self.running.pop()
1534 def _enter_backward(self, index: int, *_: Any) -> None:
1535 """Open the window of the layer at *index*, whose backward starts."""
1536 self.backward.append(index)
1537 self.windows.setdefault(index, [self.clock(), self.clock()])
1539 def _leave_backward(self, *_: Any) -> None:
1540 """Close the window of the layer whose backward has returned."""
1541 if self.backward:
1542 index = self.backward.pop()
1543 self.windows[index][1] = self.clock()
1545 @property
1546 def layer(self) -> Optional[int]:
1547 """The outermost layer running its forward, or None."""
1548 return self.running[0] if self.running else None
1550 def remove(self) -> None:
1551 """Remove every hook this placed."""
1552 for handle in self.handles:
1553 handle.remove()
1556def census_stack(config: Any, seq_length: int, replacements: Sequence[Any] = ()) -> List[Tuple[int, int]]:
1557 """What each layer of *config* keeps for its backward, and the most its backward holds, from one pass of the stack.
1559 For a model whose layers cannot run alone (:data:`_STACK_FAMILIES`): the
1560 model the trainer builds (:func:`_model_class`), on fake tensors and
1561 under its family's contracts, runs its own forward over a micro-batch of
1562 one sequence and then its backward, so each layer is called as its model
1563 calls it, with what earlier layers wrote, its hidden and mixing states
1564 and the rotary flavours the model computes. Each save counts toward the
1565 layer whose forward makes it first, and a layer keeps those of its saves
1566 still live when the forward ends (:func:`_measure`'s rule). Its backward
1567 holds what it keeps and the most it holds beyond every layer's saves,
1568 the parameters and their gradients while it runs. The embedding's and
1569 the final norm's saves count toward no layer: the output layer's census
1570 prices the end of the model.
1572 Args:
1573 config: The language model's Transformers config.
1574 seq_length: Tokens of the micro-batch of one sequence it runs.
1575 replacements: The module replacements a run installs
1576 (:func:`replacement_specs`).
1578 Returns:
1579 ``(kept, held)`` bytes of each decoder layer, in model order.
1581 Raises:
1582 CensusUnavailable: If this tree cannot build the model the config
1583 names (:func:`_model_class`), or a module a replacement installs
1584 (:func:`_replaced`).
1585 """
1586 model_cls, native = _model_class(config)
1587 config = copy.deepcopy(config)
1588 if native:
1589 # Transformers' own model runs its attention through the interface,
1590 # as _fake_layer's layer does, so it runs the kernel's contract.
1591 config._attn_implementation = _FLASH # pylint: disable=protected-access
1592 _modeling(config).ALL_ATTENTION_FUNCTIONS.register(_FLASH, _flash_attention)
1593 with npu_contracts() if replacements else contextlib.nullcontext(), _family_contracts(config), \
1594 FakeTensorMode(allow_non_fake_inputs=True), _FakeScalars():
1595 default = torch.get_default_dtype()
1596 torch.set_default_dtype(torch.bfloat16)
1597 try:
1598 model = model_cls.from_config(config)
1599 finally:
1600 torch.set_default_dtype(default)
1601 # The experts run as grouped matmuls, as every layer _fake_layer
1602 # builds does. Transformers falls back to its eager experts for a
1603 # class whose module it does not know, a loop over the experts with
1604 # token counts no fake tensor has; the experts read this per call.
1605 model.config._experts_implementation = "grouped_mm" # pylint: disable=protected-access
1606 if replacements:
1607 _replaced(model, replacements)
1608 model.train()
1609 backbone = model.get_decoder()
1610 params = list(model.parameters())
1611 stored = {p.untyped_storage()._cdata for p in params} # pylint: disable=protected-access
1612 live = _LiveBytes()
1613 follow = _Stack(backbone.layers, lambda: len(live.events))
1614 saved_by: Dict[Tuple[int, int], int] = {}
1616 def pack(tensor: torch.Tensor) -> torch.Tensor:
1617 """Count a tensor autograd saves, unless it is a parameter's, toward the layer saving it first."""
1618 if tensor.untyped_storage()._cdata not in stored and follow.layer is not None: # pylint: disable=W0212
1619 saved_by.setdefault(live.key(tensor), follow.layer)
1620 # Detached, as _measure returns it: a graph the forward drops is freed.
1621 return tensor.detach()
1623 grads: List[Tuple[int, int]] = []
1624 handles = [p.register_hook(lambda grad: grads.append(live.key(grad))) for p in params]
1625 kept = [0] * len(backbone.layers)
1626 try:
1627 with live, torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor):
1628 out = backbone(input_ids=torch.zeros(1, seq_length, dtype=torch.long)).last_hidden_state
1629 for key, index in saved_by.items():
1630 kept[index] += live.sizes.get(key, 0)
1631 out.sum().backward()
1632 finally:
1633 follow.remove()
1634 for handle in handles:
1635 handle.remove()
1636 grads += [live.key(p.grad) for p in params if p.grad is not None]
1637 # What every layer keeps, the gradients and the parameters left out, the
1638 # rest is what a layer's backward holds beyond them while it runs.
1639 excluded = set(saved_by) | set(grads) | {key for key, _ in live.events if key[0] in stored}
1640 beyond = list(itertools.accumulate(0 if key in excluded else delta for key, delta in live.events))
1641 held = []
1642 for index, size in enumerate(kept):
1643 begin, end = follow.windows.get(index, (0, 0))
1644 held.append(size + max(beyond[begin:end], default=0))
1645 return list(zip(kept, held))
1648def _stack_activations(config: Any, layers: Iterable[Mapping[str, Any]], seq_length: int,
1649 replacements: Sequence[Any] = ()) -> Dict[str, KindActivations]:
1650 """What a layer of each kind of *layers* keeps and holds, per token, from the stack at TP 1 and TP 2.
1652 A kind's record is the mean of its layers', which is what the memory
1653 model applies to each of them: with one record a kind, the stack's total
1654 comes back whole. The parts tensor parallelism splits follow from the
1655 pair as :func:`census_activations` derives them. No selective part is
1656 stated: a selective layer of these kinds is priced by the formulas.
1657 """
1658 kinds = [str(group["kind"]) for group in layers if not group.get("mtp") for _ in range(int(group["count"]))]
1659 logger.info("census of the whole stack of %d layers at %d tokens, at TP 1 and TP 2: a %s layer "
1660 "cannot run alone", len(kinds), seq_length, config.model_type)
1661 one, two = (census_stack(tp_config(config, tp), seq_length, replacements) for tp in (1, 2))
1662 if len(one) != len(kinds):
1663 raise ValueError(f"the model builds {len(one)} decoder layers, where its spec states {len(kinds)}")
1664 out: Dict[str, KindActivations] = {}
1665 for kind in dict.fromkeys(kinds):
1666 members = [index for index, each in enumerate(kinds) if each == kind]
1667 (kept_1, held_1), (kept_2, held_2) = (
1668 tuple(sum(stack[index][part] for index in members) / len(members) for part in (0, 1))
1669 for stack in (one, two))
1670 out[kind] = KindActivations(
1671 saved=max(0.0, 2 * kept_2 - kept_1) / seq_length,
1672 saved_tp=max(0.0, 2 * (kept_1 - kept_2)) / seq_length,
1673 working=max(0.0, 2 * held_2 - held_1) / seq_length,
1674 working_tp=max(0.0, 2 * (held_1 - held_2)) / seq_length,
1675 seq_length=int(seq_length),
1676 )
1677 return out
1680def census_activations(config: Any, layers: Iterable[Mapping[str, Any]], seq_length: int = 4096,
1681 replacements: Sequence[Any] = (),
1682 gdn_backend: str = _EAGER_GDN) -> Dict[str, KindActivations]:
1683 """What a layer of each kind of *layers* keeps and holds, per token, at TP 1 and TP 2.
1685 Args:
1686 config: The language model's Transformers config.
1687 layers: The model spec's layer stack, groups in model order; MTP
1688 groups are left out.
1689 seq_length: The tokens the census runs a layer at.
1690 replacements: The module replacements a run installs
1691 (:func:`replacement_specs`).
1692 gdn_backend: The gated delta rule the run takes (:func:`_fake_layer`).
1694 Returns:
1695 Each kind's :class:`KindActivations`: of the bytes a layer at TP 2
1696 holds, the part TP splits is half the part at TP 1; what it keeps
1697 under HyperParallel's selective activation checkpointing too, the
1698 shares of its matmul FLOPs that recomputes, and what it keeps for
1699 each op, its records per op. A family whose layers cannot run
1700 alone (:data:`_STACK_FAMILIES`) is measured over its whole stack
1701 instead (:func:`census_stack`), and its records state the plain
1702 layer's bytes only.
1704 Raises:
1705 CensusUnavailable: If this tree cannot build what the run builds.
1706 """
1707 if str(getattr(config, "model_type", "")) in _STACK_FAMILIES:
1708 return _stack_activations(config, layers, seq_length, replacements)
1709 firsts: Dict[str, int] = {}
1710 index = 0
1711 for group in layers:
1712 if group.get("mtp"):
1713 continue
1714 firsts.setdefault(str(group["kind"]), index)
1715 index += int(group["count"])
1716 out: Dict[str, KindActivations] = {}
1717 for kind, layer_index in firsts.items():
1718 ops_1: Dict[str, int] = {}
1719 ops_2: Dict[str, int] = {}
1720 (saved_1, working_1), (saved_2, working_2) = (
1721 census_layer(tp_config(config, tp), layer_index, seq_length, ops=ops,
1722 replacements=replacements, gdn_backend=gdn_backend)
1723 for tp, ops in ((1, ops_1), (2, ops_2)))
1724 (kept_1, _), (kept_2, _) = (
1725 census_layer(tp_config(config, tp), layer_index, seq_length, selective=True,
1726 replacements=replacements, gdn_backend=gdn_backend)
1727 for tp in (1, 2))
1728 attention_mm, ffn_mm = census_recomputed(config, layer_index, seq_length)
1729 out[kind] = KindActivations(
1730 saved=max(0.0, 2 * saved_2 - saved_1) / seq_length,
1731 saved_tp=max(0.0, 2 * (saved_1 - saved_2)) / seq_length,
1732 working=max(0.0, 2 * working_2 - working_1) / seq_length,
1733 working_tp=max(0.0, 2 * (working_1 - working_2)) / seq_length,
1734 seq_length=int(seq_length),
1735 selective=max(0.0, 2 * kept_2 - kept_1) / seq_length,
1736 selective_tp=max(0.0, 2 * (kept_1 - kept_2)) / seq_length,
1737 selective_attention_mm=attention_mm,
1738 selective_ffn_mm=ffn_mm,
1739 ops={op: max(0.0, 2 * ops_2.get(op, 0) - ops_1[op]) / seq_length for op in sorted(ops_1)},
1740 ops_tp={op: max(0.0, 2 * (ops_1[op] - ops_2.get(op, 0))) / seq_length for op in sorted(ops_1)},
1741 )
1742 return out