« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""HyperParallel native train.yaml parser (Hyper V2).
17Parses the HyperParallel YAML configuration format
18(``examples/qwen3_5_0_8b_base/train.yaml``) and populates
19a :class:`CostModelConfig` for memory estimation.
21Model hyperparameters are resolved by importing the model module's
22``_build_config`` function directly (by naming convention) rather than
23through the ``ModelSpec`` registry. This avoids coupling the
24memory-estimation pipeline to the model registration interface.
26Expected YAML structure::
28 model:
29 name: qwen3_5 # model spec key
30 config_overrides:
31 num_hidden_layers: 4
32 hidden_size: 3584
33 ...
35 train:
36 accelerator:
37 dp_shard: 4 # FSDP shard degree
38 gradient_checkpointing:
39 activation_checkpoint: none # none | full | selective
40 global_batch_size: 4
41 micro_batch_size: 1
43 data:
44 max_seq_len: 64
45"""
46# pylint: disable=too-many-locals,too-many-statements,too-many-branches
47import logging
48from typing import Any, Dict
50from hyper_parallel.auto_parallel.sapp_nd.nd.common.config import Config, YamlObject
51from hyper_parallel.auto_parallel.sapp_nd.nd.common.framework_parsers._cost_model_parser import _CostModelParser
52from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
54logger = logging.getLogger(__name__)
57class CostModelParserHyperV2(_CostModelParser):
58 """Parser for HyperParallel native train.yaml configuration format.
60 This parser replaces the placeholder ``CostModelParserHyperparallel``
61 which was written for an older TorchTitan TOML format. It reads the
62 current Hyper YAML schema (``model.name`` + ``train.accelerator``)
63 and resolves model parameters by importing ``_build_config`` directly
64 from the model module. When the model module is not available it
65 falls back to reading hyperparameters from ``config_overrides``.
66 """
68 def parse(self) -> None:
69 """Main parsing entry point."""
70 self.ccfg.config_format = "yaml"
71 self.ccfg.multimodal = False
72 self.ccfg.mm_ccfgs = None
73 self.ccfg.mm_order = None
75 # Resolve model hyperparameters via Hyper's own config pipeline
76 self._resolve_model_config_pipeline()
78 # --- Parallelism ---
79 self._parse_parallelism()
81 # --- Batch ---
82 self._parse_batch()
84 # --- Feature flags ---
85 self._parse_feature_flags()
87 # Flash attention factor
88 if self.ccfg.has_fa and self.ccfg.a > 0:
89 self.ccfg.s_fa = self.ccfg.s / self.ccfg.a
90 else:
91 self.ccfg.s_fa = self.ccfg.s
93 # --- Recompute ---
94 self._parse_recompute()
96 # --- n_s_split ---
97 self.ccfg.n_s_split = 1
99 # --- Bytes ---
100 self._init_bytes()
102 # --- Post-processing ---
103 self._init_moe_strategy()
104 self.config_optimizer_shard(self.ccfg)
105 self.config_comm_flag(self.ccfg)
106 self._init_shard()
107 self.ccfg.layer_custom_config = [(self.ccfg.n_lay + self.ccfg.n_mtp, None)]
108 self._init_offset()
109 self.ccfg.overwrite_eval_functions = {}
111 def _resolve_model_config_pipeline(self):
112 """Resolve model hyperparameters to populate ``ccfg``.
114 Tries to import ``_build_config`` directly from the model module.
115 On failure, falls back to reading hyperparameters from
116 ``config_overrides``. This avoids coupling the ModelSpec
117 registration interface to the memory-estimation pipeline.
118 """
119 try:
120 self._resolve_via_direct_import()
121 except (ValueError, TypeError) as exc:
122 logger.debug(
123 "Direct _build_config failed (%s); "
124 "falling back to config_overrides", exc
125 )
126 self._resolve_from_config_overrides()
128 def _resolve_via_direct_import(self):
129 """Import ``_build_config`` directly from the model module.
131 Each model module may expose an internal ``_build_config(cfg)``
132 function (by naming convention) that returns a model-specific
133 Config object. The parser imports it directly rather than
134 going through the ``ModelSpec`` registry.
136 Raises:
137 ValueError: If the model module or ``_build_config`` is not found.
138 """
139 # pylint: disable=import-outside-toplevel
140 from hyper_parallel.trainer.config import (
141 _instantiate_recursive, HyperTrainerConfig,
142 )
144 config_dict = self._config_to_flat_dict(self.config)
145 trainer_cfg = _instantiate_recursive(HyperTrainerConfig, config_dict)
146 self.ccfg.model_name = trainer_cfg.model.name
148 import importlib # pylint: disable=import-outside-toplevel
149 try:
150 mod = importlib.import_module(
151 f"hyper_parallel.models.{trainer_cfg.model.name}")
152 except (ImportError, OSError) as exc:
153 raise ValueError(
154 f"Model module '{trainer_cfg.model.name}' not found: {exc}"
155 ) from exc
157 build_config_fn = getattr(mod, "_build_config", None)
158 if build_config_fn is None:
159 raise ValueError(
160 f"Model module '{trainer_cfg.model.name}' has no "
161 f"_build_config function. The parser falls back to "
162 f"config_overrides."
163 )
165 model_config = build_config_fn(trainer_cfg)
166 self._map_model_config_to_ccfg(model_config)
168 def _map_model_config_to_ccfg(self, model_config) -> None:
169 """Map model-specific Config object fields to ``ccfg``."""
170 self.ccfg.h = int(model_config.hidden_size)
171 self.ccfg.n_lay = int(model_config.num_hidden_layers)
172 self.ccfg.a = int(model_config.num_attention_heads)
173 self.ccfg.hff = int(model_config.intermediate_size)
174 self.ccfg.v = int(model_config.vocab_size)
175 self.ccfg.s = int(model_config.max_position_embeddings)
177 self.ccfg.n_kv = int(getattr(model_config, "num_key_value_heads", 0))
178 if not self.ccfg.n_kv:
179 self.ccfg.n_kv = self.ccfg.a
180 self.ccfg.dh = self.ccfg.h / self.ccfg.a
181 self.ccfg.dc_kv = int(getattr(model_config, "kv_lora_rank", 0))
182 self.ccfg.dc_q = int(getattr(model_config, "q_lora_rank", 0))
183 self.ccfg.dhr = int(getattr(model_config, "qk_rope_head_dim", 0))
185 self.ccfg.n_exp = 1
186 self.ccfg.n_chosen_exp = 1
187 self.ccfg.n_shared_exp = 0
188 self.ccfg.hff_exp = self.ccfg.hff
189 self.ccfg.cap_fact = 1
190 self.ccfg.t_exp = self.ccfg.t
191 self.ccfg.d_exp = self.ccfg.d
192 self.ccfg.gmm = False
193 self.ccfg.k_1st_dense = 0
194 num_exp = int(getattr(model_config, "num_experts", 1))
195 if num_exp > 1:
196 self.ccfg.n_exp = max(1, num_exp)
197 self.ccfg.n_chosen_exp = max(
198 1, int(getattr(model_config, "num_experts_per_tok", 1)))
199 self.ccfg.n_shared_exp = int(
200 getattr(model_config, "n_shared_experts",
201 getattr(model_config, "num_shared_experts", 0)))
202 moe_inter = int(getattr(model_config, "moe_intermediate_size", 0))
203 if moe_inter:
204 self.ccfg.hff_exp = moe_inter
205 self.ccfg.k_1st_dense = int(
206 getattr(model_config, "first_k_dense_replace", 0))
207 self.ccfg.gmm = True
209 self.ccfg.n_mtp = int(getattr(model_config, "mtp_depth", 0))
210 # Match the MF parser: when ``mtp_depth > 0`` the MTP layers
211 # participate in pipeline offset balancing (default True); when there
212 # is no MTP (``n_mtp == 0``) they are excluded, mirroring the MF
213 # parser's ``num_nextn_predict_layers`` fallback which sets
214 # ``is_mtp_in_offset = False``.
215 self.ccfg.is_mtp_in_offset = bool(self.ccfg.n_mtp)
216 self.ccfg.multiple_of = int(getattr(model_config, "multiple_of", 256))
217 self.ccfg.fdm = float(getattr(model_config, "ffn_dim_multiplier", 1.0))
219 self._resolve_device_capacity()
221 def _resolve_from_config_overrides(self) -> None:
222 """Populate ``ccfg`` directly from ``config_overrides``.
224 Used when the model is not registered in ``ModelSpec``.
225 """
226 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
227 overrides = self._get_cfg_attr(model_raw, "config_overrides", Config({}))
228 data_raw = self._get_cfg_attr(self.config, "data", Config({}))
230 self.ccfg.model_name = str(
231 self._get_cfg_attr(model_raw, "name", "custom"))
232 self.ccfg.h = int(self._get_cfg_attr(overrides, "hidden_size", 0))
233 self.ccfg.n_lay = int(self._get_cfg_attr(overrides, "num_hidden_layers", 0))
234 self.ccfg.a = int(self._get_cfg_attr(overrides, "num_attention_heads", 0))
235 self.ccfg.hff = int(self._get_cfg_attr(overrides, "intermediate_size", 0))
236 self.ccfg.v = int(self._get_cfg_attr(overrides, "vocab_size", 0))
238 # seq_len: data.max_seq_len > overrides > default
239 self.ccfg.s = int(
240 self._get_cfg_attr(data_raw, "max_seq_len", 0)
241 or self._get_cfg_attr(overrides, "max_position_embeddings", 0)
242 or self._get_cfg_attr(overrides, "seq_length", 0)
243 or 4096
244 )
246 self.ccfg.n_kv = int(
247 self._get_cfg_attr(overrides, "num_key_value_heads", 0))
248 if not self.ccfg.n_kv:
249 self.ccfg.n_kv = self.ccfg.a
250 self.ccfg.dh = self.ccfg.h / self.ccfg.a if self.ccfg.a else 0
251 self.ccfg.dc_kv = int(
252 self._get_cfg_attr(overrides, "kv_lora_rank", 0))
253 self.ccfg.dc_q = int(
254 self._get_cfg_attr(overrides, "q_lora_rank", 0))
255 self.ccfg.dhr = int(
256 self._get_cfg_attr(overrides, "qk_rope_head_dim", 0))
258 # MoE
259 self.ccfg.n_exp = 1
260 self.ccfg.n_chosen_exp = 1
261 self.ccfg.n_shared_exp = 0
262 self.ccfg.hff_exp = self.ccfg.hff
263 self.ccfg.cap_fact = 1
264 self.ccfg.t_exp = self.ccfg.t
265 self.ccfg.d_exp = self.ccfg.d
266 self.ccfg.gmm = False
267 self.ccfg.k_1st_dense = 0
268 num_exp = int(self._get_cfg_attr(overrides, "num_experts", 1))
269 if num_exp > 1:
270 self.ccfg.n_exp = max(1, num_exp)
271 self.ccfg.n_chosen_exp = max(
272 1, int(self._get_cfg_attr(overrides, "num_experts_per_tok", 1)))
273 self.ccfg.n_shared_exp = int(
274 self._get_cfg_attr(overrides, "num_shared_experts", 0))
275 moe_inter = int(
276 self._get_cfg_attr(overrides, "moe_intermediate_size", 0))
277 if moe_inter:
278 self.ccfg.hff_exp = moe_inter
279 self.ccfg.k_1st_dense = int(
280 self._get_cfg_attr(overrides, "first_k_dense_replace", 0))
281 cap_val = (
282 self._get_cfg_attr(overrides, "capacity_factor", None)
283 or self._get_cfg_attr(overrides, "cap_fact", None)
284 )
285 if cap_val is not None:
286 self.ccfg.cap_fact = max(1, float(cap_val))
287 self.ccfg.gmm = (
288 self._get_cfg_attr(overrides, "use_gmm", False)
289 or self._get_cfg_attr(overrides, "gmm", False)
290 )
292 self.ccfg.n_mtp = int(self._get_cfg_attr(overrides, "mtp_depth", 0))
293 # Match the MF parser's MTP-in-offset semantics; see
294 # _map_model_config_to_ccfg for the rationale.
295 self.ccfg.is_mtp_in_offset = bool(self.ccfg.n_mtp)
296 self.ccfg.multiple_of = int(
297 self._get_cfg_attr(overrides, "multiple_of", 256))
298 self.ccfg.fdm = float(
299 self._get_cfg_attr(overrides, "ffn_dim_multiplier", 1.0))
301 self._resolve_device_capacity()
303 def _resolve_device_capacity(self) -> None:
304 """Set device capacity from config or default (64 GB)."""
305 ctx = self._get_cfg_attr(self.config, "context", Config({}))
306 device_mem_str = ctx.__dict__.get("max_device_memory", None) if isinstance(ctx, (Config, YamlObject)) else None
307 if device_mem_str:
308 self.ccfg.device_capacity = Memory.from_string(str(device_mem_str))
309 else:
310 self.ccfg.device_capacity = Memory.from_string("64GB")
312 # ── Private helpers ───────────────────────────────────────────────
314 @staticmethod
315 def _get_cfg_attr(cfg: Any, attr: str, default: Any = None) -> Any:
316 """Get an attribute from ``Config`` / ``YamlObject`` safely.
318 ``YamlObject.__getattr__`` returns ``0`` for missing attributes
319 instead of raising ``AttributeError``, which breaks Python's
320 ``getattr(obj, attr, default)`` fallback protocol. This helper
321 checks ``__dict__`` directly.
322 """
323 if isinstance(cfg, (Config, YamlObject)):
324 return cfg.__dict__.get(attr, default)
325 return getattr(cfg, attr, default)
327 @staticmethod
328 def _config_to_flat_dict(cfg: Any) -> Dict[str, Any]:
329 """Recursively convert a ``Config`` or ``YamlObject`` to a flat dict."""
330 if isinstance(cfg, (Config, YamlObject)):
331 return {k: CostModelParserHyperV2._config_to_flat_dict(v)
332 for k, v in cfg.__dict__.items()
333 if not k.startswith("_")}
334 if isinstance(cfg, (int, float, str, bool)):
335 return cfg # type: ignore[return-value]
336 if isinstance(cfg, list):
337 return [CostModelParserHyperV2._config_to_flat_dict(i) for i in cfg]
338 return cfg
340 @staticmethod
341 def _bytes_from_dtype(dtype_str: Any) -> int:
342 """Parse a dtype string (e.g. ``\"float32\"``) to byte size.
344 Returns ``4`` for float32, ``2`` for bfloat16/float16, etc.
345 Defaults to ``4`` when parsing fails.
346 """
347 import re # pylint: disable=import-outside-toplevel
348 dtype_str = str(dtype_str)
349 m = re.search(r"(\d+)", dtype_str)
350 if m:
351 return max(1, int(m.group(1)) // 8)
352 return 4
354 def _parse_parallelism(self):
355 """Extract parallelism settings from ``train.accelerator``."""
356 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
357 accel = self._get_cfg_attr(train_raw, "accelerator", Config({}))
359 dp_shard = int(self._get_cfg_attr(accel, "dp_shard", 1) or 1)
360 dp_replicate = int(self._get_cfg_attr(accel, "dp_replicate", 1) or 1)
361 tp = int(self._get_cfg_attr(accel, "tp_degree", 1) or 1)
362 pp = int(self._get_cfg_attr(accel, "pipeline_parallel_degree", 1) or 1)
363 cp = int(self._get_cfg_attr(accel, "context_parallel_degree", 1) or 1)
364 ep = int(self._get_cfg_attr(accel, "expert_parallel_degree", 1) or 1)
365 etp = int(self._get_cfg_attr(accel, "expert_tensor_parallel_degree", 0) or 0)
366 ep = max(ep, 1)
368 # FSDP: d = replicate * shard (when shard > 1)
369 self.ccfg.d = max(1, dp_replicate * dp_shard)
370 self.ccfg.t = max(1, tp)
371 self.ccfg.p = max(1, pp)
372 self.ccfg.cp = max(1, cp)
373 self.ccfg.ep = max(1, ep)
374 self.ccfg.sp = self.ccfg.t # Sequence parallel factor
375 self.ccfg.etp = etp
376 self.ccfg.vp = max(1, int(
377 self._get_cfg_attr(accel, "pp_interleave_num", 1) or 1
378 ))
379 use_sp = bool(self._get_cfg_attr(accel, "use_seq_parallel", False))
380 self.ccfg.sp = self.ccfg.t if use_sp else 1
381 self.ccfg.pp_sched = str(
382 self._get_cfg_attr(accel, "pipeline_scheduler", "1f1b")
383 )
385 # Optimizer parallel sharding
386 self.ccfg.has_op = bool(self._get_cfg_attr(accel,
387 "enable_parallel_optimizer",
388 True))
389 self.ccfg.op_weight_shard = max(1, int(
390 self._get_cfg_attr(accel, "optimizer_weight_shard_size", 0)
391 ) or (self.ccfg.d * self.ccfg.t))
392 self.ccfg.has_grad_shard = bool(self._get_cfg_attr(accel,
393 "gradient_accumulation_shard",
394 False))
395 self.ccfg.os_max_shard = (
396 self.ccfg.op_weight_shard if self.ccfg.op_weight_shard >= 1
397 else self.ccfg.d * self.ccfg.t
398 )
400 def _parse_batch(self):
401 """Extract batch settings from ``train`` section."""
402 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
403 self.ccfg.b = max(1, int(self._get_cfg_attr(train_raw, "micro_batch_size", 1) or 1))
404 m = int(self._get_cfg_attr(train_raw, "micro_batch_num", 0) or 0)
405 if m > 0:
406 self.ccfg.m = m
407 else:
408 self.ccfg.m = self.ccfg.p
409 gbs = int(self._get_cfg_attr(train_raw, "global_batch_size", 0) or 0)
410 if gbs > 0:
411 self.ccfg.gbs = gbs
412 else:
413 self.ccfg.gbs = self.ccfg.b * self.ccfg.d * self.ccfg.m
415 def _parse_feature_flags(self):
416 """Set training feature flags."""
417 self.ccfg.has_fa = True
418 self.ccfg.vocab_emb_dp = True
419 self.ccfg.tie_emb_out = False
420 self.ccfg.freeze = False
421 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
422 optimizer = self._get_cfg_attr(train_raw, "optimizer", Config({}))
423 max_grad_norm = float(
424 self._get_cfg_attr(optimizer, "max_grad_norm", 0.0) or 0.0
425 )
426 self.ccfg.has_clip = max_grad_norm > 0
427 self.ccfg.vp_less_mem = False
428 accel = self._get_cfg_attr(train_raw, "accelerator", Config({}))
429 cp_algo = self._get_cfg_attr(accel, "context_parallel_algo", None)
430 if cp_algo:
431 self.ccfg.cp_algo = cp_algo
432 else:
433 self.ccfg.cp_algo = "colossalai_cp"
434 if self.ccfg.cp and self.ccfg.cp > 1:
435 logger.warning(
436 "context_parallel_algo not set; defaulting to "
437 "'colossalai_cp' (Ring CP). Set "
438 "train.accelerator.context_parallel_algo explicitly "
439 "to 'ulysses_cp' if Ulysses CP is intended."
440 )
441 # Optimizer type — used by GlobalConfig.max_op to detect muon-based
442 # optimizers. Matches the MF parser's
443 # ``self.ccfg.optimizer = self.config.optimizer.type``.
444 opt_type = self._get_cfg_attr(optimizer, "type", None)
445 if opt_type:
446 self.ccfg.optimizer = str(opt_type)
448 def _parse_recompute(self):
449 """Parse recompute mode.
451 Reads ``activation_checkpoint`` from ``train.gradient_checkpointing``
452 first. When ``config_overrides`` supplies ``full_rec`` or ``sel_rec``
453 (Matching the MF parser's ``recompute_config.recompute`` /
454 ``recompute_config.select_recompute`` fields), those values take
455 precedence so that Hyper YAML demo files can express per-stage
456 recompute lists for side-by-side comparisons with MindFormers.
457 """
458 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
459 overrides = self._get_cfg_attr(model_raw, "config_overrides", Config({}))
460 full_rec_override = self._get_cfg_attr(overrides, "full_rec", None)
461 sel_rec_override = self._get_cfg_attr(overrides, "sel_rec", None)
463 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
464 gc = self._get_cfg_attr(train_raw, "gradient_checkpointing", Config({}))
465 ac_mode = str(self._get_cfg_attr(gc, "activation_checkpoint", "none"))
467 if full_rec_override is not None:
468 self.ccfg.full_rec = full_rec_override
469 else:
470 self.ccfg.full_rec = ac_mode == "full"
472 if sel_rec_override is not None:
473 self.ccfg.sel_rec = sel_rec_override
474 else:
475 self.ccfg.sel_rec = ac_mode == "selective"
477 self.ccfg.rec_op = Config({
478 "attBMM": 1,
479 "headCast": 1,
480 "dropout": 1,
481 "softmax": 1,
482 "normOp": 1,
483 "gather": 1,
484 "ffAct": 1,
485 })
487 def _init_bytes(self):
488 """Set FP byte sizes from dtype fields in the model section."""
489 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
490 self.ccfg.bytes_p = self._bytes_from_dtype(
491 self._get_cfg_attr(model_raw, "param_init_type", "float32"))
492 self.ccfg.bytes_compute = self._bytes_from_dtype(
493 self._get_cfg_attr(model_raw, "compute_dtype", "bfloat16"))
494 self.ccfg.bytes_softmax = self._bytes_from_dtype(
495 self._get_cfg_attr(model_raw, "softmax_compute_type", "float32"))
496 self.ccfg.bytes_grad = 4
497 self.ccfg.bytes_os = 4
498 self.ccfg.bytes_norm = 4
500 def _init_moe_strategy(self):
501 """Initialize MoE strategy variables via base helper.
503 For MoE models (``n_exp > 1``), ``etp`` defaults to 1 when
504 absent from the YAML, matching the MF parser's
505 ``expert_model_parallel`` default. For dense models the
506 existing ``etp=0`` path continues to produce ``t_exp = t,
507 d_exp = d``.
509 Catches invalid MoE combinations (e.g., ``d_exp = 0`` when
510 ``dp < ep``) so the search engine can proceed — invalid combos
511 will later be filtered by the memory budget check.
512 """
513 if self.ccfg.n_exp > 1 and self.ccfg.etp == 0:
514 self.ccfg.etp = 1
515 try:
516 self.config_dp_tp_exp(self.ccfg)
517 except TypeError:
518 logger.warning(
519 "MoE config_dp_tp_exp failed for d=%d t=%d ep=%d etp=%d "
520 "n_exp=%d — clamping to minimum values.",
521 self.ccfg.d, self.ccfg.t, self.ccfg.ep,
522 self.ccfg.etp, self.ccfg.n_exp,
523 )
524 self.ccfg.d_exp = max(1, self.ccfg.d_exp)
525 self.ccfg.t_exp = max(1, self.ccfg.t_exp)
526 self.ccfg.hff_exp = max(1, self.ccfg.hff_exp)
527 self.ccfg.n_exp = max(1, self.ccfg.n_exp)
529 def _init_offset(self):
530 """Initialize the pipeline offset.
532 The MF parser reads ``model.model_config.offset`` directly from the
533 YAML. When it is a list (e.g. ``[1, 1, ..., -1]``),
534 ``CostModelConfig.is_consistent_pp_config`` requires
535 ``len(offset) == pp``, so strategies whose pipeline degree differs
536 are rejected until ``GlobalConfig.adapt_config`` regenerates a
537 matching offset. A scalar ``0`` is always accepted.
539 To match the MF parser's *list*-based filtering behaviour (used by
540 DeepSeek-V3 and other models that declare an explicit offset), this
541 parser emits a list offset of length ``pp`` (all zeros = even
542 balancing) by default. An explicit offset supplied via
543 ``config_overrides.offset`` overrides this — a list is used as-is,
544 and a non-zero int is broadcast to ``[int] * pp``.
545 """
546 model_raw = self._get_cfg_attr(self.config, "model", Config({}))
547 overrides = self._get_cfg_attr(model_raw, "config_overrides", Config({}))
548 explicit = self._get_cfg_attr(overrides, "offset", None)
549 if isinstance(explicit, list):
550 self.ccfg.offset = list(explicit)
551 elif isinstance(explicit, int):
552 if explicit == 0:
553 self.ccfg.offset = 0
554 else:
555 self.ccfg.offset = [explicit] * self.ccfg.p
556 else:
557 self.ccfg.offset = [0] * self.ccfg.p
559 def config_shard_emb(self) -> None:
560 """Configure embedding sharding based on current parallelism.
562 Mirrors ``CostModelParserMindformers.config_shard_emb`` so that
563 ``set_strategy`` recomputes ``shard_embed`` whenever the parallel
564 configuration changes. When ``vocab_emb_dp`` is enabled and pipeline
565 parallelism is disabled (``p == 1``), the embedding is sharded only
566 along the data-parallel dimension (``d``); otherwise it is sharded
567 along ``t * d``.
569 Without this method, ``CostModelConfig.set_strategy`` skips the
570 ``config_shard_emb`` call (guarded by ``hasattr``) and the initial
571 ``shard_embed`` value computed in ``_init_shard`` is never refreshed,
572 producing an embedding-memory mismatch versus the MF parser.
573 """
574 self.ccfg.shard_embed = (
575 self.ccfg.d
576 if (self.ccfg.vocab_emb_dp and self.ccfg.p == 1)
577 else (self.ccfg.t * self.ccfg.d)
578 )
580 def config_shard_recompute(self) -> None:
581 """Recompute ``shard_recompute_input`` after strategy changes.
583 When ``recompute_slice_activation`` is ``True``, the recompute input
584 is sharded by the current tensor-parallel degree ``t``; otherwise it
585 is not sharded (value ``1``). This method is called by
586 ``set_strategy`` (via ``hasattr`` guard) so that changing ``t``
587 during search correctly updates the sharding factor.
589 Without this method, ``shard_recompute_input`` retains the value
590 computed at initial parse time (using the default ``t`` from the
591 YAML), causing memory-estimation errors when the search explores
592 strategies with different ``t`` values.
593 """
594 self.ccfg.shard_recompute_input = (
595 self.ccfg.t if self._recompute_slice_activation else 1
596 )
598 def _init_shard(self):
599 """Initialize sharding variables.
601 ``shard_embed`` is computed via :meth:`config_shard_emb` so the
602 initial value follows the same rule used on subsequent
603 ``set_strategy`` calls. ``shard_output_activ`` defaults to 1 (no
604 sharding), matching the MF parser's default; the ``custom_qwen``
605 arch hook overrides it to ``ccfg.t`` for Qwen-family models via
606 ``check_and_apply_custom_hook``.
608 ``shard_recompute_input`` mirrors the MF parser's
609 ``recompute_config.recompute_slice_activation`` flag: when the flag
610 is ``True`` (DeepSeek-V3), activations are sharded by ``ccfg.t``;
611 when ``False`` (Qwen), they are not sharded. The flag is stored
612 as ``self._recompute_slice_activation`` so that
613 :meth:`config_shard_recompute` can recompute the value after
614 ``set_strategy`` changes ``t``. Per-model arch hooks
615 (e.g. ``custom_qwen``) may override this during ``EvaluatorV2``
616 initialisation.
617 """
618 self.config_shard_emb()
619 self.ccfg.shard_output_activ = 1
620 train_raw = self._get_cfg_attr(self.config, "train", Config({}))
621 gc = self._get_cfg_attr(train_raw, "gradient_checkpointing", Config({}))
622 self._recompute_slice_activation = bool(self._get_cfg_attr(
623 gc, "recompute_slice_activation", False
624 ))
625 self.config_shard_recompute()
626 self.ccfg.is_shard_mtp_param = True