Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / compute_memory.py: 86%
264 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""Derive per-layer memory parameters from a set of dry-run stage observations."""
16import numpy as np
18import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute
19from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import Layer
20from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger
21from hyper_parallel.auto_parallel.sapp_ppb.utils.stage import Stage, filter_stage_id
24class ComputeMemory:
25 """
26 ComputeMemory class to compute the different memories with stages information running (dry) log
28 stage{A|B} means stage with different configuration A and B
29 stage{1|2} means stage same configuration but different id (can be id other than 1 or 2)
31 number_of_stage_ (int): number of stages for the LLM
32 stagesA_ (list[Stage]): list of dry run stages information, with all the same configuration A,
33 required at least staged 0, 1, (n-2), (n-1)
34 Don't set directly stagesA_, but use set_stagesA
35 stagesB_ (list[Stage]): list of dry run stages information, with all the same configuration B,
36 different from config A required at least staged 0, 1, (n-2), (n-1)
37 Don't set directly stagesB_, but use set_stagesB
38 memory_parameter_ (float): memory_parameter_ of the BODY layer, memory required to run the layer
39 memory_activation_rec_ (dict[Recompute.TYPE, float]) activation memory per recompute types
40 recompute_considered_ (dict[Recompute.TYPE, bool]) recomputation types taken into consideration
41 memory_const_ (float): constant memory required for each stages
42 memory_head_ (float): memory required to run the head layer
43 memory_tail_ (float): memory required to run the tail layer
44 """
46 number_of_stage_: int
47 stages_a: list[Stage]
48 stages_b: list[Stage]
49 memory_parameter_: float
50 memory_activation_rec_: dict[Recompute.TYPE, float]
51 recompute_considered_: dict[Recompute.TYPE, bool]
52 memory_const_: float
53 memory_head_: float
54 memory_tail_: float
56 def __init__(self, number_of_stage: int, stages_a: list[Stage] = None,
57 stages_b: list[Stage] = None) -> None:
58 """Build a :class:`ComputeMemory` solver instance.
60 Args:
61 number_of_stage: Total number of pipeline stages in the target LLM.
62 stages_a: Dry-run observations with configuration A (at least stages ``0, i, j, n-1``).
63 stages_b: Dry-run observations with configuration B (must differ from A).
64 """
65 self.number_of_stage_ = number_of_stage
66 self.set_stages_a(stages_a)
67 self.set_stages_b(stages_b)
68 # number_of_stage != len(stages) can be true
69 self.memory_parameter_ = None
70 self.memory_activation_rec_ = {r: None for r in Recompute.TYPE}
71 self.find_recompute_considered()
72 self.memory_const_ = None
73 self.memory_head_ = None
74 self.memory_tail_ = None
76 def set_stages_a(self, stages: list[Stage]) -> None:
77 """Assign dry-run observations to configuration A after a consistency check."""
78 if stages is None:
79 self.stages_a = []
80 return
81 for stage1 in stages:
82 for stage2 in stages:
83 if not stage1.same_global_config(stage2):
84 logger.error(
85 "Cannot set stagesA, all elements don't have the same configuration",)
86 self.stages_a = []
87 return
88 self.stages_a = stages
90 def set_stages_b(self, stages: list[Stage]) -> None:
91 """Assign dry-run observations to configuration B (must differ from A)."""
92 if stages is None:
93 self.stages_b = []
94 return
95 for stage1 in stages:
96 for stage2 in stages:
97 if not stage1.same_global_config(stage2):
98 logger.error(
99 "Cannot set stagesB, all elements don't have the same configuration")
100 self.stages_b = []
101 return
102 for stage_a in self.stages_b:
103 if stage1.same_global_config(stage_a):
104 logger.error(
105 "Cannot set stagesB, an elements have the same configuration than stagesA")
106 self.stages_b = []
107 return
108 self.stages_b = stages
110 def find_recompute_considered(self) -> None:
111 """Populate :attr:`recompute_considered_` from the observed ``stages_a`` data."""
112 self.recompute_considered_ = {r: False for r in Recompute.TYPE}
113 self.recompute_considered_[Recompute.TYPE.NONE] = True
115 for stage in self.stages_a:
116 for rec in Recompute.TYPE:
117 if stage.nb_layer_rec_[rec] > 0:
118 self.recompute_considered_[rec] = True
120 def _compute_memory_parameter_local_(self, stage1: Stage, stage2: Stage) -> float:
121 """
122 Given 2 stages information with the same configuration, and different id,
123 Compute the memory_parameter
124 """
125 if stage1.same_config(stage2):
126 if stage1.id_ != stage2.id_:
127 res = stage1.memory_usage_ * (stage1.nb_stage_ - stage1.id_)
128 res -= stage2.memory_usage_ * (stage2.nb_stage_ - stage2.id_)
129 res /= stage1.id_ - stage2.id_
130 res = abs(res)
131 res /= stage1.nb_layer_
132 return res
133 logger.error(
134 "stage with same characteristic, BUT SAME ID too, cannot compute memory_parameter")
135 return 0
136 logger.error("stage with different characteristic, cannot compute memory_parameter")
137 return 0
139 def _compute_memory_parameter_(self, multi_run=False) -> float:
140 """Compute memory_parameter
141 With all available stages compute all combinations of memory parameter
142 and return the mean of all the memory_parameter found
143 BEWARE: can update memory_parameter_ & memory_activation_rec_
144 because of _compute_memories_layers_()
145 return: memory_parameter
146 """
147 if multi_run or (len(self.stages_a) < 5 and len(self.stages_b) < 5):
148 memory_parameter_list = []
149 for stage1 in self.stages_a:
150 if stage1.id_ not in [0, (self.number_of_stage_ - 1)]:
151 mem_param = self._compute_memory_parameter_local_(stage1, stage2)
152 for stage2 in self.stages_a:
153 if (stage2.id_ not in [0, (self.number_of_stage_ - 1),
154 stage1.id_] and mem_param != 0):
155 memory_parameter_list.append(mem_param)
156 for stage1 in self.stages_b:
157 if stage1.id_ not in [0, (self.number_of_stage_ - 1)]:
158 for stage2 in self.stages_b:
159 mem_param = self._compute_memory_parameter_local_(stage1, stage2)
160 if (stage2.id_ not in [0, (self.number_of_stage_ - 1),
161 stage1.id_] and mem_param != 0):
162 memory_parameter_list.append(mem_param)
163 return np.mean(memory_parameter_list)
164 if self._compute_memories_layers_():
165 return self.memory_parameter_
166 logger.error("Issue with _compute_memory_parameter_!!!")
167 return 0
169 def _compute_memory_activation_(self, rec, multi_run=False) -> float:
170 """
171 Compute memory_activation for a given recomputation type
172 return: memory_activation
173 """
174 if multi_run or (len(self.stages_a) < 5 and len(self.stages_b) < 5):
175 # look at solution 4 stages
176 logger.error("Not implemented yet!!!")
177 return 0
178 if self._compute_memories_layers_():
179 return self.memory_activation_rec_[rec]
180 logger.error("Issue with _compute_memory_activation_!!!")
181 return 0
183 def zero_offset(self) -> bool:
184 """Return ``True`` if every stage in ``stages_a`` hosts the same number of layers."""
185 nb_layer = self.stages_a[0].nb_layer_
186 for s in self.stages_a:
187 if s.nb_layer_ != nb_layer:
188 return False
189 return True
191 def _compute_memories_layers_(self) -> bool:
192 """check if enough stage number is provided"""
193 used_rec = Recompute.get_used_list(self.recompute_considered_)
194 used_rec_num = len(used_rec)
195 stage_num = len(self.stages_a)
196 if stage_num == used_rec_num + 3:
197 return self._compute_memories_layer_bodies_(False)
198 if stage_num >= used_rec_num + 4:
199 logger.info("Enabled const memory component because enough stages were given")
200 if self.zero_offset():
201 logger.error(
202 "The number of layer per stage cannot be the same for all stages "
203 "when const component is enabled. Some offset must be used"
204 )
205 return False
206 return self._compute_memories_layer_bodies_(True)
208 logger.error(
209 "%s stages found and (%s) recomputation considered"
210 "is not coherent. There should be 3 or 4 more stages than recomputation considered",
211 stage_num,
212 used_rec_num,
213 )
214 return False
216 def _compute_memories_layer_bodies_local_(
217 self, unused_rec: list[Recompute.TYPE],
218 stages: list[Stage]) -> tuple[float, float, float]:
219 """Compute memory_parameter & memory activation for all recomputation types
220 Require at least 3 Stages different from first and last stage
221 """
222 variable_factor_list = []
223 constant_memory_list = []
224 unused_rec.sort(reverse=True)
225 for stage in stages:
226 if stage.id_ not in [0, self.number_of_stage_ - 1]:
227 variable_factor_list.append(stage.get_index_memory_var())
228 for rec_i in unused_rec:
229 variable_factor_list[-1].pop(1 + rec_i)
230 constant_memory_list.append(stage.memory_usage_)
231 solution = list(
232 np.linalg.solve(np.array(variable_factor_list),
233 np.array(constant_memory_list)))
234 memory_param = solution.pop(0)
235 memory_act_rec = Recompute.assign_used(solution, unused_rec)
236 return (memory_param, memory_act_rec)
240 def _compute_memories_layer_bodies_local_with_fix_(
241 self, unused_rec: list[Recompute.TYPE],
242 stages: list[Stage]) -> tuple[float, float, float]:
243 """Compute memory_const, memory_parameter & memory activation for all recomputation types
244 Require at least 4 Stages different from first and last stage
245 """
246 variable_factor_list = []
247 constant_memory_list = []
248 unused_rec.sort(reverse=True)
249 for stage in stages:
250 if stage.id_ not in [0, self.number_of_stage_ - 1]:
251 variable_factor_list.append([1] + stage.get_index_memory_var())
252 for rec_i in unused_rec:
253 variable_factor_list[-1].pop(2 + rec_i)
254 constant_memory_list.append(stage.memory_usage_)
255 logger.debug(
256 "solve(\n %s, \n %s) ",
257 np.array(variable_factor_list),
258 np.array(constant_memory_list),
259 )
260 used_rec = Recompute.get_used_list(self.recompute_considered_)
261 used_rec_num = len(used_rec)
263 if len(stages) < used_rec_num + 4:
264 raise ValueError("Stages given are not enough to solve memory constraints")
265 if len(stages) == used_rec_num + 4:
266 solution = list(
267 np.linalg.solve(np.array(variable_factor_list),
268 np.array(constant_memory_list)))
269 else:
270 logger.warning("Stages given are more than needed, switch to least sqaures method")
271 solution = list(np.linalg.lstsq(np.array(variable_factor_list),
272 np.array(constant_memory_list), rcond=None)[0])
274 memory_const = solution.pop(0)
275 memory_param = solution.pop(0)
276 memory_act_rec = Recompute.assign_used(solution, unused_rec)
277 return (memory_const, memory_param, memory_act_rec)
279 def _compute_memories_layer_bodies_(self, with_fix: bool) -> bool:
280 """
281 Compute memory_parameter, memory_recompute, memory_activation
282 Require at least 3 Stages different from first and last stage
283 BEWARE: can update memory_parameter_, memory_recompute_, memory_activation_
284 return True if success to update memory_parameter_, memory_recompute_, memory_activation_
285 """
287 memory_const_a = None
288 memory_parameter_a = None
289 memory_recompute_a = {r: None for r in Recompute.TYPE}
291 memory_const_b = None
292 memory_parameter_b = None
293 memory_recompute_b = {r: None for r in Recompute.TYPE}
295 unused_rec = Recompute.get_unused_list(self.recompute_considered_)
296 logger.info("unused recomputation: %s", unused_rec)
298 if with_fix:
299 if len(self.stages_a) >= 5:
300 (memory_const_a,
301 memory_parameter_a,
302 memory_recompute_a) = (self._compute_memories_layer_bodies_local_with_fix_(
303 unused_rec, self.stages_a))
304 if len(self.stages_b) >= 5:
305 (memory_const_b,
306 memory_parameter_b,
307 memory_recompute_b) = (self._compute_memories_layer_bodies_local_with_fix_(
308 unused_rec, self.stages_b))
310 return self._average_if_needed_fix(
311 memory_const_a,
312 memory_parameter_a,
313 memory_recompute_a,
314 memory_const_b,
315 memory_parameter_b,
316 memory_recompute_b,
317 )
318 if len(self.stages_a) >= 5:
319 (memory_parameter_a,
320 memory_recompute_a) = (self._compute_memories_layer_bodies_local_(
321 unused_rec, self.stages_a))
322 if len(self.stages_b) >= 5:
323 (memory_parameter_b,
324 memory_recompute_b) = (self._compute_memories_layer_bodies_local_(
325 unused_rec, self.stages_b))
327 return self._average_if_needed(
328 memory_parameter_a,
329 memory_recompute_a,
330 memory_parameter_b,
331 memory_recompute_b,
332 )
334 def _average_if_needed_fix(
335 self,
336 memory_const_a,
337 memory_parameter_a,
338 memory_recompute_a,
339 memory_const_b,
340 memory_parameter_b,
341 memory_recompute_b,
342 ):
343 """check if average is needed"""
344 if memory_parameter_a is not None and memory_parameter_a != 0:
345 if memory_parameter_b is not None and memory_parameter_b != 0:
346 self.memory_const_ = (memory_const_a +
347 memory_const_b) / 2
348 self.memory_parameter_ = (memory_parameter_a +
349 memory_parameter_b) / 2
350 Recompute.average([memory_recompute_a, memory_recompute_b])
351 else:
352 self.memory_const_ = memory_const_a
353 self.memory_parameter_ = memory_parameter_a
354 self.memory_activation_rec_ = memory_recompute_a
356 elif memory_parameter_b is not None and memory_parameter_b != 0:
357 self.memory_const_ = memory_const_b
358 self.memory_parameter_ = memory_parameter_b
359 self.memory_activation_rec_ = memory_recompute_b
360 else:
361 logger.error("failed to compute memories")
362 return False
363 return True
365 def _average_if_needed(self, memory_parameter_a, memory_recompute_a, memory_parameter_b,
366 memory_recompute_b,):
367 """check if average is needed"""
368 if memory_parameter_a is not None and memory_parameter_a != 0:
369 if memory_parameter_b is not None and memory_parameter_b != 0:
370 self.memory_parameter_ = (memory_parameter_a + memory_parameter_b) / 2
371 Recompute.average([memory_recompute_a, memory_recompute_b])
372 else:
373 self.memory_parameter_ = memory_parameter_a
374 self.memory_activation_rec_ = memory_recompute_a
376 elif memory_parameter_b is not None and memory_parameter_b != 0:
377 self.memory_parameter_ = memory_parameter_b
378 self.memory_activation_rec_ = memory_recompute_b
379 else:
380 logger.error("failed to compute memories")
381 return False
382 return True
384 def _compute_memory_head_(self) -> float:
385 """compute the memory for the head"""
386 head_stages = filter_stage_id(self.stages_a, 0)
387 head_stages += filter_stage_id(self.stages_b, 0)
388 memory_head_list = []
389 mem_parameter = self.get_memory_parameter()
390 for head in head_stages:
391 head_memory = head.memory_usage_
392 for rec in Recompute.TYPE:
393 if self.recompute_considered_[rec] is True:
394 head_memory -= (head.nb_layer_rec_[rec] * self.get_memory_activation(
395 rec) * self.number_of_stage_)
396 head_memory -= (head.nb_layer_) * mem_parameter
397 memory_head_list.append(head_memory)
398 return np.mean(memory_head_list)
400 def _compute_memory_tail_(self) -> float:
401 """compute the memory for the tail"""
402 tail_stages = filter_stage_id(self.stages_a, self.number_of_stage_ - 1)
403 tail_stages += filter_stage_id(self.stages_b, self.number_of_stage_ - 1)
404 memory_tail_list = []
405 for tail in tail_stages:
406 tail_memory = tail.memory_usage_
407 for rec in Recompute.TYPE:
408 if self.recompute_considered_[rec] is True:
409 tail_memory -= (tail.nb_layer_rec_[rec] * self.get_memory_activation(rec) * 1)
410 tail_memory -= (tail.nb_layer_) * self.get_memory_parameter()
411 memory_tail_list.append(tail_memory)
412 return np.mean(memory_tail_list)
414 def get_memory_const(self) -> float:
415 """Return the solver-derived constant memory component per stage."""
416 return self.memory_const_
418 def get_memory_parameter(self, force_recompute: bool = False) -> float:
419 """Return the per-body-layer parameter memory, recomputing on demand."""
420 if force_recompute or self.memory_parameter_ is None:
421 self.memory_parameter_ = self._compute_memory_parameter_()
422 return self.memory_parameter_
424 def get_memory_activation(self, rec: Recompute.TYPE,
425 force_recompute: bool = False) -> float:
426 """Return the per-layer activation memory for a given recomputation type."""
427 if force_recompute or self.memory_activation_rec_[rec] is None:
428 self.memory_activation_rec_[rec] = self._compute_memory_activation_(rec)
429 return self.memory_activation_rec_[rec]
431 def get_memory_head(self, force_recompute: bool = False) -> float:
432 """Return the HEAD-layer memory, recomputing on demand."""
433 if force_recompute or self.memory_head_ is None:
434 self.memory_head_ = self._compute_memory_head_()
435 return self.memory_head_
437 def get_memory_tail(self, force_recompute: bool = False) -> float:
438 """Return the TAIL-layer memory, recomputing on demand."""
439 if force_recompute or self.memory_tail_ is None:
440 self.memory_tail_ = self._compute_memory_tail_()
441 return self.memory_tail_
444def compute_memories(layers: list[Layer], memory_folder: str, number_of_stage: int) -> list[Layer]:
445 """compute memories"""
446 filename = ""
447 # Put some meta information in a predefine .json file like layers info?
448 with open(memory_folder + filename, encoding="utf-8"):
449 pass
450 cm = ComputeMemory(number_of_stage=number_of_stage, stages_a=[], stages_b=[])
452 for layer in layers:
453 if layer.type_ == Layer.type_enum.HEAD:
454 layer.memory_parameter_ = cm.get_memory_head()
455 elif layer.type_ == Layer.type_enum.TAIL:
456 layer.memory_parameter_ = cm.get_memory_tail()
457 elif layer.type_ == Layer.type_enum.BODY:
458 layer.memory_parameter_ = cm.get_memory_parameter()
459 for rec in Recompute.TYPE:
460 layer.memory_activation_rec_[rec] = cm.get_memory_activation(rec)
461 return layers