« 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"""HyperParallel AutoModels ``train.yaml`` parser (Hyper V2).
17Parses the HyperParallel YAML configuration format
18(``hyper_parallel/models/qwen3_moe/recipes/train.yaml``) and
19populates a :class:`CostModelConfig` for memory estimation.
21Model hyperparameters are resolved with the same Transformers
22``AutoConfig.from_pretrained`` path used by AutoModels. Legacy
23``model.config_overrides`` remains supported for standalone search configs.
25Expected YAML structure. Only two things are required: the checkpoint to
26read the dimensions from, and the training sequence length::
28 model:
29 _target_: hyper_parallel.models._transformers.HyperAutoModelForCausalLM.from_pretrained
30 pretrained_model_name_or_path: Qwen/Qwen3-30B-A3B
31 torch_dtype: bfloat16 # optional, default bfloat16
33 dataset:
34 data_transform:
35 max_seq_len: 4096 # or the legacy data.max_seq_len
36 # absent: the model's CONTEXT LIMIT is
37 # costed, which is rarely what you meant
39Everything below is optional, and on the search path it is ignored: the
40search varies these dimensions itself, and the device count, the batch size
41and the memory budget come from ``-d``, ``-b`` and ``-M`` (or from the
42search config). Give them only to cost one fixed strategy::
44 training:
45 global_batch_size: 4
46 micro_batch_size: 1
48 accelerator:
49 tp_size: 1
50 pp_size: 1
52 fsdp_config:
53 dp_shard_size: 4
55 activation_checkpoint:
56 mode: full
58 context:
59 max_device_memory: "64GB"
60 device_num: 64
61 census: true # price each layer kind's activations
62 # from a census of a fake layer of it
63 expert_shard: group # shard each strategy's routed experts
64 # over its whole expert DP group
66``model.config_overrides`` stays supported for standalone search configs,
67and wins over anything read from the checkpoint.
68"""
69# pylint: disable=too-many-locals,too-many-statements,too-many-branches
70import logging
71from typing import Any, Dict, Optional, Tuple
73from hyper_parallel.auto_parallel.sapp_nd.nd.balancing_adapter import front_loaded_offset
74from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config, YamlObject
75from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser
76from hyper_parallel.auto_parallel.sapp_nd.nd.common.apply_exec import apply_exec
77from hyper_parallel.auto_parallel._exec_spec import ExecSpec
78from hyper_parallel.auto_parallel.sapp_nd.nd.recompute_dimension import read_activation_checkpoint_mode
79from hyper_parallel.auto_parallel._hf_model_spec import (
80 is_auto_models_schema,
81 resolve_hf_model_spec,
82)
83from hyper_parallel.auto_parallel._layer_census import replacement_specs
84from hyper_parallel.auto_parallel._layer_stack import LinearAttentionDims, resolve_layers
85from hyper_parallel.auto_parallel._model_spec import (
86 KindActivations,
87 activations_from_dict,
88 layers_from_list,
89 ops_from_dict,
90)
91from hyper_parallel.auto_parallel._op_profiles import VISION_ARCH
93logger = logging.getLogger(__name__)
95# The legacy precision keys of a model section, which the AutoModels trainer
96# hands on to the model's Transformers config and never reads, by what it
97# reads instead.
98_UNREAD_PRECISION_KEYS = {
99 "param_init_type": "the run keeps its weights in model_init_dtype, else model.torch_dtype",
100 "compute_dtype": (
101 "the run computes in fsdp_config.mix_precision.param_dtype, else model_init_dtype, "
102 "else model.torch_dtype"
103 ),
104}
107class CostModelParserHyperV2(_CostModelParser):
108 """Parser for HyperParallel native train.yaml configuration format.
110 This parser replaces the placeholder ``CostModelParserHyperparallel``
111 which was written for an older TorchTitan TOML format. It reads the
112 AutoModels Trainer schema and resolves model parameters through the
113 Transformers configuration pipeline. Legacy ``config_overrides`` input
114 remains available for standalone cost-model search files.
115 """
117 def parse(self) -> None:
118 """Main parsing entry point."""
119 self.ccfg.config_format = "yaml"
120 self.ccfg.multimodal = False
121 self.ccfg.mm_ccfgs = None
122 self.ccfg.mm_order = None
123 self._vision_spec = None
124 self._model_seq_len = 0
125 self._tie_word_embeddings = False
127 # The model, through AutoModels' Transformers pipeline.
128 self._resolve_model_config_pipeline()
130 # The run, as the train yaml states it. An expert layout the
131 # degrees cannot hold is clamped rather than refused: the search can
132 # still reach it, and rejects it by memory.
133 apply_exec(self.ccfg, self._exec_spec(), strict=False)
134 self.ccfg.moe_dispatch = self._moe_dispatch()
135 self.ccfg.overwrite_eval_functions = {}
137 # --- Multimodal split (vision-language models only) ---
138 self._resolve_multimodal()
140 def _resolve_model_config_pipeline(self):
141 """Resolve model hyperparameters to populate ``ccfg``.
143 Delegates to the shared AutoModels resolver, which reads the same
144 Transformers configuration as the Trainer and falls back to
145 ``model.config_overrides`` for standalone cost-model search files.
146 A vision-language config additionally yields a ``vision`` sub-spec,
147 held here until :meth:`_resolve_multimodal` can build its submodule.
148 The AutoModels trainer trains the model Transformers builds, which
149 has no MTP layer (:meth:`_without_mtp`), and a vision tower only
150 where its model class builds one (:meth:`_builds_vision_tower`).
151 ``context.census`` has the resolver run a census of each layer kind
152 (:meth:`_config_census`), of the modules the run's ``plan_overrides``
153 install (:meth:`_replacements`), and ``context.census_spec`` has it
154 state the census's spec of the model rather than its own.
155 """
156 spec = resolve_hf_model_spec(
157 self._model_section(), self._visual_seq_len_override(), self._census_seq_len(), self._census_spec(),
158 self._replacements(),
159 )
160 if is_auto_models_schema(self.config):
161 # A stated depth looked priced and was not, and nothing said so
162 #: a search repeats this under its ranking.
163 self.ccfg.mtp_unpriced = self._spec_int(spec, "mtp_depth")
164 if self.ccfg.mtp_unpriced:
165 logger.warning(
166 "The model states %d MTP layer(s), which the AutoModels trainer does not build: "
167 "ND prices none", self.ccfg.mtp_unpriced,
168 )
169 spec = self._without_mtp(spec)
170 self._vision_spec = spec.pop("vision", None)
171 if self._vision_spec and not self._builds_vision_tower():
172 self._vision_spec = None
173 self._tie_word_embeddings = bool(spec.get("tie_word_embeddings"))
174 # The dtype the spec states, the checkpoint's own where the section
175 # states none: the runtime loads the weights in it (_model_dtype).
176 self._spec_dtype = spec.get("torch_dtype")
177 self._apply_spec(self.ccfg, spec)
178 ops = None if spec.get("ops") is None else ops_from_dict(spec["ops"])
179 self.config_op_counts(self.ccfg, spec["arch"], ops)
180 self.config_layer_stack(self.ccfg, resolve_layers(
181 spec["arch"], layers_from_list(spec["layers"]), ops,
182 LinearAttentionDims.from_fields(spec),
183 ))
184 self._config_census(spec)
185 self._model_seq_len = self._spec_int(spec, "max_position_embeddings")
187 def _config_census(self, spec: Dict[str, Any]) -> None:
188 """Hold the census the spec states, the record of a stack of one kind, and the output layer's.
190 The memory model prices a layer whose kind has a record with it,
191 rather than with its formulas, and the output layer with its own.
192 A stack of several kinds binds each kind's record as the layer
193 priced becomes one of it (``arch_hooks``).
194 """
195 census, output = spec.get("activations"), spec.get("output_activations")
196 self.ccfg.census = activations_from_dict(census) if census else None
197 self.ccfg.output_census = KindActivations.from_dict(output, "output_activations") if output else None
198 kinds = self.ccfg.layer_stack.distinct_kinds()
199 single = len(kinds) == 1 and self.ccfg.census
200 self.ccfg.kind_activations = self.ccfg.census.get(kinds[0].name) if single else None
202 def _census_seq_len(self) -> int:
203 """The tokens to run a census of each layer kind at, 0 unless ``context.census`` asks for one.
205 A census gives bytes per token, which hardly depend on the length:
206 it runs at the dataset's, else at 4096 tokens, never at the model's
207 context limit.
208 """
209 ctx = self._get_cfg_attr(self.config, "context", Config({}))
210 if not self._get_cfg_attr(ctx, "census", False):
211 return 0
212 return self._dataset_seq_len() or 4096
214 def _census_spec(self) -> bool:
215 """Whether ``context.census_spec`` asks for the census's spec of the model rather than the resolver's."""
216 ctx = self._get_cfg_attr(self.config, "context", Config({}))
217 return bool(self._get_cfg_attr(ctx, "census_spec", False))
219 def _replacements(self) -> Tuple[Any, ...]:
220 """The module replacements the run's ``plan_overrides`` install, which a census runs its layers with."""
221 entries = self._config_to_flat_dict(self._get_cfg_attr(self.config, "plan_overrides", None))
222 return replacement_specs(entries if isinstance(entries, list) else ())
224 # The AutoModels classes that build a vision-language checkpoint's
225 # vision tower; every other named class builds the language model alone.
226 _VISION_TARGETS = ("ImageTextToText", "Vision2Seq", "ConditionalGeneration")
228 def _builds_vision_tower(self) -> bool:
229 """Whether the run's model class builds the vision tower its checkpoint has.
231 HyperParallel's AutoModels trainer builds ``model._target_``:
232 ``HyperAutoModelForImageTextToText`` builds the tower, while
233 ``HyperAutoModelForCausalLM`` and the recipes' builders build the
234 language model alone and load the tower's weights as unexpected.
235 A yaml that names no class (the trainer requires one) and the
236 legacy schema price what the checkpoint has.
237 """
238 if not is_auto_models_schema(self.config):
239 return True
240 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
241 target = self._get_cfg_attr(model_raw, "_target_", None)
242 return not target or any(name in str(target) for name in self._VISION_TARGETS)
244 @staticmethod
245 def _without_mtp(spec: Dict[str, Any]) -> Dict[str, Any]:
246 """Return *spec* without its MTP layers, as the AutoModels trainer builds the model.
248 It builds the Transformers causal LM, which has no MTP layer whatever
249 the checkpoint declares: Transformers loads their weights as
250 unexpected and never trains them (Qwen3.5's ``mtp.*``, DeepSeek-V3's
251 last layer). The legacy schema's trainer is not this one.
252 """
253 layers = spec.get("layers")
254 if isinstance(layers, list):
255 spec["layers"] = [group for group in layers if not (isinstance(group, dict) and group.get("mtp"))]
256 spec["mtp_depth"] = 0
257 return spec
259 def _model_section(self) -> Dict[str, Any]:
260 """Return the ``model`` section as a plain mapping, ``_target_`` included.
262 ``_config_to_flat_dict`` drops every key beginning with an
263 underscore, which is right for the private fields a ``Config``
264 carries and wrong for exactly one: ``_target_`` names the factory
265 that builds the model. Stripping it left the resolver reading the
266 directory beside it as a checkpoint, and for a model Transformers
267 cannot build, falling back to a hand-written description of a config
268 that was one call away.
269 """
270 section = self._get_cfg_attr(self.config, "model", Config({}))
271 model_raw = self._config_to_flat_dict(section)
272 if not isinstance(model_raw, dict):
273 return {}
274 target = self._get_cfg_attr(section, "_target_", None)
275 if target:
276 model_raw["_target_"] = str(target)
277 return model_raw
279 def _visual_seq_len_override(self) -> int:
280 """Return ``context.visual_seq_len`` when the config declares one."""
281 ctx = self._get_cfg_attr(self.config, "context", Config({}))
282 return int(self._get_cfg_attr(ctx, "visual_seq_len", 0) or 0)
284 @staticmethod
285 def _spec_int(spec: Dict[str, Any], name: str, default: int = 0) -> int:
286 """Return an integer spec field, treating ``None`` as absent."""
287 return int(spec.get(name, default) or default)
289 def _apply_spec(self, ccfg: Any, spec: Dict[str, Any]) -> None:
290 """Populate one cost-model config object from a canonical model spec.
292 Shared by the main model, the ``config_overrides`` fallback and every
293 multimodal submodule, so all three agree on field semantics.
294 """
295 self._apply_core_spec(ccfg, spec)
296 self._apply_moe_spec(ccfg, spec)
297 self._apply_scaling_spec(ccfg, spec)
299 def _apply_core_spec(self, ccfg: Any, spec: Dict[str, Any]) -> None:
300 """Map common transformer geometry and attention fields."""
301 ccfg.model_name = str(spec.get("name", "custom"))
302 ccfg.h = self._spec_int(spec, "hidden_size")
303 ccfg.n_lay = self._spec_int(spec, "num_hidden_layers")
304 ccfg.a = self._spec_int(spec, "num_attention_heads")
305 ccfg.hff = self._spec_int(spec, "intermediate_size")
306 ccfg.v = self._spec_int(spec, "vocab_size")
307 ccfg.s = self._spec_int(spec, "max_position_embeddings") or 4096
308 ccfg.n_kv = self._spec_int(spec, "num_key_value_heads") or ccfg.a
309 # Transformers exposes head_dim explicitly and it is not always
310 # hidden_size / num_attention_heads (Qwen3 has h/a = 64, head_dim 128).
311 head_dim = self._spec_int(spec, "head_dim")
312 ccfg.dh = head_dim if head_dim else (ccfg.h / ccfg.a if ccfg.a else 0)
313 ccfg.v_head_dim = self._spec_int(spec, "v_head_dim") or None
314 ccfg.qk_nope_head_dim = self._spec_int(spec, "qk_nope_head_dim") or None
315 ccfg.dc_kv = self._spec_int(spec, "kv_lora_rank")
316 ccfg.dc_q = self._spec_int(spec, "q_lora_rank")
317 ccfg.dhr = self._spec_int(spec, "qk_rope_head_dim")
318 # Qwen3.5 fuses the output gate into q_proj, doubling its width.
319 ccfg.attn_output_gate = bool(spec.get("attn_output_gate", False))
320 # Qwen3 normalizes each head's queries and keys.
321 ccfg.qk_norm = bool(spec.get("qk_norm", False))
322 # The biases and norms the spec states; unstated, the parameter
323 # formulas count their own.
324 for name in ("qkv_bias", "o_bias", "mlp_bias", "norm_bias", "shared_expert_gate"):
325 setattr(ccfg, name, None if spec.get(name) is None else bool(spec[name]))
326 ccfg.layer_norms = self._spec_int(spec, "layer_norms") or None
328 def _apply_moe_spec(self, ccfg: Any, spec: Dict[str, Any]) -> None:
329 """Map dense defaults and optional MoE fields."""
330 ccfg.n_exp = 1
331 ccfg.n_chosen_exp = 1
332 ccfg.n_shared_exp = 0
333 ccfg.hff_exp = ccfg.hff
334 ccfg.k_1st_dense = 0
335 num_exp = self._spec_int(spec, "num_experts", 1)
336 if num_exp <= 1:
337 return
338 ccfg.n_exp = num_exp
339 ccfg.n_chosen_exp = max(1, self._spec_int(spec, "num_experts_per_tok", 1))
340 ccfg.n_shared_exp = self._spec_int(spec, "num_shared_experts")
341 moe_inter = self._spec_int(spec, "moe_intermediate_size")
342 if moe_inter:
343 ccfg.hff_exp = moe_inter
344 ccfg.k_1st_dense = self._spec_int(spec, "first_k_dense_replace")
346 def _apply_scaling_spec(self, ccfg: Any, spec: Dict[str, Any]) -> None:
347 """Map MTP and feed-forward scaling fields."""
348 ccfg.n_mtp = self._spec_int(spec, "mtp_depth")
349 ccfg.multiple_of = self._spec_int(spec, "multiple_of", 256)
350 ccfg.fdm = float(spec.get("ffn_dim_multiplier", 1.0) or 1.0)
352 # -- Multimodal ----------------------------------------------------
354 def _resolve_multimodal(self) -> None:
355 """Split a vision-language model into cost-model submodules.
357 The vision tower and the language model run on one shared pipeline,
358 so both submodules inherit the parent's strategy and only their
359 geometry, sequence length and layer placement differ. The parent
360 keeps the strategy and drops to ``n_lay = 0`` so the backbone sums
361 the submodules rather than its own layer count.
362 """
363 if not self._vision_spec:
364 return
366 text_ccfg = self._clone_submodule(self.ccfg.model_name)
367 vision_ccfg = self._build_vision_submodule(self._vision_spec)
369 self.ccfg.multimodal = True
370 self.ccfg.mm_ccfgs = {"vision": vision_ccfg, "text": text_ccfg}
371 # Vision runs first; the language model drives the search space.
372 self.ccfg.mm_order = ["vision", "text"]
373 self.ccfg.mm_main = "text"
374 # Each submodule is priced by its own arch.
375 self.ccfg.hooks_dict = None
376 self.ccfg.n_lay = 0
377 self.ccfg.layer_stack = None
378 logger.info(
379 "Multimodal cost model: vision %s (%d layers, s=%d) + text %s "
380 "(%d layers, s=%d)",
381 vision_ccfg.model_name, vision_ccfg.n_lay, vision_ccfg.s,
382 text_ccfg.model_name, text_ccfg.n_lay, text_ccfg.s,
383 )
385 def _clone_submodule(self, name: str) -> Any:
386 """Return a submodule cost config seeded from the parsed parent.
388 The loop below copies references, so every mutable container a
389 submodule writes to has to be rebuilt: each submodule is priced by
390 its own family and must not see the others' overrides.
391 """
392 cc = type(self.ccfg)({})
393 for key, value in self.ccfg.__dict__.items():
394 if key in ("mm_ccfgs", "mm_order", "mm_main", "hooks_dict"):
395 continue
396 setattr(cc, key, value)
397 cc.multimodal = False
398 cc.mm_ccfgs = None
399 cc.mm_order = None
400 cc.parser = self
401 cc.model_name = name
402 cc.rec_op = Config(dict(self.ccfg.rec_op.__dict__))
403 cc.overwrite_eval_functions = dict(self.ccfg.overwrite_eval_functions)
404 cc.offset = self._even_offset()
405 return cc
407 def _build_vision_submodule(self, vision_spec: Dict[str, Any]) -> Any:
408 """Return the vision-tower submodule config.
410 The tower is dense, carries no vocabulary embedding, and consumes the
411 visual token count rather than the text sequence length. It is placed
412 entirely on the first pipeline stage, which is where the runtime puts
413 it unless an MPipe-style schedule moves it.
414 """
415 cc = self._clone_submodule(str(vision_spec.get("name", "vision")))
416 self._apply_spec(cc, vision_spec)
417 # A tower is priced with the vision profile. Its name carries the
418 # language model's type: the family it implies is the one whose
419 # activation sharding the tower takes.
420 self.config_op_counts(cc)
421 cc.inherited_arch, cc.arch = cc.arch, VISION_ARCH
422 cc.v = 0 # patch embedding, not a vocabulary table
423 cc.n_mtp = 0
424 cc.layer_binding = None
425 # The language model's census is not the tower's.
426 cc.census = cc.kind_activations = cc.output_census = None
427 self.config_layer_stack(cc, resolve_layers(
428 VISION_ARCH, layers_from_list(vision_spec.get("layers"), "vision.layers"),
429 ))
430 # The tower runs the language model's strategy with no vocabulary
431 # embedding, no MTP layer and no experts, all on the first stage.
432 # Its own facts set its derived fields, not those of the language
433 # model it was cloned from.
434 apply_exec(cc, ExecSpec(
435 vocab_emb_dp=False, tie_embeddings=False, mtp_in_offset=False, grouped_gemm=False,
436 capacity_factor=1, offset=self._front_loaded_offset(cc.n_lay), loss_parallel=True,
437 ), strict=False)
438 return cc
440 def _even_offset(self):
441 """Return a balanced offset of the shape ``is_consistent_pp_config`` wants."""
442 if self.ccfg.vp > 1:
443 return [[0] * self.ccfg.p for _ in range(self.ccfg.vp)]
444 return [0] * self.ccfg.p
446 def _front_loaded_offset(self, n_lay: int):
447 """Return an offset placing every layer on the first pipeline stage."""
448 return front_loaded_offset(n_lay, self.ccfg.p, self.ccfg.vp)
450 def _dataset_seq_len(self) -> int:
451 """The sequence length the dataset states, the legacy ``data.max_seq_len`` too, else 0.
453 The Online path states ``dataset.data_transform.max_seq_len`` and an
454 Indexed Dataset ``dataset.data_config.seq_length``; both describe the
455 dataset that runs, so they outrank the deprecated ``data.max_seq_len``.
456 """
457 data_raw = self._get_cfg_attr(self.config, "data", Config({}))
458 legacy_seq_len = self._get_cfg_attr(data_raw, "max_seq_len", 0)
460 dataset_raw = self._get_cfg_attr(self.config, "dataset", Config({}))
461 transform_raw = self._get_cfg_attr(
462 dataset_raw, "data_transform", Config({}),
463 )
464 trainer_seq_len = self._get_cfg_attr(transform_raw, "max_seq_len", 0)
465 # Only the Online path carries data_transform.max_seq_len; an Indexed
466 # Dataset states its length as data_config.seq_length.
467 data_config_raw = self._get_cfg_attr(dataset_raw, "data_config", Config({}))
468 indexed_seq_len = self._get_cfg_attr(data_config_raw, "seq_length", 0)
469 return int(trainer_seq_len or indexed_seq_len or legacy_seq_len or 0)
471 def _resolve_sequence_length(self) -> int:
472 """The training sequence length: the dataset's, else the model's limit.
474 A model section that states no limit falls back on
475 ``config_overrides.seq_length``, then on 4096.
476 """
477 model_seq_len = self._model_seq_len or int(
478 self._get_cfg_attr(self._config_overrides(), "seq_length", 0) or 0
479 )
480 stated = self._dataset_seq_len()
481 # The ranking says so beside its table, where a reader looks.
482 self.ccfg.seq_len_stated = bool(stated)
483 if not stated:
484 # The fallback is the model's context limit, orders of magnitude
485 # above any real training length on a long-context model, which
486 # makes the quadratic attention term meaningless.
487 logger.warning(
488 "no training sequence length in the config; costing the model "
489 "context limit of %d. Set dataset.data_transform.max_seq_len "
490 "(Online) or dataset.data_config.seq_length (Indexed).",
491 int(model_seq_len or 4096),
492 )
493 return int(stated or model_seq_len or 4096)
495 def _resolve_device_capacity(self) -> str:
496 """The device's memory: ``context.max_device_memory``, else 64 GB."""
497 ctx = self._get_cfg_attr(self.config, "context", Config({}))
498 device_mem_str = (
499 ctx.__dict__.get("max_device_memory", None)
500 if isinstance(ctx, (Config, YamlObject))
501 else None
502 )
503 return str(device_mem_str) if device_mem_str else "64GB"
505 def _moe_dispatch(self) -> Optional[float]:
506 """What a MoE layer's token dispatch costs on this cluster, ``context.moe_dispatch``.
508 A run that measured its own states the number, and a model whose
509 compute does not grow with EP states 0, which switches the term off;
510 unstated (None), the estimate keeps the one measured on A3
511 (``estimate.MOE_DISPATCH``).
512 """
513 ctx = self._get_cfg_attr(self.config, "context", Config({}))
514 stated = self._get_cfg_attr(ctx, "moe_dispatch", None)
515 if stated is not None and float(stated) < 0:
516 raise ValueError(f"context.moe_dispatch takes a cost, not {stated!r}")
517 return float(stated) if stated is not None else None
519 def _exec_spec(self) -> ExecSpec:
520 """State the run the train yaml describes.
522 Read after the model, whose expert count, MTP depth and context limit
523 set the run's defaults.
524 """
525 stated = self._parse_parallelism()
526 stated.update(self._parse_batch(stated["dp"], stated["pp"]))
527 stated.update(self._parse_feature_flags(stated["cp"]))
528 stated.update(self._parse_recompute())
529 stated.update(self._init_bytes())
530 stated.update(self._init_moe_run(stated["etp"]))
531 stated.update(self._init_shard())
532 stated["offset"] = self._init_offset(stated["pp"], stated["vpp"])
533 stated["seq_split"] = 1
534 # Match the MF parser: MTP layers the model declares take part in
535 # pipeline offset balancing, and without any the offset leaves them
536 # out, as the MF parser's num_nextn_predict_layers fallback does.
537 stated["mtp_in_offset"] = bool(self.ccfg.n_mtp)
538 stated["seq_length"] = self._resolve_sequence_length()
539 stated["device_memory"] = self._resolve_device_capacity()
540 return ExecSpec(**stated)
542 def _config_overrides(self) -> Any:
543 """The model section's ``config_overrides``, empty when it has none."""
544 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
545 return self._get_cfg_attr(model_raw, "config_overrides", Config({}))
547 # ── Private helpers ───────────────────────────────────────────────
549 @staticmethod
550 def _get_cfg_attr(cfg: Any, attr: str, default: Any = None) -> Any:
551 """Get an attribute from ``Config`` / ``YamlObject`` safely.
553 ``YamlObject.__getattr__`` returns ``0`` for missing attributes
554 instead of raising ``AttributeError``, which breaks Python's
555 ``getattr(obj, attr, default)`` fallback protocol. This helper
556 checks ``__dict__`` directly.
557 """
558 if isinstance(cfg, (Config, YamlObject)):
559 return cfg.__dict__.get(attr, default)
560 return getattr(cfg, attr, default)
562 @staticmethod
563 def _config_to_flat_dict(cfg: Any) -> Dict[str, Any]:
564 """Recursively convert a ``Config`` or ``YamlObject`` to a flat dict."""
565 if isinstance(cfg, (Config, YamlObject)):
566 return {k: CostModelParserHyperV2._config_to_flat_dict(v)
567 for k, v in cfg.__dict__.items()
568 if not k.startswith("_")}
569 if isinstance(cfg, (int, float, str, bool)):
570 return cfg # type: ignore[return-value]
571 if isinstance(cfg, list):
572 return [CostModelParserHyperV2._config_to_flat_dict(i) for i in cfg]
573 return cfg
575 @staticmethod
576 def _bytes_from_dtype(dtype_str: Any) -> int:
577 """Parse a dtype string (e.g. ``\"float32\"``) to byte size.
579 Returns ``4`` for float32, ``2`` for bfloat16/float16, etc.
580 Defaults to ``4`` when parsing fails.
581 """
582 import re # pylint: disable=import-outside-toplevel
583 dtype_str = str(dtype_str)
584 m = re.search(r"(\d+)", dtype_str)
585 if m:
586 return max(1, int(m.group(1)) // 8)
587 return 4
589 def _parse_parallelism(self) -> Dict[str, Any]:
590 """The degrees and sharding the AutoModels or legacy Trainer schema states."""
591 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
592 legacy_accel = self._get_cfg_attr(train_raw, "accelerator", Config({}))
593 accel = self._get_cfg_attr(self.config, "accelerator", legacy_accel)
594 fsdp = self._get_cfg_attr(self.config, "fsdp_config", Config({}))
596 stated, dp_shard = self._parse_parallel_dimensions(accel, fsdp)
597 stated.update(self._parse_sequence_parallelism(accel))
598 stated.update(self._parse_optimizer_parallelism(accel, dp_shard, stated["dp"]))
599 stated.update(self._parse_expert_sharding(fsdp))
600 return stated
602 def _parse_expert_sharding(self, fsdp) -> Dict[str, Any]:
603 """How HyperParallel's FSDP shards a routed expert under expert parallelism.
605 Over ``edp_shard_size`` ranks of its expert data-parallel group, 1 by
606 default: a run that states none keeps each of its experts whole on
607 every rank holding it. The legacy schema states nothing, and the
608 family's rule applies.
610 One ``edp_shard_size`` cannot state what a launcher does that sets it,
611 strategy by strategy, to the whole expert data-parallel group, since
612 that group changes with EP; ``context.expert_shard: group`` states it.
613 """
614 if not is_auto_models_schema(self.config):
615 return {}
616 ctx = self._get_cfg_attr(self.config, "context", Config({}))
617 rule = self._get_cfg_attr(ctx, "expert_shard", None)
618 if rule not in (None, "group"):
619 raise ValueError(f"context.expert_shard takes 'group' or nothing, not {rule!r}")
620 return {
621 "expert_shard": max(1, int(self._get_cfg_attr(fsdp, "edp_shard_size", 1) or 1)),
622 "expert_shard_group": rule == "group",
623 }
625 def _parse_parallel_dimensions(self, accel, fsdp) -> Tuple[Dict[str, Any], int]:
626 """Return the mesh's degrees, and the data shard degree."""
628 dp_shard = int(
629 self._get_cfg_attr(fsdp, "dp_shard_size", 0)
630 or self._get_cfg_attr(accel, "dp_shard", 1)
631 or 1
632 )
633 dp_replicate = int(self._get_cfg_attr(accel, "dp_replicate", 1) or 1)
634 tp = int(
635 self._get_cfg_attr(accel, "tp_size", 0)
636 or self._get_cfg_attr(accel, "tp_degree", 1)
637 or 1
638 )
639 pp = int(
640 self._get_cfg_attr(accel, "pp_size", 0)
641 or self._get_cfg_attr(accel, "pipeline_parallel_degree", 1)
642 or 1
643 )
644 cp = int(
645 self._get_cfg_attr(accel, "cp_size", 0)
646 or self._get_cfg_attr(accel, "context_parallel_degree", 1)
647 or 1
648 )
649 ep = int(
650 self._get_cfg_attr(accel, "ep_size", 0)
651 or self._get_cfg_attr(accel, "expert_parallel_degree", 1)
652 or 1
653 )
654 etp = int(self._get_cfg_attr(accel, "expert_tensor_parallel_degree", 0) or 0)
655 ep = max(ep, 1)
657 degrees = {"tp": max(1, tp), "pp": max(1, pp), "cp": max(1, cp), "ep": max(1, ep)}
658 degrees["dp"] = self._resolve_data_parallel(
659 dp_replicate, dp_shard, degrees["tp"] * degrees["pp"] * degrees["cp"]
660 )
661 degrees["etp"] = etp
662 degrees["vpp"] = max(1, int(
663 self._get_cfg_attr(accel, "pp_interleave_num", 1) or 1
664 ))
665 return degrees, dp_shard
667 def _resolve_data_parallel(self, dp_replicate: int, dp_shard: int, denom: int) -> int:
668 """Return the data-parallel degree, preferring an explicit device count.
670 ND reads ``d * t * cp * p`` back as the cluster size whenever no
671 device count is supplied on the command line. The AutoModels schema
672 has no replicate field, since the runtime derives it from the world
673 size, so without ``context.device_num`` an HSDP run would understate
674 the cluster by exactly its replicate factor.
676 Without ``context.device_num``, an AutoModels config takes the device
677 count its caller states (``run_nd -d``), as the runtime takes the
678 world size its launcher gives it. Read from ``dp_shard_size`` alone,
679 a one-node yaml searched on more devices kept its one node's degree,
680 and its micro-batch count with it, in every candidate.
682 Args:
683 dp_replicate: The replicate degree the config states.
684 dp_shard: The shard degree the config states.
685 denom: The product ``t * p * cp`` of the other degrees.
686 """
687 ctx = self._get_cfg_attr(self.config, "context", Config({}))
688 device_num = int(self._get_cfg_attr(ctx, "device_num", 0) or 0)
689 source = "context.device_num"
690 if not device_num and is_auto_models_schema(self.config):
691 device_num = int(getattr(self.ccfg, "devices", 0) or 0)
692 source = "the device count (-d)"
693 if not device_num:
694 if is_auto_models_schema(self.config) and dp_replicate == 1:
695 logger.warning(
696 "AutoModels config carries no context.device_num; assuming "
697 "dp_replicate=1 (d = dp_shard_size = %d). Pass -d/--devices "
698 "for HSDP runs.", dp_shard,
699 )
700 return max(1, dp_replicate * dp_shard)
701 if denom < 1 or device_num % denom:
702 raise ValueError(
703 f"{source}={device_num} is not divisible by "
704 f"t*p*cp={denom}"
705 )
706 return max(1, device_num // denom)
708 def _parse_sequence_parallelism(self, accel) -> Dict[str, Any]:
709 """Sequence parallelism and the pipeline scheduler."""
710 use_sp = bool(
711 self._get_cfg_attr(accel, "sequence_parallel", False)
712 or self._get_cfg_attr(accel, "use_seq_parallel", False)
713 )
714 return {
715 "sequence_parallel": use_sp,
716 "pp_schedule": str(self._get_cfg_attr(accel, "pipeline_scheduler", "1f1b")),
717 }
719 def _parse_optimizer_parallelism(self, accel, dp_shard: int, dp: int) -> Dict[str, Any]:
720 """Optimizer and gradient sharding.
722 HyperParallel's FSDP holds every gradient sharded as its parameter,
723 from the first backward to the optimizer step, whatever the pipeline
724 degree.
725 """
726 is_auto_models = is_auto_models_schema(self.config)
727 optimizer_parallel = (
728 dp_shard > 1
729 if is_auto_models
730 else bool(self._get_cfg_attr(
731 accel, "enable_parallel_optimizer", True,
732 ))
733 )
734 weight_shard = max(1, int(
735 self._get_cfg_attr(accel, "optimizer_weight_shard_size", 0)
736 ) or (dp_shard if is_auto_models else dp))
737 return {
738 "optimizer_parallel": optimizer_parallel,
739 "optimizer_shard": weight_shard,
740 "grad_shard": bool(self._get_cfg_attr(accel, "gradient_accumulation_shard", False)),
741 "grad_shard_as_params": True,
742 # Its dp_shard_size spans the DP x CP domain, so CP shards no
743 # parameter state beyond it (derive.param_cp).
744 "shard_spans_cp": is_auto_models,
745 "grad_accumulation": True,
746 # It adds each layer's reduce-scatter output to the accumulated
747 # gradient only in the root's backward hook.
748 "deferred_grad_accumulation": True,
749 # It reduce-scatters a layer's gradients while the next layer's
750 # backward runs, and the root's in its backward hook.
751 "overlapped_grad_reduce": True,
752 "reshard_params": self._reshards_params(),
753 # It gathers the weights in the compute dtype, and a layer keeps
754 # no cast of them, whatever dp_shard.
755 "param_casts": False,
756 }
758 def _reshards_params(self) -> bool:
759 """Whether HyperParallel's FSDP frees a layer's gathered parameters once it has run.
761 It does after the layer's forward and after its backward, unless the
762 run keeps them gathered through either.
763 """
764 fsdp = self._get_cfg_attr(self.config, "fsdp_config", Config({}))
765 return bool(
766 self._get_cfg_attr(fsdp, "reshard_after_forward", True)
767 and self._get_cfg_attr(fsdp, "reshard_after_backward", True)
768 )
770 def _parse_batch(self, dp: int, pp: int) -> Dict[str, Any]:
771 """Batch settings from ``training`` or legacy ``train``."""
772 legacy_train = self._get_cfg_attr(self.config, "train", Config({}))
773 train_raw = self._get_cfg_attr(self.config, "training", legacy_train)
774 micro = max(1, int(self._get_cfg_attr(train_raw, "micro_batch_size", 1) or 1))
775 num = int(self._get_cfg_attr(train_raw, "micro_batch_num", 0) or 0)
776 gbs = int(self._get_cfg_attr(train_raw, "global_batch_size", 0) or 0)
777 if num <= 0:
778 if gbs > 0 and gbs % (micro * dp) == 0:
779 num = max(1, gbs // (micro * dp))
780 else:
781 num = pp
782 return {
783 "micro_batch_size": micro,
784 "micro_batch_num": num,
785 "global_batch_size": gbs if gbs > 0 else micro * dp * num,
786 }
788 def _parse_feature_flags(self, cp: int) -> Dict[str, Any]:
789 """Training features: attention kernel, clipping, CP algorithm, optimizer."""
790 legacy_train = self._get_cfg_attr(self.config, "train", Config({}))
791 training = self._get_cfg_attr(self.config, "training", legacy_train)
792 optimizer = self._get_cfg_attr(
793 self.config,
794 "optimizer",
795 self._get_cfg_attr(legacy_train, "optimizer", Config({})),
796 )
797 max_grad_norm = float(
798 self._get_cfg_attr(training, "max_grad_norm", None)
799 or self._get_cfg_attr(optimizer, "max_grad_norm", 0.0)
800 or 0.0
801 )
802 accel = self._get_cfg_attr(
803 self.config,
804 "accelerator",
805 self._get_cfg_attr(legacy_train, "accelerator", Config({})),
806 )
807 cp_algo = self._get_cfg_attr(accel, "context_parallel_algo", None)
808 if not cp_algo:
809 cp_algo = "colossalai_cp"
810 if cp and cp > 1:
811 logger.warning(
812 "context_parallel_algo not set; defaulting to "
813 "'colossalai_cp' (Ring CP). Set "
814 "train.accelerator.context_parallel_algo explicitly "
815 "to 'ulysses_cp' if Ulysses CP is intended."
816 )
817 # The optimizer, by its target, which decides the states it keeps a
818 # parameter; the MindFormers parser names it by its type.
819 opt_name, opt_states = self.stated_optimizer(
820 self._get_cfg_attr(optimizer, "_target_", None) or self._get_cfg_attr(optimizer, "type", None)
821 )
822 return {
823 "flash_attention": True,
824 "vocab_emb_dp": True,
825 "tie_embeddings": self._tie_word_embeddings,
826 "frozen": False,
827 "grad_clip": max_grad_norm > 0,
828 # The trainer's lm_head gathers the logits whole on every TP rank
829 # unless the loss runs on them sharded.
830 "loss_parallel": bool(self._get_cfg_attr(accel, "loss_parallel", False)),
831 "cp_algo": cp_algo,
832 "optimizer": opt_name,
833 **self._optimizer_states(optimizer, opt_states),
834 }
836 def _optimizer_states(self, optimizer: Any, states: int) -> Dict[str, Any]:
837 """What HyperParallel's optimizer keeps per parameter.
839 Its AdamW keeps two moments and its Muon one momentum per matrix
840 (``stated_optimizer`` counts them), each ``zeros_like`` the gradient,
841 which FSDP casts to the stored parameter's dtype. With
842 ``fp32_main_params`` the optimizer keeps an fp32 copy of each
843 narrower parameter, and its states in fp32.
844 """
845 stored = self._stored_param_bytes()
846 fp32_main = bool(self._get_cfg_attr(optimizer, "fp32_main_params", False))
847 return {
848 "optimizer_states": states,
849 "optimizer_state_bytes": 4 if fp32_main else stored,
850 "main_param_bytes": 4 if fp32_main and stored < 4 else 0,
851 }
853 def _stored_param_bytes(self) -> int:
854 """The width FSDP stores the parameters in: the model's, whatever FSDP gathers them in."""
855 return self._bytes_from_dtype(
856 self._get_cfg_attr(self.config, "model_init_dtype", None)
857 or self._model_dtype("param_init_type")
858 )
860 def _model_dtype(self, legacy_key: str) -> str:
861 """The dtype the runtime loads the model's weights in, by name.
863 ``model.torch_dtype`` where the section states one other than
864 ``auto``; else, in a legacy config, the legacy *legacy_key*
865 (``param_init_type`` for the parameters, ``compute_dtype`` for the
866 compute), which the AutoModels trainer does not read; else the
867 dtype the spec states, which is the checkpoint's own, as the runtime
868 loads the weights under its default ``auto``; else bfloat16
869 (``_init_bytes`` warns). A model stating none was priced in float32
870 parameters doing bfloat16 arithmetic, where the runtime loads its
871 checkpoint's, mostly two bytes.
872 """
873 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
874 stated = self._get_cfg_attr(model_raw, "torch_dtype", None)
875 if stated and stated != "auto":
876 return str(stated)
877 legacy = None if is_auto_models_schema(self.config) else self._get_cfg_attr(model_raw, legacy_key, None)
878 return str(legacy or getattr(self, "_spec_dtype", None) or "bfloat16")
880 def _parse_recompute(self) -> Dict[str, Any]:
881 """Parse recompute mode.
883 Reads ``activation_checkpoint.mode`` from the AutoModels schema, with
884 the legacy ``train.gradient_checkpointing`` path as a fallback. When
885 ``config_overrides`` supplies ``full_rec`` or ``sel_rec`` (matching
886 the MF parser's ``recompute_config.recompute`` /
887 ``recompute_config.select_recompute`` fields), those values take
888 precedence so that Hyper YAML demo files can express per-stage
889 recompute lists for side-by-side comparisons with MindFormers.
890 """
891 overrides = self._config_overrides()
892 full_rec_override = self._get_cfg_attr(overrides, "full_rec", None)
893 sel_rec_override = self._get_cfg_attr(overrides, "sel_rec", None)
895 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
896 gc = self._get_cfg_attr(train_raw, "gradient_checkpointing", Config({}))
897 activation_checkpoint = self._get_cfg_attr(
898 self.config, "activation_checkpoint", Config({}),
899 )
900 # Read as the trainer reads it: an unquoted off, which YAML reads as
901 # False, is off rather than a fall through to the legacy key, and a
902 # name that is no mode is refused rather than priced as off.
903 stated, where = self._get_cfg_attr(activation_checkpoint, "mode", None), "activation_checkpoint.mode"
904 if stated is None:
905 stated = self._get_cfg_attr(gc, "activation_checkpoint", None)
906 where = "train.gradient_checkpointing.activation_checkpoint"
907 ac_mode = read_activation_checkpoint_mode(stated, where)
908 # The mode the run trains with, which a search deriving its own
909 # recompute names under its ranking where the two differ.
910 overridden = full_rec_override is not None or sel_rec_override is not None
911 self.ccfg.stated_ac_mode = None if overridden else ac_mode
912 return {
913 "full_recompute": full_rec_override if full_rec_override is not None else ac_mode == "full",
914 "selective_recompute": sel_rec_override if sel_rec_override is not None else ac_mode == "selective",
915 "selective_rule": "hyperparallel",
916 }
918 def _init_bytes(self) -> Dict[str, Any]:
919 """FP byte sizes from AutoModels or legacy dtype fields.
921 ``model_init_dtype`` is a top-level AutoModels key applied after the
922 weights are loaded, so it outranks ``model.torch_dtype`` for the
923 stored parameters but not an explicit FSDP ``param_dtype``, and the
924 run computes in the same width: FSDP gathers each parameter in its
925 ``param_dtype``, else in the width it is kept in. Where
926 nothing states a dtype, the checkpoint's own is taken
927 (:meth:`_model_dtype`), and bfloat16 with a warning where there is
928 no checkpoint either.
929 """
930 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
931 fsdp = self._get_cfg_attr(self.config, "fsdp_config", Config({}))
932 mix_precision = self._get_cfg_attr(fsdp, "mix_precision", Config({}))
933 init_dtype = self._get_cfg_attr(self.config, "model_init_dtype", None)
934 param_dtype = self._get_cfg_attr(mix_precision, "param_dtype", None)
935 param_bytes = self._bytes_from_dtype(
936 param_dtype
937 or init_dtype
938 or self._model_dtype("param_init_type")
939 )
940 auto_models = is_auto_models_schema(self.config)
941 stated = [self._get_cfg_attr(model_raw, key, None)
942 for key in (("torch_dtype",) if auto_models else ("torch_dtype", "param_init_type"))]
943 if (not any(value and value != "auto" for value in stated) and not init_dtype
944 and not getattr(self, "_spec_dtype", None)):
945 logger.warning(
946 "neither the config nor a checkpoint states the model's dtype: its weights are priced in "
947 "bfloat16; state model.torch_dtype if they are held in another"
948 )
949 for key, read_instead in _UNREAD_PRECISION_KEYS.items():
950 if auto_models and self._get_cfg_attr(model_raw, key, None):
951 logger.warning(
952 "the AutoModels trainer reads no model.%s, and neither does ND: %s", key, read_instead,
953 )
954 return {
955 "param_bytes": param_bytes,
956 # FSDP keeps each gradient in its parameter's dtype.
957 "grad_bytes": param_bytes,
958 "compute_bytes": self._bytes_from_dtype(param_dtype or init_dtype or self._model_dtype("compute_dtype")),
959 "softmax_bytes": self._bytes_from_dtype(
960 self._get_cfg_attr(model_raw, "softmax_compute_type", "float32")),
961 }
963 def _init_moe_run(self, etp: int) -> Dict[str, Any]:
964 """The run of the expert layers: grouped GEMM, capacity, expert TP.
966 A dense model runs no grouped GEMM and a capacity factor of 1. For
967 a MoE model (``n_exp > 1``) both come from ``config_overrides``, and
968 ``etp`` defaults to 1 when the YAML states none, matching the MF
969 parser's ``expert_model_parallel`` default; a dense model's ``etp=0``
970 leaves ``t_exp = t, d_exp = d``.
971 """
972 if self.ccfg.n_exp <= 1:
973 return {"grouped_gemm": False, "capacity_factor": 1}
974 overrides = self._config_overrides()
975 cap_val = self._get_cfg_attr(
976 overrides, "capacity_factor", self._get_cfg_attr(overrides, "cap_fact", None),
977 )
978 run = {
979 "grouped_gemm": bool(self._get_cfg_attr(
980 overrides, "use_gmm", self._get_cfg_attr(overrides, "gmm", True),
981 )),
982 "capacity_factor": 1 if cap_val is None else max(1, float(cap_val)),
983 }
984 if etp == 0:
985 run["etp"] = 1
986 return run
988 def _init_offset(self, pp: int, vpp: int = 1) -> Any:
989 """The pipeline offset.
991 The MF parser reads ``model.model_config.offset`` directly from the
992 YAML. When it is a list (e.g. ``[1, 1, ..., -1]``),
993 ``CostModelConfig.is_consistent_pp_config`` requires
994 ``len(offset) == pp``, so strategies whose pipeline degree differs
995 are rejected until ``GlobalConfig.adapt_config`` regenerates a
996 matching offset. A scalar ``0`` is always accepted.
998 To match the MF parser's *list*-based filtering behaviour (used by
999 DeepSeek-V3 and other models that declare an explicit offset), this
1000 parser states a list offset of length ``pp`` by default, one that
1001 places every layer (:meth:`_balanced_offset`). An explicit offset supplied via
1002 ``config_overrides.offset`` overrides this: a list is used as-is,
1003 and a non-zero int is broadcast to ``[int] * pp``.
1004 """
1005 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
1006 explicit = self._get_cfg_attr(self._config_overrides(), "offset", None)
1007 if explicit is None:
1008 explicit = self._get_cfg_attr(model_raw, "offset", None)
1009 if isinstance(explicit, list):
1010 return list(explicit)
1011 if isinstance(explicit, int):
1012 return 0 if explicit == 0 else [explicit] * pp
1013 return self._balanced_offset(pp, vpp)
1015 def _balanced_offset(self, pp: int, vpp: int = 1) -> list:
1016 """An offset of length *pp* that places every layer the pipeline balances.
1018 Each stage runs the layers per stage, and the first ones one more
1019 each until the remainder has a stage, as the search's balancing
1020 places them; all zeros where the pipeline divides the layers, or
1021 interleaves, which the search balances itself.
1022 """
1023 layers = self.ccfg.n_lay + (self.ccfg.n_mtp or 0)
1024 extra = layers % max(1, pp) if vpp <= 1 else 0
1025 return [1 if stage < extra else 0 for stage in range(pp)]
1027 def _init_shard(self) -> Dict[str, Any]:
1028 """How the embedding and the recompute input are sharded.
1030 The embedding is sharded over data parallelism, as under the MF
1031 parser. ``recompute_slice_activation`` mirrors the MF parser's
1032 ``recompute_config`` flag: when it is set (DeepSeek-V3), a recomputed
1033 layer keeps its input sliced over ``ccfg.t``, and when it is not
1034 (Qwen), whole. :func:`derive` computes the sharding factors from
1035 them, and from the activation sharding of the model's family, which
1036 for Qwen shards them over ``ccfg.t`` in any case.
1037 """
1038 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
1039 gc = self._get_cfg_attr(train_raw, "gradient_checkpointing", Config({}))
1040 fsdp = self._get_cfg_attr(self.config, "fsdp_config", Config({}))
1041 ac = self._get_cfg_attr(self.config, "activation_checkpoint", Config({}))
1042 return {
1043 "emb_dp_sharded": True,
1044 "recompute_slice_activation": bool(self._get_cfg_attr(
1045 ac,
1046 "recompute_slice_activation",
1047 self._get_cfg_attr(
1048 fsdp,
1049 "recompute_slice_activation",
1050 self._get_cfg_attr(gc, "recompute_slice_activation", False),
1051 ),
1052 )),
1053 "shard_mtp_param": True,
1054 }