« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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"""cost model parser module"""
16from __future__ import annotations
17from typing import TYPE_CHECKING, Any, Tuple
19import logging
20from abc import ABC
21from abc import abstractmethod
23from hyper_parallel.auto_parallel._layer_stack import derive_layers, resolve_layers
24from hyper_parallel.auto_parallel._op_profiles import infer_arch, load_op_profile, resolve_ops
25from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
27# What a parse says a run should read: the memory logger is off below -v 4.
28parse_logger = logging.getLogger(__name__)
30# The optimizer a config that states none is priced as, with a warning: the
31# shipped recipes state AdamW and Muon both, so none is assumed silently.
32DEFAULT_OPTIMIZER = "adamw"
34if TYPE_CHECKING:
35 from typing import Mapping, Optional
36 from hyper_parallel.auto_parallel._layer_stack import LayerStack
37 from hyper_parallel.auto_parallel._model_spec import OpCounts
38 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import _CostModVar
41class _CostModelParser(ABC):
42 """abstract parser class"""
44 def __init__(self, ccfg: _CostModVar):
45 self.ccfg = ccfg
46 self.config = ccfg.config
48 @abstractmethod
49 def parse(self):
50 """Parse the cost model configuration and populate the cost model variables.
52 Subclasses must implement this method to read framework-specific
53 configuration values into the shared _CostModVar instance.
54 """
56 @staticmethod
57 def stated_optimizer(stated: Any) -> Tuple[str, int]:
58 """Name the optimizer a config states, and the states it keeps a parameter, which its name decides.
60 Muon keeps one momentum a matrix and AdamW two moments a parameter.
61 A config that states no optimizer is priced as AdamW, and the parser
62 says so: the HyperParallel parser named one silently and charged its
63 two states, which a Muon run described without its optimizer does
64 not keep, and the MindFormers parser failed on the missing section
65. The name is what the state count is read from.
67 Returns:
68 The optimizer's name and the states it keeps a parameter.
69 """
70 if not stated:
71 parse_logger.warning("the config states no optimizer: it is priced as %s, two states a parameter",
72 DEFAULT_OPTIMIZER)
73 stated = DEFAULT_OPTIMIZER
74 name = str(stated)
75 return name, 1 if "muon" in name.lower() else 2
77 @staticmethod
78 def state_optimizer(ccfg: Any, stated: Any) -> None:
79 """Set the optimizer a config states and its state count on *ccfg* (``stated_optimizer``)."""
80 ccfg.optimizer, ccfg.optimizer_states = _CostModelParser.stated_optimizer(stated)
82 def config_op_counts(
83 self,
84 ccfg: _CostModVar,
85 arch: Optional[str] = None,
86 ops: Optional[Mapping[str, OpCounts]] = None,
87 ) -> None:
88 """Settle the op profile the model is priced with.
90 A parser that reads a declared arch passes it, with any op counts the
91 spec declares; one that has only a model name leaves the family to
92 :func:`infer_arch`.
93 """
94 ccfg.arch = arch or infer_arch(ccfg.model_name)
95 ccfg.op_counts = resolve_ops(ccfg.arch, ops)
97 @staticmethod
98 def config_layer_stack(ccfg: _CostModVar, stack: Optional[LayerStack] = None) -> None:
99 """Settle the layer stack the estimators price.
101 Args:
102 ccfg: The config, with its arch, op counts and layer counts parsed.
103 stack: The stack the model states. By default, the one its layer
104 counts imply: dense layers first, two halves, or one kind.
105 """
106 if stack is None:
107 layers = derive_layers(
108 load_op_profile(ccfg.arch), int(ccfg.n_lay), int(ccfg.n_mtp or 0),
109 first_k_dense=int(ccfg.k_1st_dense or 0),
110 )
111 stack = resolve_layers(ccfg.arch, layers, ccfg.op_counts)
112 ccfg.layer_stack = stack
113 if len(stack.groups) > 1:
114 logger.info(
115 "%s layer stack: %s", ccfg.model_name,
116 ", ".join(f"{group.count}x{group.kind.name}" for group in stack.groups),
117 )
119 # def config_op_level(self, ccfg, strategy):
120 # def full_partial():
121 # return Config({"full":0, "partial":0})
122 # def exp_or_not():
123 # return Config({
124 # "non_exp":full_partial(),
125 # "exp":full_partial()
126 # })
127 # ccfg.op = Config({
128 # "p":exp_or_not(),
129 # "os":exp_or_not(),
130 # "grad"exp_or_not()
131 # })
132 # shard_strat = {
133 # "grad":0, #zero 1
134 # "os+grad":0, #zero 2
135 # "p+os+grad":0, # zero 3
136 # "p+os":0 # zero2 mindspore
137 # }
138 # shard_strat[strategy]
140 def init_hff(self):
141 """MindFormers format for FFn hidden size"""
142 # Assuming following 3 variables are already parsed
143 hidden_size = self.ccfg.h
144 ffn_dim_multiplier = self.ccfg.fdm
145 multiple_of = self.ccfg.multiple_of
146 hff = 4 * hidden_size
147 if ffn_dim_multiplier:
148 hff = int((ffn_dim_multiplier + 0.01) * hff)
149 hff = int(2 * hff / 3)
150 hff = multiple_of * ((hff + multiple_of - 1) // multiple_of)
151 return hff