Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / _layer_census.py: 98%
696 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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_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_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 an entry states ``replace_module`` without
662 ``module_type``, or names a module or factory that cannot be
663 imported.
664 """
665 from hyper_parallel.models.replacement import ModuleReplacementSpec # 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) or entry.get("replace_module") is None or entry.get("when") is not None:
671 continue
672 if entry.get("module_type") is None:
673 raise ValueError(f"{where} states replace_module without module_type")
674 match = entry.get("match")
675 patterns = (match,) if isinstance(match, str) else tuple(match or ())
676 factory = entry["replace_module"]
677 specs.append(ModuleReplacementSpec(
678 match=patterns,
679 factory=_import(factory.get("_target_") if isinstance(factory, Mapping) else factory, where),
680 module_type=_import(entry["module_type"], where),
681 exact_type=bool(entry.get("exact_type", False)),
682 ))
683 return tuple(specs)
686class _Holder(torch.nn.Module):
687 """A module tree holding one part where a model holds it, so a rule's patterns match its path."""
689 def __init__(self, **parts: Any) -> None:
690 """Hold each part of *parts* under ``model``, a layer in a list as a model's stack holds it."""
691 super().__init__()
692 self.model = torch.nn.Module()
693 for name, part in parts.items():
694 setattr(self.model, name, torch.nn.ModuleList([part]) if name == "layers" else part)
697def _matching(spec: Any, names: Sequence[str]) -> Any:
698 """*spec* with the patterns that match a module of *names*, or ``None`` where none does."""
699 import fnmatch # pylint: disable=C0415
701 import dataclasses # pylint: disable=C0415
703 kept = tuple(pattern for pattern in spec.match
704 if any(fnmatch.fnmatchcase(name, pattern) for name in names))
705 return None if not kept else dataclasses.replace(spec, match=kept)
708def _replaced(holder: _Holder, specs: Sequence[Any]) -> None:
709 """Install on *holder* the replacements of *specs* whose patterns match one of its modules.
711 A rule's other patterns name modules no fake part holds, such as a
712 model's final norm beside a layer, and are left out rather than refused.
714 Raises:
715 CensusUnavailable: If a factory needs a library this host lacks, such
716 as an attention built on a native extension. Transformers' own
717 modules are not what the run trains, and measuring them prices a
718 layer nobody runs: a DeepSeek-V3.2 layer whose indexer runs eager
719 keeps, under the selective policy, 1192192 bytes a token against
720 the 516368 it keeps plain, its scores an fp32 tensor of one
721 sequence by another that the policy holds and no backward takes.
722 """
723 from hyper_parallel.models.replacement import ( # pylint: disable=C0415
724 apply_module_replacements,
725 compile_module_replacements,
726 )
728 names = [name for name, _ in holder.named_modules()]
729 kept = [found for found in (_matching(spec, names) for spec in specs) if found is not None]
730 if not kept:
731 return
732 try:
733 apply_module_replacements(holder, compile_module_replacements(holder, kept),
734 weights_mapping=[], context={}, capture_checkpoint_metadata=False)
735 except (ImportError, OSError) as exc:
736 raise CensusUnavailable(
737 f"this host cannot build the modules the run installs "
738 f"({[spec.factory.__name__ for spec in kept]}: {exc}); "
739 f"run the cost model where they build, or drop context.census and let the formulas price them"
740 ) from exc
743# Model families whose attention a census cannot price here, by model type,
744# with what is missing. Their layers build and run on this host, through
745# Transformers' own eager path, and that path is not the one the run takes:
746# measuring it prices a layer nobody runs, as a replacement this host cannot
747# build does (:func:`_replaced`, F53). A family belongs here only while no
748# contract in :mod:`~hyper_parallel.auto_parallel._npu_contracts` states what
749# its kernels save; stating one is what takes it out again.
750_NO_ACTIVATION_CONTRACT: Dict[str, str] = {
751 # Empty while every family with kernels of its own has its contracts
752 # stated. DeepSeek-V4 was the entry that forced this guard and now has
753 # them (:data:`_FAMILY_CONTRACTS`); the next family to arrive with a
754 # kernel and no contract belongs here rather than being priced from the
755 # path its run does not take.
756}
758# The families whose own modules call kernels this host has not got, with the
759# contracts their census runs under. A family here is measured as its run
760# behaves; a family in :data:`_NO_ACTIVATION_CONTRACT` is refused instead.
761_FAMILY_CONTRACTS = ("deepseek_v4",)
763# The families whose layers a census cannot run one at a time, by model type,
764# with why. Their census runs the model's own forward over the whole stack
765# (:func:`census_stack`), which calls each layer as the model calls it.
766_STACK_FAMILIES: Dict[str, str] = {
767 "deepseek_v4": (
768 "31 of the DeepSeek-V4.1 crop's 40 layers consume compressed keys and values, or an Indexer "
769 "key, that a source layer wrote earlier in the same forward and refuse to run without it; the "
770 "hidden state carries a hyper-connection dimension beside a mixing state; and the rotary "
771 "embedding is a pair of flavours the model picks between"
772 ),
773}
776@contextlib.contextmanager
777def _family_contracts(config: Any) -> Iterator[None]:
778 """The contracts the modules of *config*'s family run under, for what runs inside.
780 DeepSeek-V4.1's sparse attention is chosen on the device the tensors are
781 on, so a census, whose tensors are fake and on none, would take the
782 reference beside the kernel and measure a layer no run builds: the
783 attention's contract stands in. The Indexer's selection keeps nothing a
784 backward reads on either path, fused or reference, since a run frees its
785 scores' graph as soon as the indices are taken; its contract runs
786 the selection without recording that graph, and the census measures the
787 same bytes with it as without.
789 Whether this tree builds the model at all is the stack census's to
790 refuse (:func:`_model_class`): a layer's parameters do not depend on
791 which kernel runs it and are counted from the same fake layer.
792 """
793 if str(getattr(config, "model_type", "")) not in _FAMILY_CONTRACTS:
794 yield
795 return
796 from hyper_parallel.auto_parallel._npu_contracts import ( # pylint: disable=C0415
797 v41_attention_contract,
798 v41_indexer_contract,
799 )
800 rope = int(getattr(config, "qk_rope_head_dim", 0) or 0)
801 with v41_attention_contract(rope), v41_indexer_contract():
802 yield
805def _no_contract(config: Any) -> None:
806 """Refuse a census of one layer of a model whose kernels no contract states, or whose layers cannot run alone.
808 Raises:
809 CensusUnavailable: If *config*'s family is one of
810 :data:`_NO_ACTIVATION_CONTRACT`, or one of
811 :data:`_STACK_FAMILIES`, whose census is the whole stack's
812 (:func:`census_stack`). A layer's parameters do not depend on
813 which kernel runs it and are counted as ever
814 (:func:`census_parameters`); what it keeps for its backward
815 does, so this refuses rather than answer with the bytes of a
816 layer the run does not build.
817 """
818 family = str(getattr(config, "model_type", ""))
819 missing = _NO_ACTIVATION_CONTRACT.get(family)
820 if missing is not None:
821 raise CensusUnavailable(
822 f"no census of a {config.model_type} layer: {missing}; state the kernels' contracts, "
823 f"or drop context.census and read the formulas' memory as an estimate, not as a bound"
824 )
825 alone = _STACK_FAMILIES.get(family)
826 if alone is not None:
827 raise CensusUnavailable(
828 f"no census of one {family} layer, whose layers cannot be priced one at a time: {alone}. "
829 f"census_activations measures the whole stack instead (census_stack)"
830 )
833def _modeling(config: Any) -> Any:
834 """The Transformers modeling module of *config*'s model."""
835 model_type = str(config.model_type)
836 names = [model_type]
837 if model_type.endswith("_text"):
838 names.append(model_type[:-len("_text")])
839 for name in names:
840 try:
841 return importlib.import_module(f"transformers.models.{name}.modeling_{name}")
842 except ImportError:
843 continue
844 raise ValueError(f"no Transformers modeling module for model type {model_type!r}")
847def _classes(modeling: Any) -> Tuple[type, type]:
848 """The decoder layer and the language model's rotary embedding of *modeling*."""
849 classes = {name: cls for name, cls in vars(modeling).items() if isinstance(cls, type)}
850 layers = [cls for name, cls in classes.items() if name.endswith("DecoderLayer")]
851 rotary = [cls for name, cls in classes.items()
852 if name.endswith("RotaryEmbedding") and "Vision" not in name]
853 rotary.sort(key=lambda cls: "Text" not in cls.__name__)
854 if not layers or not rotary:
855 raise ValueError(f"{modeling.__name__} states no decoder layer or rotary embedding")
856 return layers[0], rotary[0]
859def tp_config(config: Any, tp: int) -> Any:
860 """*config* as one of *tp* tensor-parallel ranks holds its layers: heads and widths divided.
862 Two of the fields divided can be one field. A Transformers config class
863 may declare an ``attribute_map``, and DeepSeek-V4's maps
864 ``intermediate_size`` onto ``moe_intermediate_size``, so a model whose
865 feed-forward is all experts states one width under two names. Dividing
866 each name in turn then divides that width by *tp* twice, and the layer
867 built as a rank's share is a quarter as wide at TP 2 rather than half,
868 which makes the split the census derives from the pair too large and ND
869 over-reward tensor parallelism. Qwen3-MoE keeps the two apart, which is
870 why this stayed hidden. Each underlying attribute is divided once.
871 """
872 config = copy.deepcopy(config)
873 config.head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
874 aliases = getattr(type(config), "attribute_map", None) or {}
875 divided = set()
876 for name in _TP_FIELDS:
877 target = aliases.get(name, name)
878 if target in divided:
879 continue
880 value = getattr(config, name, None)
881 if isinstance(value, int) and not isinstance(value, bool) and value > 0:
882 setattr(config, name, max(1, value // tp))
883 divided.add(target)
884 return config
887def _grid_rotary(rotary: Any) -> bool:
888 """Whether *rotary* embeds a position on three grids, time, height and width, as a multimodal model's does.
890 Qwen3.5's and Qwen3-VL's hold the grids' ``mrope_section``; Qwen2-VL's
891 config states it.
892 """
893 if getattr(rotary, "mrope_section", None) is not None:
894 return True
895 config = getattr(rotary, "config", None)
896 return any(isinstance(getattr(config, name, None), Mapping) and "mrope_section" in getattr(config, name)
897 for name in ("rope_parameters", "rope_scaling"))
900def _positions(rotary: Any, hidden: torch.Tensor, seq_length: int) -> Tuple[torch.Tensor, Any]:
901 """The position ids of a sequence of *seq_length* tokens, and the embeddings *rotary* gives them.
903 A rotary on three grids (:func:`_grid_rotary`) is handed a row of ids
904 for each grid, as its model hands them, a text token's index on all
905 three; the layers keep the one row. Transformers 5.17 and 5.18 refuse
906 the single row the versions before them expanded themselves.
907 """
908 position_ids = torch.arange(seq_length).unsqueeze(0)
909 rotary_ids = position_ids[None].expand(3, *position_ids.shape) if _grid_rotary(rotary) else position_ids
910 with torch.no_grad():
911 return position_ids, rotary(hidden, rotary_ids)
914def _run(layer: Any, hidden: torch.Tensor, positions: Tuple[torch.Tensor, Any],
915 call: Optional[Callable[..., Any]] = None) -> torch.Tensor:
916 """One forward of *layer*, as its model calls it at *positions*, through *call* if given."""
917 position_ids, position_embeddings = positions
918 accepted = inspect.signature(layer.forward).parameters
919 kwargs = {name: value for name, value in (
920 ("position_embeddings", position_embeddings), ("position_ids", position_ids),
921 ("attention_mask", None), ("use_cache", False)) if name in accepted}
922 out = (call or layer)(hidden, **kwargs)
923 return out[0] if isinstance(out, tuple) else out
926# The function a gated-delta-rule attention of Transformers 5.15 and later
927# calls by name, which its modeling module defines; earlier versions hold the
928# rule in the module's ``chunk_gated_delta_rule``.
929_DELTA_RULE_FUNCTION = "torch_chunk_gated_delta_rule"
932def delta_rule_module(module: Any) -> bool:
933 """Whether *module* runs a gated delta rule, whichever way its Transformers version calls it.
935 Up to 5.14 the module holds the rule in ``chunk_gated_delta_rule``; from
936 5.15 its forward calls its modeling module's ``torch_chunk_gated_delta_rule``
937 by name, and the module holds nothing.
938 """
939 if hasattr(module, "chunk_gated_delta_rule"):
940 return True
941 forward = inspect.unwrap(getattr(type(module), "forward", None) or (lambda: None))
942 return _DELTA_RULE_FUNCTION in getattr(getattr(forward, "__code__", None), "co_names", ())
945@contextlib.contextmanager
946def _delta_rule_held(modeling: Any, layer: Any, rule: Callable[..., Any]) -> Iterator[None]:
947 """Have *layer*'s delta rule module run *rule* through its ``chunk_gated_delta_rule``, while the context lasts.
949 The census swaps and follows the rule through that attribute. A module
950 of Transformers 5.15 or later, which calls its modeling module's function
951 by name, gets the attribute too, and the name runs it.
953 Raises:
954 ValueError: If the layer holds more than one such module, which the
955 name could not tell apart.
956 """
957 modules = [module for module in layer.modules() if delta_rule_module(module)]
958 for module in modules:
959 module.chunk_gated_delta_rule = rule
960 by_name = getattr(modeling, _DELTA_RULE_FUNCTION, None)
961 if not modules or by_name is None:
962 yield
963 return
964 if len(modules) > 1:
965 raise ValueError(f"{type(layer).__name__} holds {len(modules)} gated delta rule modules, and "
966 f"{_DELTA_RULE_FUNCTION} cannot tell which one calls it")
967 module = modules[0]
969 def by_attribute(*args: Any, **kwargs: Any) -> Any:
970 """The module's rule as the attribute holds it when called: the census wraps it afterwards."""
971 return module.chunk_gated_delta_rule(*args, **kwargs)
973 setattr(modeling, _DELTA_RULE_FUNCTION, by_attribute)
974 try:
975 yield
976 finally:
977 setattr(modeling, _DELTA_RULE_FUNCTION, by_name)
980@contextlib.contextmanager
981def _fake_layer(config: Any, layer_index: int, gdn_backend: str = _EAGER_GDN,
982 replacements: Sequence[Any] = ()) -> Iterator[Tuple[Any, Any]]:
983 """Layer *layer_index* of *config* and its model's rotary embedding, on fake tensors and the runtime's kernels.
985 *gdn_backend* names the gated delta rule the layer runs, as the runtime
986 names it: ``"eager"`` for Transformers' own chunked implementation,
987 ``"triton"`` for HyperParallel's kernel's contract, which saves the rule's
988 inputs, the cumulated gates, beta and one chunk matrix and nothing more.
989 Eager is the default because it is the rule a run reaches: the runtime's
990 dispatcher defaults to it (``components/modules/gated_delta_net.py``,
991 ``chunk_gated_delta_rule``), the only caller that can select the kernel is
992 the context-parallel wrapper, which defaults to eager as well, and no
993 recipe states a backend. The two are not interchangeable: a layer keeping
994 the kernel's saved set keeps less than half as much.
996 Where *replacements* are given, the modules they name are HyperParallel's
997 fused ones, running under the kernels' contracts for as long as the layer
998 does. Eager is Transformers' own Python rule, never a kernel a hub or the
999 ``fla`` package would put in its place.
1001 Raises:
1002 ValueError: If *gdn_backend* names no backend the runtime has.
1003 """
1004 if gdn_backend not in _GDN_BACKENDS:
1005 raise ValueError(
1006 f"unsupported GDN backend {gdn_backend!r}; "
1007 f"expected one of {sorted(_GDN_BACKENDS)}."
1008 )
1009 modeling = _modeling(config)
1010 layer_cls, rotary_cls = _classes(modeling)
1011 config = copy.deepcopy(config)
1012 config._attn_implementation = _FLASH # pylint: disable=protected-access
1013 config._experts_implementation = "grouped_mm" # pylint: disable=protected-access
1014 modeling.ALL_ATTENTION_FUNCTIONS.register(_FLASH, _flash_attention)
1015 with npu_contracts() if replacements else contextlib.nullcontext(), \
1016 _family_contracts(config), \
1017 FakeTensorMode(allow_non_fake_inputs=True):
1018 default = torch.get_default_dtype()
1019 torch.set_default_dtype(torch.bfloat16)
1020 try:
1021 layer = layer_cls(config, layer_index)
1022 rotary = rotary_cls(config=config)
1023 finally:
1024 torch.set_default_dtype(default)
1025 rule = (_gated_delta_rule(modeling) if gdn_backend == _HYPER_GDN
1026 else inspect.unwrap(getattr(modeling, _DELTA_RULE_FUNCTION, None) or (lambda: None)))
1027 with _delta_rule_held(modeling, layer, rule):
1028 if replacements:
1029 holder = _Holder(layers=layer)
1030 _replaced(holder, replacements)
1031 layer = holder.model.layers[0]
1032 layer.train()
1033 yield layer, rotary
1036class _ByPart(TorchDispatchMode):
1037 """A walk over a layer's forward that attributes each op to the part of the layer running it."""
1039 def __init__(self, layer: Any, recurrent: bool = True) -> None:
1040 """Follow which of *layer*'s modules runs, and which tensors are its parameters.
1042 *recurrent* is whether an attention of this layer has a recurrence to
1043 name: where it has none, an op of the attention that no weight takes
1044 part in is the attention's own, its mask or its rotary. A walk that
1045 counts only matmuls leaves it at the default, because a weightless
1046 matmul under an attention is a recurrence or a score, and a score is
1047 counted on its own.
1048 """
1049 super().__init__()
1050 self.recurrent = recurrent
1051 self.params = {param.untyped_storage()._cdata # pylint: disable=protected-access
1052 for param in layer.parameters()}
1053 self.running: List[str] = []
1054 for name, module in layer.named_modules():
1055 if name:
1056 module.register_forward_pre_hook(functools.partial(self._enter, name))
1057 module.register_forward_hook(self._leave)
1059 def _enter(self, name: str, *_: Any) -> None:
1060 """Note that the module at *name* runs."""
1061 self.running.append(name)
1063 def _leave(self, *_: Any) -> None:
1064 """Note that the innermost module running has returned."""
1065 self.running.pop()
1067 def _part(self, args: Sequence[Any]) -> str:
1068 """The part an op on *args* runs for: its module's, or an attention's recurrence where no weight takes part."""
1069 part = _parameter_part(f"{self.running[-1]}.weight") if self.running else "ffn"
1070 weighted = any(isinstance(arg, torch.Tensor)
1071 and arg.untyped_storage()._cdata in self.params # pylint: disable=protected-access
1072 for arg in args)
1073 if part == "attention" and not weighted and self.recurrent:
1074 return "linrec"
1075 return part
1078class _PartFlops(_ByPart):
1079 """A layer's forward FLOPs by part: its matmuls' toward the part running them, flash attention's the scores'."""
1081 def __init__(self, layer: Any) -> None:
1082 """Count no FLOPs yet, and follow *layer*'s modules."""
1083 super().__init__(layer)
1084 self.flops: Dict[str, int] = {}
1086 def _add(self, part: str, flops: int) -> None:
1087 """Count *flops* toward *part*."""
1088 self.flops[part] = self.flops.get(part, 0) + flops
1090 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
1091 out = func(*args, **(kwargs or {}))
1092 if func is torch.ops.nd_census.flash_attention.default:
1093 query, key, value = args[:3]
1094 batch, heads, seq, width = query.shape
1095 self._add("scores", 2 * batch * heads * seq * key.shape[2] * (width + value.shape[-1]))
1096 elif flops := _matmul_flops(func, args):
1097 self._add(self._part(args), flops)
1098 return out
1101class _PartTraffic(_ByPart):
1102 """The bytes a layer's forward moves by part, the parameters among them, and the ops it dispatches.
1104 An op reads the tensors it is given and writes the tensors it returns.
1105 A view writes nothing and reads nothing: it renames bytes its producer
1106 already wrote, and no kernel runs for it, so it counts for neither. A
1107 tensor given to an op twice is read twice, as the count is of the op's
1108 arguments rather than of the storages behind them.
1110 The parameters are counted again on their own, because a strategy
1111 divides their traffic where it leaves the activations' alone: the degree
1112 sharding a parameter divides the bytes read of it, and the routed
1113 experts' part is the one the expert-parallel degree divides.
1114 """
1116 def __init__(self, layer: Any, recurrent: bool = True) -> None:
1117 """Move nothing yet, and follow *layer*'s modules."""
1118 super().__init__(layer, recurrent=recurrent)
1119 self.moved: Dict[str, float] = {}
1120 self.parameter: Dict[str, float] = {}
1121 self.launches: Dict[str, int] = {}
1123 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
1124 kwargs = kwargs or {}
1125 out = func(*args, **kwargs)
1126 if func.namespace == "prim" or getattr(func, "is_view", False):
1127 return out
1128 part = ("scores" if func is torch.ops.nd_census.flash_attention.default
1129 else self._part(args))
1130 moved = parameter = 0
1131 for tensor in tree_flatten((args, kwargs))[0]:
1132 if not isinstance(tensor, torch.Tensor):
1133 continue
1134 size = tensor.numel() * tensor.element_size()
1135 moved += size
1136 if tensor.untyped_storage()._cdata in self.params: # pylint: disable=protected-access
1137 parameter += size
1138 for tensor in tree_flatten(out)[0]:
1139 if isinstance(tensor, torch.Tensor):
1140 moved += tensor.numel() * tensor.element_size()
1141 self.moved[part] = self.moved.get(part, 0.0) + moved
1142 self.parameter[part] = self.parameter.get(part, 0.0) + parameter
1143 self.launches[part] = self.launches.get(part, 0) + 1
1144 return out
1147def census_flops(config: Any, layer_index: int, seq_length: int) -> Dict[str, int]:
1148 """The forward FLOPs of layer *layer_index* of *config* on one sequence, by part.
1150 Args:
1151 config: The language model's Transformers config.
1152 layer_index: The layer to build, which settles its kind.
1153 seq_length: Tokens of the sequence it runs.
1155 Returns:
1156 The FLOPs of each part that runs a matmul: the attention's
1157 projections (``attention``), its scores and values at every pair of
1158 tokens (``scores``), a linear attention's recurrence (``linrec``),
1159 and the parts of the feed-forward :func:`_parameter_part` names.
1160 """
1161 with _fake_layer(config, layer_index) as (layer, rotary):
1162 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)
1163 positions = _positions(rotary, hidden, seq_length)
1164 counter = _PartFlops(layer)
1165 with counter:
1166 _run(layer, hidden, positions)
1167 return counter.flops
1170def census_traffic(config: Any, layer_index: int, seq_length: int,
1171 gdn_backend: str = _EAGER_GDN) -> LayerTraffic:
1172 """What one forward of layer *layer_index* of *config* moves, by the part of the layer moving it.
1174 The quantity the time model counts is matmul FLOPs, and the cluster's
1175 kernel tables put those at 6.3 to 7.1% of a profiled step of this model.
1176 This counts the other quantity a step could be priced by, on the same
1177 forward and with no device, so that a measured round can set each
1178 beside the time it took.
1180 Args:
1181 config: The language model's Transformers config.
1182 layer_index: The layer to build, which settles its kind.
1183 seq_length: Tokens of the micro-batch of one sequence it runs.
1184 gdn_backend: The gated delta rule the run takes (:func:`_fake_layer`).
1185 It is the largest single choice here: Transformers' own chunked
1186 rule runs a matmul a chunk and the kernel's contract is one op.
1188 Returns:
1189 A :class:`LayerTraffic` of the whole layer at tensor and context
1190 parallelism 1, by the parts :func:`census_flops` prices
1191 (:class:`_PartTraffic`).
1192 """
1193 with _fake_layer(config, layer_index, gdn_backend=gdn_backend) as (layer, rotary):
1194 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)
1195 positions = _positions(rotary, hidden, seq_length)
1196 # The rule's own module is what a layer has a recurrence by, as
1197 # _fake_layer finds it to settle which rule the layer runs.
1198 recurrent = any(delta_rule_module(module) for module in layer.modules())
1199 counter = _PartTraffic(layer, recurrent=recurrent)
1200 with counter:
1201 _run(layer, hidden, positions)
1202 return LayerTraffic(counter.moved, counter.parameter, counter.launches)
1205def census_layer(config: Any, layer_index: int, seq_length: int, selective: bool = False,
1206 ops: Optional[Dict[str, int]] = None,
1207 replacements: Sequence[Any] = (),
1208 gdn_backend: str = _EAGER_GDN) -> Tuple[int, int]:
1209 """Bytes layer *layer_index* of *config* keeps for its backward, and the most its backward holds.
1211 Args:
1212 config: The language model's Transformers config.
1213 layer_index: The layer to build, which settles its kind.
1214 seq_length: Tokens of the micro-batch of one sequence it runs.
1215 selective: Whether the layer runs under HyperParallel's selective
1216 activation checkpointing, as its trainer wraps a layer.
1217 ops: Where given, and the layer runs without checkpointing, takes
1218 the bytes it saves by the op saving them (:class:`_SavedOps`).
1219 replacements: The module replacements a run installs
1220 (:func:`replacement_specs`).
1221 gdn_backend: The gated delta rule the run takes (:func:`_fake_layer`).
1223 Returns:
1224 The bytes the forward keeps for the backward, its input included,
1225 and the bytes of activations the backward holds when it holds the
1226 most, its own gradients included, less the parameters and those
1227 gradients.
1229 Raises:
1230 CensusUnavailable: If no contract states what this model's kernels
1231 save (:func:`_no_contract`).
1232 """
1233 _no_contract(config)
1234 with _fake_layer(config, layer_index, gdn_backend=gdn_backend,
1235 replacements=replacements) as (layer, rotary):
1236 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16, requires_grad=True)
1237 grad = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16)
1238 if replacements:
1239 # A fused kernel keeps a constant per device, such as the mask
1240 # the attention runs its causal sparse mode with, which the
1241 # runtime allocates once for every layer and micro-batch: one
1242 # forward allocates them before the layer's own bytes are
1243 # counted, so they count toward no layer.
1244 with torch.no_grad():
1245 _run(layer, hidden, _positions(rotary, hidden, seq_length))
1246 if not selective:
1247 saved_ops = None if ops is None else _SavedOps(layer)
1248 measured = _measure(list(layer.parameters()), (hidden,),
1249 lambda: _run(layer, hidden, _positions(rotary, hidden, seq_length)),
1250 lambda out: out.backward(grad), grad_inputs=(grad,), ops=saved_ops)
1251 if saved_ops is not None:
1252 ops.update(saved_ops.saved)
1253 return measured
1254 # The model embeds the positions once for all its layers, outside
1255 # their checkpoints, and the trainer checkpoints each layer's call.
1256 positions = _positions(rotary, hidden, seq_length)
1257 call = functools.partial(activation_memory.checkpoint, layer, swap_inputs=False,
1258 context_fn=_selective_contexts)
1259 return _measure(list(layer.parameters()), (hidden,), lambda: _run(layer, hidden, positions, call),
1260 lambda out: out.backward(grad), grad_inputs=(grad,), checkpointed=True,
1261 shared=tree_flatten(positions)[0])
1264def census_saved_ops(config: Any, layer_index: int, seq_length: int,
1265 replacements: Sequence[Any] = ()) -> Dict[str, float]:
1266 """What layer *layer_index* of *config* keeps for its backward for each op, per token, the whole layer.
1268 Args:
1269 config: The language model's Transformers config.
1270 layer_index: The layer to build, which settles its kind.
1271 seq_length: Tokens of the micro-batch of one sequence it runs.
1272 replacements: The module replacements a run installs
1273 (:func:`replacement_specs`).
1275 Returns:
1276 ``{op: bytes per token}``, by the op saving them
1277 (:class:`_SavedOps`), which sum to what the layer keeps.
1278 """
1279 ops: Dict[str, int] = {}
1280 census_layer(config, layer_index, seq_length, ops=ops, replacements=replacements)
1281 return {op: size / seq_length for op, size in ops.items()}
1284def census_recomputed(config: Any, layer_index: int, seq_length: int) -> Tuple[float, float]:
1285 """The shares of layer *layer_index*'s matmul FLOPs its backward runs again under selective checkpointing.
1287 Args:
1288 config: The language model's Transformers config.
1289 layer_index: The layer to build, which settles its kind.
1290 seq_length: Tokens of the micro-batch of one sequence it runs.
1292 Returns:
1293 The share of its attention's projections' FLOPs, and of the rest of
1294 the layer's, that HyperParallel's selective policy recomputes: it
1295 saves every other matmul's output, in the order the layer runs
1296 them, and a biased projection's always.
1297 """
1298 with _fake_layer(config, layer_index) as (layer, rotary):
1299 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16, requires_grad=True)
1300 ledger = _RecomputedMatmuls(layer)
1301 call = functools.partial(activation_memory.checkpoint, layer, swap_inputs=False,
1302 context_fn=functools.partial(_selective_contexts, ledger))
1303 with _checkpoints_on_cpu():
1304 _run(layer, hidden, _positions(rotary, hidden, seq_length), call)
1305 return ledger.share(attention=True), ledger.share(attention=False)
1308def _parameter_part(name: str) -> str:
1309 """The part of a decoder layer the parameter at path *name* belongs to, as ND prices the parts.
1311 The layer's norms and its attention's per-head query and key norms are
1312 its norms; the rest of a child named for attention is its attention.
1313 Of the feed-forward, a shared expert's parameters, the routed experts'
1314 and the router's are parts of their own, and the rest is the dense
1315 feed-forward's.
1316 """
1317 segments = name.split(".")
1318 module = segments[-2] if len(segments) > 1 else ""
1319 if "norm" in segments[0] or module in _QK_NORMS:
1320 return "norm"
1321 if "attn" in segments[0] or "attention" in segments[0]:
1322 return "attention"
1323 if any("shared" in segment for segment in segments):
1324 return "shared"
1325 if "experts" in segments:
1326 return "routed"
1327 if module in ("gate", "router"):
1328 return "router"
1329 return "ffn"
1332def census_parameters(config: Any, layer_index: int, replacements: Sequence[Any] = ()) -> Dict[str, int]:
1333 """The parameters of layer *layer_index* of *config*, by part.
1335 Args:
1336 config: The language model's Transformers config.
1337 layer_index: The layer to build, which settles its kind.
1338 replacements: The module replacements a run installs
1339 (:func:`replacement_specs`). A replacement holds the
1340 parameters its source held, fused or renamed, so it counts the
1341 same.
1343 Returns:
1344 The parameter count of each part the layer has
1345 (:func:`_parameter_part`).
1346 """
1347 parts: Dict[str, int] = {}
1348 with _fake_layer(config, layer_index, replacements=replacements) as (layer, _):
1349 for name, param in layer.named_parameters():
1350 part = _parameter_part(name)
1351 parts[part] = parts.get(part, 0) + param.numel()
1352 return parts
1355def census_final_norm(config: Any, replacements: Sequence[Any] = ()) -> int:
1356 """The parameters of *config*'s final norm, of the class of a layer's input norm."""
1357 with _fake_layer(config, 0, replacements=replacements) as (layer, _):
1358 return sum(param.numel() for param in layer.input_layernorm.parameters())
1361def census_output(config: Any, seq_length: int, replacements: Sequence[Any] = ()) -> Tuple[int, int]:
1362 """Bytes the output layer of *config* keeps for its backward, and the most its backward holds.
1364 Args:
1365 config: The language model's Transformers config.
1366 seq_length: Tokens of the micro-batch of one sequence it runs.
1367 replacements: The module replacements a run installs
1368 (:func:`replacement_specs`), of which the final norm's is the
1369 one that reaches the output layer.
1371 Returns:
1372 As :func:`census_layer`: the final norm's input included, the
1373 output table and its gradient left out.
1374 """
1375 modeling = _modeling(config)
1376 layer_cls, _ = _classes(modeling)
1377 loss_function = importlib.import_module("transformers.loss.loss_utils").ForCausalLMLoss
1378 with npu_contracts() if replacements else contextlib.nullcontext(), \
1379 FakeTensorMode(allow_non_fake_inputs=True):
1380 default = torch.get_default_dtype()
1381 torch.set_default_dtype(torch.bfloat16)
1382 try:
1383 # The final norm is of the class of a layer's input norm.
1384 norm = layer_cls(config, 0).input_layernorm
1385 head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1386 finally:
1387 torch.set_default_dtype(default)
1388 if replacements:
1389 holder = _Holder(norm=norm)
1390 _replaced(holder, replacements)
1391 norm = holder.model.norm
1392 hidden = torch.randn(1, seq_length, config.hidden_size, dtype=torch.bfloat16, requires_grad=True)
1393 labels = torch.randint(0, config.vocab_size, (1, seq_length))
1394 # The loss alone holds the logits: the trainer drops the model's
1395 # output before the backward.
1396 return _measure([*norm.parameters(), *head.parameters()], (hidden,),
1397 lambda: loss_function(head(norm(hidden)), labels, config.vocab_size),
1398 lambda loss: loss.backward())
1401def census_output_activations(config: Any, seq_length: int = 4096,
1402 replacements: Sequence[Any] = ()) -> KindActivations:
1403 """What the output layer of *config* keeps and holds, per token, with its vocabulary and half of it.
1405 Args:
1406 config: The language model's Transformers config.
1407 seq_length: The tokens the census runs the layer at.
1408 replacements: The module replacements a run installs
1409 (:func:`replacement_specs`).
1411 Returns:
1412 The layer's :class:`KindActivations`, whose TP part is the part a
1413 vocabulary-parallel loss over two ranks halves.
1414 """
1415 half = copy.deepcopy(config)
1416 half.vocab_size = max(1, config.vocab_size // 2)
1417 (saved_1, working_1), (saved_2, working_2) = (
1418 census_output(each, seq_length, replacements) for each in (config, half))
1419 return KindActivations(
1420 saved=max(0.0, 2 * saved_2 - saved_1) / seq_length,
1421 saved_tp=max(0.0, 2 * (saved_1 - saved_2)) / seq_length,
1422 working=max(0.0, 2 * working_2 - working_1) / seq_length,
1423 working_tp=max(0.0, 2 * (working_1 - working_2)) / seq_length,
1424 seq_length=int(seq_length),
1425 )
1428class _FakeScalars(TorchDispatchMode):
1429 """A scalar read of a fake tensor answered with zero, as :class:`_LiveBytes` answers it.
1431 A model can read a value out of a tensor while it builds, as
1432 DeepSeek-V4.1's Engram reads its token map's pad entry to build its hash
1433 mapping, which a fake tensor refuses. The value settles a padding id,
1434 never a shape the census counts.
1435 """
1437 def __torch_dispatch__(self, func, types, args=(), kwargs=None): # pylint: disable=unused-argument
1438 if func is _SCALAR_READ and isinstance(args[0], FakeTensor):
1439 return 0
1440 return func(*args, **(kwargs or {}))
1443def _model_class(config: Any) -> Tuple[Any, bool]:
1444 """What the trainer builds *config* as, by its ``from_config``: its family's own class, else Transformers'.
1446 The trainer's builder discovers the family's registration, asks the
1447 model registry for a class registered by the architecture the config
1448 names, and builds Transformers' own model where none is.
1450 Returns:
1451 The class, and whether it is Transformers' own.
1453 Raises:
1454 CensusUnavailable: If the config names an architecture neither the
1455 registry nor Transformers builds, as a DeepSeek-V4.1 config does
1456 on a tree that does not carry that model: Transformers'
1457 DeepSeek-V4 in its place would price a different model.
1458 """
1459 import transformers # pylint: disable=C0415
1460 from hyper_parallel.models._transformers.config_resolver import get_is_hf_model # pylint: disable=C0415
1461 from hyper_parallel.models.registry import _resolve_custom_model_cls # pylint: disable=C0415
1463 architecture = (getattr(config, "architectures", None) or [""])[0] or ""
1464 if not get_is_hf_model(config):
1465 return _resolve_custom_model_cls(architecture), False
1466 if architecture and not hasattr(transformers, architecture):
1467 raise CensusUnavailable(
1468 f"the config names {architecture}, which neither this tree's model registry nor Transformers "
1469 f"builds; price it from a tree carrying the model, or drop context.census and read the "
1470 f"formulas' memory as an estimate, not as a bound"
1471 )
1472 return transformers.AutoModelForCausalLM, True
1475class _Stack:
1476 """Which layer of a stack runs, forward and backward, and where in a tracker's log each backward runs."""
1478 def __init__(self, layers: Sequence[Any], clock: Callable[[], int]) -> None:
1479 """Follow *layers* through both passes, placing their backwards by *clock*."""
1480 self.running: List[int] = []
1481 self.windows: Dict[int, List[int]] = {}
1482 self.clock = clock
1483 self.handles: List[Any] = []
1484 for index, layer in enumerate(layers):
1485 self.handles += [
1486 layer.register_forward_pre_hook(functools.partial(self._enter, index)),
1487 layer.register_forward_hook(self._leave),
1488 layer.register_full_backward_pre_hook(functools.partial(self._enter_backward, index)),
1489 layer.register_full_backward_hook(self._leave_backward),
1490 ]
1491 self.backward: List[int] = []
1493 # A hook's return replaces what the module takes or gives, so each of
1494 # these returns None.
1496 def _enter(self, index: int, *_: Any) -> None:
1497 """Note that the layer at *index* runs its forward."""
1498 self.running.append(index)
1500 def _leave(self, *_: Any) -> None:
1501 """Note that the innermost layer running its forward has returned."""
1502 if self.running:
1503 self.running.pop()
1505 def _enter_backward(self, index: int, *_: Any) -> None:
1506 """Open the window of the layer at *index*, whose backward starts."""
1507 self.backward.append(index)
1508 self.windows.setdefault(index, [self.clock(), self.clock()])
1510 def _leave_backward(self, *_: Any) -> None:
1511 """Close the window of the layer whose backward has returned."""
1512 if self.backward:
1513 index = self.backward.pop()
1514 self.windows[index][1] = self.clock()
1516 @property
1517 def layer(self) -> Optional[int]:
1518 """The outermost layer running its forward, or None."""
1519 return self.running[0] if self.running else None
1521 def remove(self) -> None:
1522 """Remove every hook this placed."""
1523 for handle in self.handles:
1524 handle.remove()
1527def census_stack(config: Any, seq_length: int, replacements: Sequence[Any] = ()) -> List[Tuple[int, int]]:
1528 """What each layer of *config* keeps for its backward, and the most its backward holds, from one pass of the stack.
1530 For a model whose layers cannot run alone (:data:`_STACK_FAMILIES`): the
1531 model the trainer builds (:func:`_model_class`), on fake tensors and
1532 under its family's contracts, runs its own forward over a micro-batch of
1533 one sequence and then its backward, so each layer is called as its model
1534 calls it, with what earlier layers wrote, its hidden and mixing states
1535 and the rotary flavours the model computes. Each save counts toward the
1536 layer whose forward makes it first, and a layer keeps those of its saves
1537 still live when the forward ends (:func:`_measure`'s rule). Its backward
1538 holds what it keeps and the most it holds beyond every layer's saves,
1539 the parameters and their gradients while it runs. The embedding's and
1540 the final norm's saves count toward no layer: the output layer's census
1541 prices the end of the model.
1543 Args:
1544 config: The language model's Transformers config.
1545 seq_length: Tokens of the micro-batch of one sequence it runs.
1546 replacements: The module replacements a run installs
1547 (:func:`replacement_specs`).
1549 Returns:
1550 ``(kept, held)`` bytes of each decoder layer, in model order.
1552 Raises:
1553 CensusUnavailable: If this tree cannot build the model the config
1554 names (:func:`_model_class`), or a module a replacement installs
1555 (:func:`_replaced`).
1556 """
1557 model_cls, native = _model_class(config)
1558 config = copy.deepcopy(config)
1559 if native:
1560 # Transformers' own model runs its attention through the interface,
1561 # as _fake_layer's layer does, so it runs the kernel's contract.
1562 config._attn_implementation = _FLASH # pylint: disable=protected-access
1563 _modeling(config).ALL_ATTENTION_FUNCTIONS.register(_FLASH, _flash_attention)
1564 with npu_contracts() if replacements else contextlib.nullcontext(), _family_contracts(config), \
1565 FakeTensorMode(allow_non_fake_inputs=True), _FakeScalars():
1566 default = torch.get_default_dtype()
1567 torch.set_default_dtype(torch.bfloat16)
1568 try:
1569 model = model_cls.from_config(config)
1570 finally:
1571 torch.set_default_dtype(default)
1572 # The experts run as grouped matmuls, as every layer _fake_layer
1573 # builds does. Transformers falls back to its eager experts for a
1574 # class whose module it does not know, a loop over the experts with
1575 # token counts no fake tensor has; the experts read this per call.
1576 model.config._experts_implementation = "grouped_mm" # pylint: disable=protected-access
1577 if replacements:
1578 _replaced(model, replacements)
1579 model.train()
1580 backbone = model.get_decoder()
1581 params = list(model.parameters())
1582 stored = {p.untyped_storage()._cdata for p in params} # pylint: disable=protected-access
1583 live = _LiveBytes()
1584 follow = _Stack(backbone.layers, lambda: len(live.events))
1585 saved_by: Dict[Tuple[int, int], int] = {}
1587 def pack(tensor: torch.Tensor) -> torch.Tensor:
1588 """Count a tensor autograd saves, unless it is a parameter's, toward the layer saving it first."""
1589 if tensor.untyped_storage()._cdata not in stored and follow.layer is not None: # pylint: disable=W0212
1590 saved_by.setdefault(live.key(tensor), follow.layer)
1591 # Detached, as _measure returns it: a graph the forward drops is freed.
1592 return tensor.detach()
1594 grads: List[Tuple[int, int]] = []
1595 handles = [p.register_hook(lambda grad: grads.append(live.key(grad))) for p in params]
1596 kept = [0] * len(backbone.layers)
1597 try:
1598 with live, torch.autograd.graph.saved_tensors_hooks(pack, lambda tensor: tensor):
1599 out = backbone(input_ids=torch.zeros(1, seq_length, dtype=torch.long)).last_hidden_state
1600 for key, index in saved_by.items():
1601 kept[index] += live.sizes.get(key, 0)
1602 out.sum().backward()
1603 finally:
1604 follow.remove()
1605 for handle in handles:
1606 handle.remove()
1607 grads += [live.key(p.grad) for p in params if p.grad is not None]
1608 # What every layer keeps, the gradients and the parameters left out, the
1609 # rest is what a layer's backward holds beyond them while it runs.
1610 excluded = set(saved_by) | set(grads) | {key for key, _ in live.events if key[0] in stored}
1611 beyond = list(itertools.accumulate(0 if key in excluded else delta for key, delta in live.events))
1612 held = []
1613 for index, size in enumerate(kept):
1614 begin, end = follow.windows.get(index, (0, 0))
1615 held.append(size + max(beyond[begin:end], default=0))
1616 return list(zip(kept, held))
1619def _stack_activations(config: Any, layers: Iterable[Mapping[str, Any]], seq_length: int,
1620 replacements: Sequence[Any] = ()) -> Dict[str, KindActivations]:
1621 """What a layer of each kind of *layers* keeps and holds, per token, from the stack at TP 1 and TP 2.
1623 A kind's record is the mean of its layers', which is what the memory
1624 model applies to each of them: with one record a kind, the stack's total
1625 comes back whole. The parts tensor parallelism splits follow from the
1626 pair as :func:`census_activations` derives them. No selective part is
1627 stated: a selective layer of these kinds is priced by the formulas.
1628 """
1629 kinds = [str(group["kind"]) for group in layers if not group.get("mtp") for _ in range(int(group["count"]))]
1630 logger.info("census of the whole stack of %d layers at %d tokens, at TP 1 and TP 2: a %s layer "
1631 "cannot run alone", len(kinds), seq_length, config.model_type)
1632 one, two = (census_stack(tp_config(config, tp), seq_length, replacements) for tp in (1, 2))
1633 if len(one) != len(kinds):
1634 raise ValueError(f"the model builds {len(one)} decoder layers, where its spec states {len(kinds)}")
1635 out: Dict[str, KindActivations] = {}
1636 for kind in dict.fromkeys(kinds):
1637 members = [index for index, each in enumerate(kinds) if each == kind]
1638 (kept_1, held_1), (kept_2, held_2) = (
1639 tuple(sum(stack[index][part] for index in members) / len(members) for part in (0, 1))
1640 for stack in (one, two))
1641 out[kind] = KindActivations(
1642 saved=max(0.0, 2 * kept_2 - kept_1) / seq_length,
1643 saved_tp=max(0.0, 2 * (kept_1 - kept_2)) / seq_length,
1644 working=max(0.0, 2 * held_2 - held_1) / seq_length,
1645 working_tp=max(0.0, 2 * (held_1 - held_2)) / seq_length,
1646 seq_length=int(seq_length),
1647 )
1648 return out
1651def census_activations(config: Any, layers: Iterable[Mapping[str, Any]], seq_length: int = 4096,
1652 replacements: Sequence[Any] = (),
1653 gdn_backend: str = _EAGER_GDN) -> Dict[str, KindActivations]:
1654 """What a layer of each kind of *layers* keeps and holds, per token, at TP 1 and TP 2.
1656 Args:
1657 config: The language model's Transformers config.
1658 layers: The model spec's layer stack, groups in model order; MTP
1659 groups are left out.
1660 seq_length: The tokens the census runs a layer at.
1661 replacements: The module replacements a run installs
1662 (:func:`replacement_specs`).
1663 gdn_backend: The gated delta rule the run takes (:func:`_fake_layer`).
1665 Returns:
1666 Each kind's :class:`KindActivations`: of the bytes a layer at TP 2
1667 holds, the part TP splits is half the part at TP 1; what it keeps
1668 under HyperParallel's selective activation checkpointing too, the
1669 shares of its matmul FLOPs that recomputes, and what it keeps for
1670 each op, its records per op. A family whose layers cannot run
1671 alone (:data:`_STACK_FAMILIES`) is measured over its whole stack
1672 instead (:func:`census_stack`), and its records state the plain
1673 layer's bytes only.
1675 Raises:
1676 CensusUnavailable: If this tree cannot build what the run builds.
1677 """
1678 if str(getattr(config, "model_type", "")) in _STACK_FAMILIES:
1679 return _stack_activations(config, layers, seq_length, replacements)
1680 firsts: Dict[str, int] = {}
1681 index = 0
1682 for group in layers:
1683 if group.get("mtp"):
1684 continue
1685 firsts.setdefault(str(group["kind"]), index)
1686 index += int(group["count"])
1687 out: Dict[str, KindActivations] = {}
1688 for kind, layer_index in firsts.items():
1689 ops_1: Dict[str, int] = {}
1690 ops_2: Dict[str, int] = {}
1691 (saved_1, working_1), (saved_2, working_2) = (
1692 census_layer(tp_config(config, tp), layer_index, seq_length, ops=ops,
1693 replacements=replacements, gdn_backend=gdn_backend)
1694 for tp, ops in ((1, ops_1), (2, ops_2)))
1695 (kept_1, _), (kept_2, _) = (
1696 census_layer(tp_config(config, tp), layer_index, seq_length, selective=True,
1697 replacements=replacements, gdn_backend=gdn_backend)
1698 for tp in (1, 2))
1699 attention_mm, ffn_mm = census_recomputed(config, layer_index, seq_length)
1700 out[kind] = KindActivations(
1701 saved=max(0.0, 2 * saved_2 - saved_1) / seq_length,
1702 saved_tp=max(0.0, 2 * (saved_1 - saved_2)) / seq_length,
1703 working=max(0.0, 2 * working_2 - working_1) / seq_length,
1704 working_tp=max(0.0, 2 * (working_1 - working_2)) / seq_length,
1705 seq_length=int(seq_length),
1706 selective=max(0.0, 2 * kept_2 - kept_1) / seq_length,
1707 selective_tp=max(0.0, 2 * (kept_1 - kept_2)) / seq_length,
1708 selective_attention_mm=attention_mm,
1709 selective_ffn_mm=ffn_mm,
1710 ops={op: max(0.0, 2 * ops_2.get(op, 0) - ops_1[op]) / seq_length for op in sorted(ops_1)},
1711 ops_tp={op: max(0.0, 2 * (ops_1[op] - ops_2.get(op, 0))) / seq_length for op in sorted(ops_1)},
1712 )
1713 return out