Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / derive.py: 97%
192 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"""The derived fields of a cost-model config, computed in one place.
17A parser states primary facts: the model's dimensions, the parallel strategy
18and the fixed facts of the run. :func:`derive` computes what follows from
19them, once when a config is parsed and again whenever its strategy changes,
20where each parser used to compute it. What a producer does not state, its
21model's family says, from its op profile: :func:`derive_family`.
22"""
23import logging
24import math
25from typing import Any, Dict, Mapping, Union
27from hyper_parallel.auto_parallel._exec_spec import RECOMPUTE_OPS
28from hyper_parallel.auto_parallel._op_profiles import VISION_ARCH, family_profile
29from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config
30from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_order import stated_recompute
32logger = logging.getLogger(__name__)
34# The recompute switches of HyperParallel's selective activation checkpointing
35# (hyper_parallel/distributed/activation_checkpoint.py). It keeps the outputs of
36# matmul and attention kernels and of reduce-scatter, all-to-all and all-reduce,
37# and recomputes everything else, all-gathers included. A switch at 1 keeps the
38# op's activation and 0 recomputes it. The policy also recomputes every other
39# projection matmul, which has no switch, so that part is not priced; an MLA
40# layer's up-projections, matmuls too, are kept.
41HYPER_SELECTIVE_REC_OP = {
42 "attBMM": 1,
43 "headCast": 0,
44 "dropout": 0,
45 "softmax": 0,
46 "normOp": 0,
47 "gather": 0,
48 "ffAct": 0,
49 "attUp": 1,
50}
53def runs_hyper_selective(ccfg: Any, switches: Any = None) -> bool:
54 """Whether *ccfg*'s selective layers run HyperParallel's policy: their switches, or *switches*, are its switches.
56 A switch they leave out keeps its op.
57 """
58 stated = getattr(ccfg, "rec_op", None) if switches is None else switches
59 values = stated if isinstance(stated, Mapping) else vars(stated) if stated is not None else {}
60 return all(values.get(name, 1) == state for name, state in HYPER_SELECTIVE_REC_OP.items())
63def derive_sequence_parallel(ccfg: Any) -> None:
64 """Set the sequence-parallel factor ``sp``: the TP degree, or 1 without."""
65 ccfg.sp = ccfg.t if ccfg.sequence_parallel else 1
68def derive_expert_degrees(ccfg: Any, strict: bool = True) -> None:
69 """Set the degrees an expert layer runs with, ``t_exp`` and ``d_exp``.
71 Args:
72 ccfg: The config, with its degrees and expert layout set.
73 strict: Refuse degrees that leave an expert group, the expert width
74 or the expert count below 1. When false, clamp them to 1
75 instead, with a warning, so that a search can still reach the
76 strategy and reject it by memory.
78 Raises:
79 TypeError: When *strict* and the degrees cannot hold the experts.
80 """
81 if ccfg.etp > 1:
82 ccfg.t_exp = ccfg.etp
83 # d * t = inner dp * outer dp * etp
84 # inner dp = EP, outer dp = the rest
85 ccfg.d_exp = ccfg.d * ccfg.t * ccfg.cp // ccfg.t_exp // ccfg.ep
86 else:
87 ccfg.t_exp = ccfg.t
88 if ccfg.d >= ccfg.ep:
89 ccfg.d_exp = ccfg.d // ccfg.ep
90 else:
91 ccfg.d_exp = ccfg.d * ccfg.t // ccfg.ep
92 if ccfg.t_exp * ccfg.ep > ccfg.d * ccfg.t:
93 ccfg.t_exp = 1
95 # The group the runtime forms, which is what every consumer wants; a
96 # strategy is impossible only when the stage holds fewer ranks than it
97 # spreads experts over. d_exp stays the older form, and is clamped
98 # because a group of one still leaves it at zero under CP.
99 ccfg.edp_group = expert_dp_group(ccfg)
100 ccfg.d_exp = max(1, ccfg.d_exp)
101 # A MoE layer dispatches its tokens once and takes them back once.
102 ccfg.n_dispatch = 1 if ccfg.n_exp > 1 else 0
104 exp_group1_invalid = ccfg.edp_group < 1 or ccfg.t_exp < 1
105 exp_group2_invalid = ccfg.hff_exp < 1 or ccfg.n_exp < 1
106 if not (exp_group1_invalid or exp_group2_invalid):
107 return
108 if strict:
109 raise TypeError(
110 f"MoE parsing error: edp_group({ccfg.edp_group})/"
111 f"t_exp({ccfg.t_exp})/"
112 f"hff_exp({ccfg.hff_exp})/n_exp({ccfg.n_exp})/"
113 f"DP = {ccfg.d}, TP = {ccfg.t}, EP = {ccfg.ep}, CP = {ccfg.cp}/"
114 )
115 logger.warning(
116 "Expert degrees invalid for d=%d t=%d ep=%d etp=%d n_exp=%d; clamping to minimum values.",
117 ccfg.d, ccfg.t, ccfg.ep, ccfg.etp, ccfg.n_exp,
118 )
119 ccfg.d_exp = max(1, ccfg.d_exp)
120 ccfg.edp_group = max(1, ccfg.edp_group)
121 ccfg.t_exp = max(1, ccfg.t_exp)
122 ccfg.hff_exp = max(1, ccfg.hff_exp)
123 ccfg.n_exp = max(1, ccfg.n_exp)
126def optimizer_ranks(ccfg: Any) -> int:
127 """How many data-parallel ranks the optimizer shards a parameter over.
129 ``os_max_shard`` counts them, as MindSpore's ``optimizer_weight_shard_size``
130 and HyperParallel's ``dp_shard`` do, on top of TP's sharding. A count
131 that does not divide DP shards over all of it, as MindSpore does.
132 HyperParallel's domain is DP times CP, which its shard spans
133 (:func:`param_cp`).
134 """
135 domain = ccfg.d * (ccfg.cp if getattr(ccfg, "shard_spans_cp", False) else 1)
136 ranks = int(ccfg.os_max_shard or 0)
137 return ranks if ranks >= 1 and domain % ranks == 0 else domain
140def param_cp(ccfg: Any) -> float:
141 """What CP's ranks divide a parameter's state by, beyond the shard stated for it.
143 The legacy schemas shard a parameter, its gradient and its optimizer
144 states over CP's ranks on top of the optimizer's, and a table they gather
145 is a CP rank's share of it. HyperParallel's ``dp_shard_size`` already
146 spans the DP x CP domain, its dense FSDP mesh being
147 ``(dp_replicate, dp_shard, tp)`` over ``dp * cp`` ranks
148 (``MeshContext.build_meshs``), and the table it gathers is whole: so 1
149 there (``shard_spans_cp``), CP elsewhere.
150 """
151 return 1 if getattr(ccfg, "shard_spans_cp", False) else ccfg.cp
154def expert_dp_group(ccfg: Any) -> int:
155 """How many ranks hold the same experts.
157 The runtime spreads the experts over a whole pipeline stage: its expert
158 mesh is ``(edp_replicate, edp_shard, ep)`` over the ``dp * cp * tp``
159 ranks of the device mesh, so ``dp * cp * tp / ep`` ranks hold the same
160 experts, and an expert tensor shard divides that group again
161 (``hyper_parallel/distributed/mesh.py``,
162 ``MeshContext._build_expert_parallel_mesh``).
164 ``d_exp`` cannot stand in for this group. It is floored by EP against
165 DP alone, so it loses the remainder when EP does not divide DP and
166 falls to zero once EP passes DP, which is what context parallelism
167 does: DP counts the stage's ranks apart from CP and TP, so a strategy
168 keeps its group while its DP shrinks with CP. Returning zero made a
169 runnable strategy look invalid, and the consumers that rebuild the
170 group as ``d_exp * cp * t_exp`` cannot express a group narrower than
171 CP.
172 """
173 expert_tp = int(ccfg.etp) if ccfg.etp > 1 else 1
174 return int(ccfg.d * ccfg.t * ccfg.cp) // int(expert_tp * ccfg.ep)
177def expert_dp_ranks(ccfg: Any) -> int:
178 """How many ranks hold the same slice of a routed expert, so reduce its gradient together.
180 The stage's ranks over EP and over the expert's tensor shard
181 ``t_exp``, whose ranks hold different slices: context parallelism's
182 ranks hold the same weights and reduce their gradients with the
183 data-parallel ones, where ``d_exp`` counts the latter alone outside an
184 expert tensor shard, and :func:`expert_dp_group` counts TP's ranks
185 whatever slice they hold.
186 """
187 return max(1, int(ccfg.d * ccfg.t * ccfg.cp) // (max(1, int(ccfg.t_exp)) * max(1, int(ccfg.ep))))
190def routed_expert_shard(ccfg: Any, ranks: int) -> int:
191 """How many ranks a routed expert's parameters and optimizer states are sharded over.
193 A run that states its expert shard, as HyperParallel's ``edp_shard_size``,
194 shards an expert under expert parallelism over as many ranks of its expert
195 data-parallel group: the stage's ranks over EP, whatever their DP, CP or
196 TP, the group of ranks that hold the same experts. Without expert
197 parallelism its FSDP shards the experts with the other parameters, over
198 the optimizer's *ranks*. Stated by no one, the optimizer shards them over
199 the whole group. A run that shards each strategy's experts over its
200 whole group (``expert_shard_group``) does so under expert parallelism
201 whatever shard it states, since the group changes with the strategy.
202 """
203 stated = getattr(ccfg, "expert_shard", None)
204 # The group as the runtime forms it, once a strategy has been laid over
205 # the config; a dense model and the parser's own doubles state none, and
206 # keep the older reconstruction from d_exp, which agrees with it
207 # wherever EP divides DP.
208 whole = max(1, int(ccfg.d * ccfg.cp * ccfg.t // max(1, int(ccfg.ep))))
209 group = max(1, int(ccfg.edp_group)) if ccfg.edp_group else whole
210 if getattr(ccfg, "expert_shard_group", False) and ccfg.ep > 1:
211 return group
212 if stated is None:
213 if not ccfg.has_op:
214 return param_cp(ccfg) * ccfg.t_exp
215 return group if ccfg.edp_group else ccfg.d_exp * ccfg.cp * ccfg.t_exp
216 if ccfg.ep > 1:
217 return math.gcd(int(stated), group)
218 return ranks * param_cp(ccfg) * ccfg.t_exp
221def derive_optimizer_sharding(ccfg: Any) -> None:
222 """Set how parameters, optimizer states and gradients are sharded.
224 With optimizer sharding, a parameter is sharded over TP and then over
225 :func:`optimizer_ranks` data-parallel ranks, and a routed expert as
226 :func:`routed_expert_shard` says. Gradients are sharded as
227 the parameters are when the run says so (``grad_shard_as_params``), as
228 FSDP shards them; else over the whole optimizer shard when the run
229 shards them (``has_grad_shard``), and over TP alone otherwise.
230 """
231 ranks = optimizer_ranks(ccfg) if ccfg.has_op else 1
232 # Non expert params, sharded over CP's ranks too where the stated shard
233 # does not span them already (param_cp)
234 ccfg.shard_p_os_non_exp_partial = ranks * ccfg.t * param_cp(ccfg)
235 ccfg.shard_p_os_non_exp = (
236 (ccfg.d if ccfg.has_op else 1) * ccfg.cp * ccfg.t
237 )
239 # Expert params
240 ccfg.shard_p_os_exp_partial = math.gcd(ccfg.n_exp, ranks * ccfg.t_exp)
241 ccfg.shard_p_os_exp = routed_expert_shard(ccfg, ranks)
243 # Gradients
244 if getattr(ccfg, "grad_shard_as_params", False):
245 grads = (ccfg.shard_p_os_non_exp_partial, ccfg.shard_p_os_exp, ccfg.shard_p_os_exp_partial)
246 elif ccfg.has_grad_shard:
247 grads = (ccfg.shard_p_os_non_exp, ccfg.shard_p_os_exp, ccfg.shard_p_os_exp_partial)
248 else:
249 grads = (ccfg.t, ccfg.t_exp, ccfg.t_exp)
250 ccfg.shard_grad_non_exp, ccfg.shard_grad_exp, ccfg.shard_grad_exp_partial = grads
253def derive_param_casts(ccfg: Any) -> None:
254 """Set whether a layer keeps the weight casts the formulas price beside its matmuls, ``keeps_param_casts``.
256 The run states it as ``param_casts``; unstated, a layer keeps them where
257 the optimizer does not shard, the formulas' convention.
258 """
259 stated = getattr(ccfg, "param_casts", None)
260 ccfg.keeps_param_casts = (not ccfg.has_op) if stated is None else bool(stated)
263def derive_comm_flags(ccfg: Any) -> None:
264 """Set the communication factors and the transitional overlaps."""
265 ccfg.comm_d_non_exp = (
266 0
267 if ((ccfg.d == 1) or not ccfg.has_op)
268 else (2 if not ccfg.has_grad_shard else 3)
269 ) # data parallel comm factor
270 # Whether another rank holds the same slice of a routed expert, as a
271 # shard or a copy, so that its gradient is reduced: d_exp left CP's
272 # ranks out and read 1 where they reduce it.
273 ccfg.comm_d_exp = (
274 0
275 if ((expert_dp_ranks(ccfg) == 1) or not ccfg.has_op)
276 else (2 if not ccfg.has_grad_shard else 3)
277 ) # data parallel comm factor
278 ccfg.comm_t = float(ccfg.t > 1) # tensor parallel comm factor
279 ccfg.comm_ep = float(
280 ccfg.ep > 1 or ccfg.n_exp > 1
281 ) # expert parallel comm factor
282 ccfg.comm_cp = float(ccfg.cp > 1) # context parallel comm factor
283 ccfg.comm_dp_overlap = 0.9 # transitional overlap, see _cost_model_variables.py
284 ccfg.comm_tp_overlap = 0.5 # transitional overlap, see _cost_model_variables.py
287def derive_embedding_sharding(ccfg: Any) -> None:
288 """Set how the embedding table is sharded, ``shard_embed``, and over how many of those
289 ranks it is gathered to compute with it, ``gather_embed``.
291 The table is split over tensor parallelism unless the vocabulary
292 embedding runs data parallel without pipelining, and over data
293 parallelism unless the config says it is not. FSDP, which shards the
294 gradients as the parameters (``grad_shard_as_params``), shards the
295 table as every parameter, over :func:`optimizer_ranks`, its FSDP group
296 and not the whole of data parallelism under HSDP, and gathers it over
297 them to compute with it.
298 """
299 tp = 1 if (ccfg.vocab_emb_dp and ccfg.p == 1) else ccfg.t
300 fsdp = getattr(ccfg, "grad_shard_as_params", False)
301 ranks = (optimizer_ranks(ccfg) if ccfg.has_op else 1) if fsdp else ccfg.d
302 ccfg.shard_embed = (ranks if ccfg.emb_dp_sharded else 1) * tp
303 ccfg.gather_embed = ranks if fsdp and ccfg.emb_dp_sharded else 1
306def hyper_rec_op(selective: Union[bool, list]) -> dict[str, int]:
307 """Recompute switches for a HyperParallel activation checkpoint mode.
309 Args:
310 selective: The parsed ``sel_rec``: truthy when the run uses
311 selective activation checkpointing.
313 Returns:
314 The switches for ``ccfg.rec_op``: :data:`HYPER_SELECTIVE_REC_OP`
315 for a selective run, and every op kept otherwise.
316 """
317 if selective:
318 return dict(HYPER_SELECTIVE_REC_OP)
319 return dict.fromkeys(HYPER_SELECTIVE_REC_OP, 1)
322def mindformers_rec_op(ccfg: Any, selective: Any = None) -> dict[str, int]:
323 """Recompute switches of MindFormers' selective recompute.
325 What ``select_recompute`` recomputes depends on flash attention and
326 sequence parallelism; ``select_comm_recompute`` recomputes the
327 sequence-parallel all-gather.
329 Args:
330 ccfg: The config, with ``sel_rec``, ``sel_comm_rec``, ``has_fa`` and
331 ``sp`` set.
332 selective: Whether the run recomputes selectively, ``sel_rec`` by
333 default.
335 Returns:
336 The switches for ``ccfg.rec_op``: 1 keeps the op's activation and 0
337 recomputes it.
338 """
339 selective = ccfg.sel_rec if selective is None else selective
340 return {
341 "attBMM": int(not (selective and not ccfg.has_fa and ccfg.sp > 1)),
342 "headCast": int(not (selective and ccfg.has_fa)),
343 "dropout": 1,
344 "softmax": int(not (selective and not ccfg.has_fa)),
345 "normOp": int(not (selective and ccfg.sp > 1)),
346 "gather": int(not (ccfg.sel_comm_rec and ccfg.sp > 1)),
347 "ffAct": int(not (selective and ccfg.sp > 1)),
348 }
351def derive_recompute_ranges(ccfg: Any) -> None:
352 """Check the recompute ranges a config states against its layers.
354 Raises:
355 ValueError: When a range reaches past the model's last layer.
356 """
357 ranges = stated_recompute(ccfg)
358 if not ranges:
359 return
360 layers = int(ccfg.n_lay + ccfg.n_mtp)
361 for item in ranges:
362 last = item.first + (item.count or 1) - 1
363 if last >= layers:
364 raise ValueError(
365 f"{ccfg.model_name}: recompute range {item.to_dict()} reaches layer {last}, "
366 f"past the model's {layers} layers"
367 )
370def derive_recompute_switches(ccfg: Any) -> None:
371 """Set which activations a recomputed layer keeps: ``rec_op``, and each layer's own, ``layer_switches``.
373 ``sel_rec_rule`` names the framework whose selective recompute the run
374 uses: HyperParallel's recomputes a fixed set of ops, MindFormers' a set
375 that depends on flash attention and sequence parallelism. A selective
376 recompute range that states its ops sets them instead. Where the
377 ranges state several selective settings, ``layer_switches`` gives each
378 layer in model order its range's, None for a layer that is not
379 selective, and the estimates read each layer's through its context,
380 the per-layer channel; ``rec_op`` holds the first range's. Otherwise
381 ``layer_switches`` is None and ``rec_op`` holds the one setting.
383 Raises:
384 ValueError: When the config names no known rule, or it is a
385 multimodal config whose ranges state several selective
386 settings.
387 """
388 ranges = stated_recompute(ccfg)
389 selective = [item for item in ranges or () if item.option == "selective"]
390 run_selective = ccfg.sel_rec if ranges is None else bool(selective)
391 if ccfg.sel_rec_rule == "hyperparallel":
392 switches = hyper_rec_op(run_selective)
393 elif ccfg.sel_rec_rule == "mindformers":
394 switches = mindformers_rec_op(ccfg, run_selective)
395 else:
396 raise ValueError(f"Unknown selective recompute rule {ccfg.sel_rec_rule!r}")
397 # Every switch, as a range's own state them: one the rule leaves out keeps its op.
398 switches = {op: switches.get(op, 1) for op in RECOMPUTE_OPS}
399 settings = {tuple(sorted((item.switches() or switches).items())) for item in selective}
400 ccfg.layer_switches = None
401 if len(settings) > 1:
402 if getattr(ccfg, "multimodal", False):
403 raise ValueError(
404 f"{ccfg.model_name}: its recompute ranges state {len(settings)} selective settings; "
405 "a multimodal config prices one"
406 )
407 ccfg.layer_switches = _per_layer_switches(ccfg, selective, switches)
408 switches = dict(selective[0].switches() or switches)
409 elif settings:
410 switches = dict(next(iter(settings)))
411 ccfg.rec_op = Config(switches)
414def _per_layer_switches(ccfg: Any, selective: list, rule: Dict[str, int]) -> tuple:
415 """Each layer's switches in model order, its selective range's or else *rule*'s; None where none covers it."""
416 layers = int(ccfg.n_lay + ccfg.n_mtp)
417 per_layer = [None] * layers
418 for item in selective:
419 stop = layers if item.count is None else min(layers, item.first + item.count)
420 per_layer[item.first:stop] = [dict(item.switches() or rule)] * max(0, stop - item.first)
421 return tuple(per_layer)
424def derive_flash_attention_factor(ccfg: Any) -> None:
425 """Set the flash attention factor, ``s_fa``.
427 The attention scores are priced over ``s / a`` rather than ``s`` when
428 flash attention is on.
429 """
430 ccfg.s_fa = ccfg.s / ccfg.a if ccfg.has_fa and ccfg.a > 0 else ccfg.s
433def family_run(ccfg: Any) -> Dict[str, Any]:
434 """Return the run defaults of the config's family, from its op profile.
436 A vision tower takes the vision profile's, but for its activation
437 sharding, which is its language model's family's.
438 """
439 arch = getattr(ccfg, "arch", None)
440 run = dict(family_profile(arch).run)
441 inherited = getattr(ccfg, "inherited_arch", None)
442 if arch == VISION_ARCH and inherited is not None:
443 run["shard_activations"] = family_profile(inherited).run["shard_activations"]
444 return run
447def _stated(ccfg: Any, name: str, run: Mapping[str, Any]) -> Any:
448 """Return the run fact *name* the config states, else its family's."""
449 value = getattr(ccfg, name, None)
450 return run[name] if value is None else value
453def derive_byte_widths(ccfg: Any, run: Mapping[str, Any]) -> None:
454 """Set the byte widths the estimators read from those the run states, else its family's.
456 Gradients take memory only under pipeline parallelism, unless the run
457 accumulates them without it too (``grad_accumulation``, which it gives as
458 ``accumulates_grads``: the search then gives PP 1 several micro-batches,
459 its accumulation steps). The optimizer
460 keeps, per parameter, its states and any copy of the parameters: a
461 layer's parameter as many states as its optimizer has, the embedding
462 and output tables' AdamW's two.
463 """
464 ccfg.accumulates_grads = bool(_stated(ccfg, "grad_accumulation", run))
465 accumulates = ccfg.p > 1 or ccfg.accumulates_grads
466 ccfg.bytes_grad = _stated(ccfg, "grad_bytes", run) if accumulates else 0
467 ccfg.bytes_os = _stated(ccfg, "optimizer_state_bytes", run)
468 main_copy = _stated(ccfg, "main_param_bytes", run)
469 ccfg.bytes_optim = _stated(ccfg, "optimizer_states", run) * ccfg.bytes_os + main_copy
470 ccfg.bytes_optim_table = 2 * ccfg.bytes_os + main_copy
471 ccfg.bytes_norm = _stated(ccfg, "norm_bytes", run)
472 ccfg.bytes_dropout = _stated(ccfg, "dropout_bytes", run)
475def derive_resharding(ccfg: Any, run: Mapping[str, Any]) -> None:
476 """Set whether FSDP frees a layer's gathered parameters once it has run, ``reshards``,
477 whether it holds each layer's reduce-scatter output until the backward ends,
478 ``defers_grads``, and whether it reduces a layer's gradients while the next
479 layer's backward runs, ``overlaps_grad_reduce``."""
480 ccfg.reshards = bool(_stated(ccfg, "reshard_params", run))
481 ccfg.defers_grads = bool(_stated(ccfg, "deferred_grad_accumulation", run))
482 ccfg.overlaps_grad_reduce = bool(_stated(ccfg, "overlapped_grad_reduce", run))
485def derive_activation_sharding(ccfg: Any, run: Mapping[str, Any]) -> None:
486 """Set how tensor parallelism shards what a layer keeps.
488 ``shard_recompute_input`` divides the input a fully recomputed layer
489 keeps, and ``shard_output_activ`` the output layer's activations. Both
490 are the TP degree when the run shards the activations between layers
491 (``shard_activations``), and 1 otherwise; a recomputed layer's input is
492 sliced over TP too when the run's recompute slices it, and the output
493 layer's activations only where the loss runs on logits sharded over the
494 vocabulary (``loss_parallel``, which derive gives as ``shards_logits``):
495 otherwise every rank gathers them whole.
496 """
497 sharded = _stated(ccfg, "shard_activations", run)
498 sliced = sharded or getattr(ccfg, "recompute_slice_activation", False)
499 ccfg.shards_logits = bool(_stated(ccfg, "loss_parallel", run))
500 ccfg.shard_recompute_input = ccfg.t if sliced else 1
501 ccfg.shard_output_activ = ccfg.t if sharded and ccfg.shards_logits else 1
504def derive_qk_norm(ccfg: Any) -> None:
505 """Count the QK-norm each layer runs, ``n_qknorm``: one where the model has one.
507 A layer kind whose attention has none, a linear one, states 0 instead.
508 """
509 ccfg.n_qknorm = 1 if getattr(ccfg, "qk_norm", False) else 0
512def derive_head_dim(ccfg: Any) -> None:
513 """Price an MLA family's heads at the value heads' width, ``dh``.
515 The model's ``v_head_dim``, else its family's. A family whose profile
516 states no value-head width keeps the head width its parser read.
517 """
518 width = family_profile(getattr(ccfg, "arch", None)).model.get("v_head_dim")
519 if width is not None:
520 ccfg.dh = getattr(ccfg, "v_head_dim", None) or width
523def derive_family(ccfg: Any) -> None:
524 """Set what the config's family decides where its producer states nothing.
526 The byte widths, the resharding and the activation sharding, from the
527 run facts the config states and else its family's (:func:`family_run`),
528 and an MLA family's head width.
529 """
530 run = family_run(ccfg)
531 derive_byte_widths(ccfg, run)
532 derive_resharding(ccfg, run)
533 derive_activation_sharding(ccfg, run)
534 derive_head_dim(ccfg)
537def derive(ccfg: Any, strict: bool = True) -> None:
538 """Compute the config's derived fields from its primary ones.
540 Args:
541 ccfg: The config, with its model, strategy and run facts set.
542 strict: As in :func:`derive_expert_degrees`.
544 Raises:
545 TypeError: When *strict* and the degrees cannot hold the experts.
546 ValueError: When the config names no known selective recompute rule.
547 """
548 derive_sequence_parallel(ccfg)
549 derive_expert_degrees(ccfg, strict)
550 derive_optimizer_sharding(ccfg)
551 derive_param_casts(ccfg)
552 derive_comm_flags(ccfg)
553 derive_embedding_sharding(ccfg)
554 derive_recompute_ranges(ccfg)
555 derive_recompute_switches(ccfg)
556 derive_flash_attention_factor(ccfg)
557 derive_qk_norm(ccfg)
558 derive_family(ccfg)