Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / sapp / sapp_solver.py: 84%
859 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"""Solver Class"""
17import os
18from dataclasses import dataclass
19from enum import IntEnum
20from typing import Any, Dict, List, Optional
22import pulp as lpSolver
24import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute
25from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import Layer
26from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger
28# seqpipe const
29TENSOR_FLOAT_16 = 2
30TENSOR_FLOAT_32 = 4
31const_from_byte_to_mb = 1024 * 1024
32# llama intermideate_size
33LLAMA_INTERMEDIATE_SIZE = 11008
36@dataclass
37class PipelineMemoryConstraint:
38 """constraint struct"""
39 prob: Any
40 variables: Any
41 layers_sorted: dict[Any]
42 num_of_stage: int
43 num_of_interleave: int
44 micro_batch: int
45 memory_limit: int
48class SappSolver:
49 """solver for pipeline balance"""
51 BIG_M = 1000000
53 MEM_OVERHEAD_NAME = "memory_overhead"
54 TOTAL_SUM = "var_sum_FPi_BPi"
55 CHUNKS_SUM = "chunks_sum"
56 PREV_DIFF = "prev_diff"
57 NEXT_DIFF = "next_diff"
58 MAX_STAGE_TIME = "max_stage_time"
59 MAX_LAST_CHUNK = "max_last_chunk"
60 LAYER_FRONTIER = "layer_frontier"
61 REC_FRONTIER = "recompute_frontier"
62 PROP_PHASE = IntEnum("Propagation", ["FW", "BW"], start=0)
64 def __init__(
65 self,
66 num_of_stage: int,
67 num_of_interleave: int,
68 num_of_micro_batch: int,
69 max_memory: int,
70 layers: list[Layer],
71 layers_sorted: dict[Layer.type_enum, list[Layer]],
72 vpp_less_memory: bool = False,
73 # add dualpipe_v arg
74 dual: bool = False,
75 constant_memory: int = 0,
76 optimization_level: int = 1,
77 description: str = "Pipeline_execution_time_minimize",
78 extracted_training_params: dict[str, int] = None,
79 seq_split_num: int = 1,
80 ) -> None:
81 """Build the ILP variables and the empty problem skeleton.
83 Args:
84 num_of_stage: Number of physical pipeline stages.
85 num_of_interleave: Virtual-pipeline (VPP) chunk count.
86 num_of_micro_batch: Number of micro-batches.
87 max_memory: Per-device memory budget (MB).
88 layers: Flat list of :class:`Layer` descriptors covering the full model.
89 layers_sorted: ``layers`` indexed by HEAD / BODY / TAIL classification.
90 vpp_less_memory: Use the less-memory VPP scheduler variant.
91 dual: Enable dualpipe-V scheduling support.
92 constant_memory: Constant per-stage memory overhead (MB).
93 optimization_level: Solver optimization level (``0-2``).
94 description: Problem description used when exporting the LP model.
95 extracted_training_params: Optional training params for sequence-pipeline mode.
96 seq_split_num: Number of sequence splits (``>1`` enables sequence pipeline).
97 """
99 self.num_of_stage_ = num_of_stage
100 self.num_of_interleave_ = num_of_interleave
101 self.num_of_micro_batch_ = num_of_micro_batch
102 self.max_memory_ = max_memory
103 self.vpp_less_memory_ = vpp_less_memory
104 # Add dualpipe_v
105 self.dual_ = dual
106 self.constant_memory_ = constant_memory
107 self.optimization_level_ = optimization_level
108 self.layers_ = layers
109 self.layers_sorted_ = layers_sorted
111 self.recompute_considered_ = self.find_recompute_considered(
112 layers_sorted)
113 self.extracted_training_params_ = extracted_training_params
114 self.seq_split_num_ = seq_split_num
115 self.seq_pipe = self.seq_split_num_ > 1
116 if self.seq_pipe:
117 self._initialize_seq_pipe_layers()
119 self.variables_ = self._create_variables_to_solve_(
120 num_of_stage, num_of_interleave, layers_sorted)
121 self.problem_ = self._create_problem_(description)
123 def _initialize_seq_pipe_layers(self):
124 """Update memory and time metadata for sequence pipeline mode."""
125 self._update_seq_pipe_memory()
126 self.num_of_micro_batch_ *= self.seq_split_num_
127 self._update_seq_pipe_time()
129 def _update_seq_pipe_memory(self):
130 """Update layer memory values for sequence pipeline mode."""
131 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
132 self._update_body_seq_memory(layer)
133 for head in self.layers_sorted_[Layer.type_enum.HEAD]:
134 self._update_head_seq_memory(head)
135 for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
136 self._update_tail_seq_memory(tail)
138 def _update_body_seq_memory(self, layer):
139 """Update body layer memory values for sequence pipeline mode."""
140 if layer.memory_parameter_ is not None:
141 logger.info("Body Layer 1f1b Parameter Memory: %s", layer.memory_parameter_)
142 layer.memory_parameter_ = self.compute_seq_mem_parameter(
143 layer.memory_parameter_, self.extracted_training_params_)
144 logger.info("Body Layer Seq Parameter Memory: %s", layer.memory_parameter_)
145 for rec in Recompute.TYPE:
146 if not self.recompute_considered_[rec]:
147 continue
148 if rec.name == "FULL":
149 self.recompute_considered_[rec] = False
150 layer.memory_activation_rec_[rec] = None
151 logger.error("Seqpipe doesn't support full recomputation, "
152 "recompute_activation is set as None for seqpp")
153 continue
154 logger.info(
155 "Body Layer 1f1b %s activation Memory: %s",
156 rec,
157 layer.memory_activation_rec_[rec],
158 )
159 layer.memory_activation_rec_[rec] = self.compute_seq_mem_activation(
160 layer.memory_activation_rec_[rec],
161 self.extracted_training_params_,
162 self.seq_split_num_
163 )
164 logger.info(
165 "Body Layer seq %s activation Memory: %s",
166 rec,
167 layer.memory_activation_rec_[rec],
168 )
170 def _update_head_seq_memory(self, head):
171 """Update head layer memory values for sequence pipeline mode."""
172 if head.memory_parameter_ is None:
173 return
174 logger.info("Head cost 1f1b: %s", head.memory_parameter_)
175 head.memory_parameter_ = self.compute_seq_mem_head_cost(
176 head.memory_parameter_, self.extracted_training_params_, self.seq_split_num_)
177 logger.info("Head cost Seq: %s", head.memory_parameter_)
179 def _update_tail_seq_memory(self, tail):
180 """Update tail layer memory values for sequence pipeline mode."""
181 if tail.memory_parameter_ is None:
182 return
183 logger.info("Tail cost 1f1b: %s", tail.memory_parameter_)
184 tail.memory_parameter_ = self.compute_seq_mem_tail_cost(
185 tail.memory_parameter_, self.extracted_training_params_, self.seq_split_num_)
186 logger.info("Tail cost seq: %s", tail.memory_parameter_)
188 def _update_seq_pipe_time(self):
189 """Update layer times for sequence pipeline mode."""
190 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
191 self._update_layer_seq_time(layer, "Body")
192 for head in self.layers_sorted_[Layer.type_enum.HEAD]:
193 self._update_layer_seq_time(head, "Head")
194 for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
195 self._update_layer_seq_time(tail, "Tail")
197 def _update_layer_seq_time(self, layer, layer_name):
198 """Scale one layer's time by the sequence split number."""
199 logger.info("%s Layer 1f1b fp time: %s", layer_name, layer.forward_time_)
200 logger.info("%s Layer 1f1b bp time:", layer_name)
201 for key, value in layer.backward_time_rec_.items():
202 logger.output("%s: %s", key, value)
203 layer.time_ = layer.time_ / self.seq_split_num_
204 layer.forward_time_ = layer.forward_time_ / self.seq_split_num_
205 layer.update_internal_time_for_seqpp()
206 logger.info("%s Layer seq fp time: %s", layer_name, layer.forward_time_)
207 logger.info("%s Layer seq bp time:", layer_name)
208 for key, value in layer.backward_time_rec_.items():
209 logger.output("%s: %s", key, value)
211 @staticmethod
212 def compute_forward_in_backward(num_of_stage: int,
213 micro_batch: int) -> list[int]:
214 """Computes the number of forward propagation happening after a backward"""
215 n = num_of_stage - 1
216 factors = []
217 for _ in range(num_of_stage):
218 factors.append(abs(n))
219 n -= 2
220 if micro_batch < 2 * num_of_stage:
221 for i in range(num_of_stage // 2):
222 factors[i] = 0
223 return factors
225 @staticmethod
226 def compute_lm_forward_in_backward(num_of_stage: int) -> list[int]:
227 """Function compute_forward_in_backward in less_memory schedule"""
228 return list(range(num_of_stage))
230 @staticmethod
231 def compute_activation_nums(num_of_stage: int, num_of_interleave: int,
232 micro_batch: int) -> list[list[int]]:
233 """compute the number of activation"""
234 activation_nums = []
236 if num_of_interleave > 1:
237 for i in range(num_of_interleave):
238 activation_nums.append([])
239 for _ in range(num_of_stage):
240 activation_nums[i].append(num_of_stage)
241 for s in range(num_of_stage):
242 activation_nums[0][s] += max(0, num_of_stage - 2 * s - 1)
243 for s in range(num_of_stage):
244 activation_nums[num_of_interleave - 1][s] += min(
245 0, num_of_stage - 2 * s - 1)
246 for i in range(num_of_interleave):
247 for s in range(num_of_stage):
248 activation_nums[i][s] = min(activation_nums[i][s],
249 micro_batch)
250 else:
251 for i in range(num_of_interleave):
252 activation_nums.append([])
253 for s in range(num_of_stage):
254 activation_nums[i].append(num_of_stage - s)
256 return activation_nums
258 @staticmethod
259 def compute_activation_nums_dual(num_of_stage: int, num_of_interleave: int,
260 micro_batch: int) -> list[list[int]]:
261 """compute the number of activation for dualpipe_v"""
262 activation_nums = []
264 for i in range(num_of_interleave):
265 activation_nums.append([])
266 for _ in range(num_of_stage):
267 activation_nums[i].append(0)
268 for s in range(num_of_stage):
269 activation_nums[0][s] += max(0, 2 * num_of_stage - s)
270 for s in range(num_of_stage):
271 activation_nums[num_of_interleave - 1][s] += max(
272 0, s + 1)
273 for i in range(num_of_interleave):
274 for s in range(num_of_stage):
275 activation_nums[i][s] = min(activation_nums[i][s],
276 micro_batch)
278 return activation_nums
280 @staticmethod
281 def compute_less_activation_nums(
282 num_of_stage: int, num_of_interleave: int) -> list[list[int]]:
283 """compute number of less_mem activation"""
284 activation_nums = []
285 if num_of_interleave > 1:
286 for i in range(num_of_interleave):
287 activation_nums.append([])
288 for _ in range(num_of_stage):
289 activation_nums[i].append(num_of_stage)
290 for s in range(num_of_stage):
291 activation_nums[num_of_interleave - 1][s] -= s
292 else:
293 for i in range(num_of_interleave):
294 activation_nums.append([])
295 for s in range(num_of_stage):
296 activation_nums[i].append(num_of_stage - s)
297 return activation_nums
299 #######################################################################
300 ## ##
301 ## SeqPipe ##
302 ## ##
303 #######################################################################
304 @staticmethod
305 def compute_activation_seq_nums(num_of_stage: int, num_of_interleave: int,
306 seq_split_num: int, micro_batch: int, less_memory: False) -> list[list[int]]:
307 """compute the number of activation for seq chunks"""
308 activation_nums = []
309 if less_memory:
310 act_gap = 1
311 else:
312 act_gap = 2
313 if num_of_interleave > 1:
314 for i in range(num_of_interleave):
315 activation_nums.append([])
316 for _ in range(num_of_stage):
317 activation_nums[i].append(num_of_stage)
318 for s in range(num_of_stage):
319 activation_nums[num_of_interleave - 1][s] = seq_split_num
321 loop_index = 1
322 for stage_index in range(num_of_stage - 2, -1, -1):
323 flag_added = False
324 for chunk_index in range(num_of_interleave):
325 condition1 = activation_nums[chunk_index][stage_index + 1] % num_of_stage != 0
326 condition2 = activation_nums[chunk_index][stage_index + 1] // num_of_stage < loop_index
327 if condition1 or condition2:
328 for update in range(stage_index + 1):
329 activation_nums[chunk_index][update] += act_gap
330 flag_added = True
331 break
332 if not flag_added:
333 for update in range(stage_index + 1):
334 activation_nums[0][update] += act_gap
335 loop_index += 1
336 # microbatch
337 for i in range(num_of_interleave):
338 for s in range(num_of_stage):
339 activation_nums[i][s] = min(activation_nums[i][s],
340 micro_batch)
341 else:
342 for i in range(num_of_interleave):
343 activation_nums.append([])
344 for s in range(num_of_stage):
345 activation_nums[i].append(num_of_stage - s + seq_split_num - 1)
347 logger.output("compute_activation_seq_nums: %s", activation_nums)
348 return activation_nums
350 @staticmethod
351 def compute_seq_mem_activation(original_memory_activation: float,
352 extracted_training_params: dict[str, int],
353 seq_split_num: int) -> float:
354 """compute activation memory for seqpipe"""
355 # context parallel? cp?
356 batch_size = extracted_training_params['batch_size']
357 heads = extracted_training_params['num_heads']
358 seq_length = extracted_training_params['seq_length']
359 head_dim = extracted_training_params['head_dim']
360 mp = extracted_training_params['model_parallel']
361 # cp = extracted_training_params['context_parallel']
362 # 2*Kv add
363 # cp?
364 kv_update_mem_byte = 2 * ((TENSOR_FLOAT_16 * batch_size * heads * seq_length * head_dim) / (mp))
365 kv_update_mem = kv_update_mem_byte / const_from_byte_to_mb
366 # Attention Key,Value
367 # cp?
368 key_mem_byte = (TENSOR_FLOAT_16 * batch_size * heads * seq_length * head_dim) / (mp)
369 key_mem = key_mem_byte / const_from_byte_to_mb
370 # cp?
371 value_mem_byte = (TENSOR_FLOAT_16 * batch_size * heads * seq_length * head_dim) / (mp)
372 value_mem = value_mem_byte / const_from_byte_to_mb
374 seq_memory_activation = (original_memory_activation - key_mem - value_mem) / seq_split_num + kv_update_mem
375 return seq_memory_activation
377 @staticmethod
378 def compute_seq_mem_parameter(original_memory_parameter: float, extracted_training_params: dict[str, int]) -> float:
379 """compute layer parameter memory for seqpipe"""
380 # context parallel? cp?
381 batch_size = extracted_training_params['batch_size']
382 heads = extracted_training_params['num_heads']
383 seq_length = extracted_training_params['seq_length']
384 head_dim = extracted_training_params['head_dim']
385 mp = extracted_training_params['model_parallel']
386 # cp = extracted_training_params['context_parallel']
387 # cp?
388 kv_cache_parameter_mem_byte = 4 * (TENSOR_FLOAT_16 * batch_size * heads * seq_length * head_dim / (mp))
389 kv_cache_parameter_mem = kv_cache_parameter_mem_byte / const_from_byte_to_mb
390 seq_memory_parameter = original_memory_parameter + kv_cache_parameter_mem
391 return seq_memory_parameter
393 @staticmethod
394 def compute_seq_mem_head_cost(original_head_cost: float,
395 extracted_training_params: dict[str, int],
396 seq_split_num: int) -> float:
397 """compute head stage extra cost for seqpipe"""
398 batch_size = extracted_training_params['batch_size']
399 seq_length = extracted_training_params['seq_length']
400 hidden_size = extracted_training_params['hidden_size']
401 mp = extracted_training_params['model_parallel']
402 # cp = extracted_training_params['context_parallel']
403 if mp > 1:
404 # comm operator Mem (recv+reduceScatter)
405 # cp?
406 comm_operator_mem_byte = 2 * (TENSOR_FLOAT_16 * batch_size * seq_length * hidden_size / (mp))
407 comm_operator_mem = comm_operator_mem_byte / const_from_byte_to_mb
408 # StridedSliceGrad Operator Mem
409 stridslice_operator_mem_byte = TENSOR_FLOAT_16 * batch_size * seq_length * hidden_size
410 stridslice_operator_mem = stridslice_operator_mem_byte / const_from_byte_to_mb
411 seq_head_cost = original_head_cost - (1 - 1 / seq_split_num) * (comm_operator_mem + stridslice_operator_mem)
412 else:
413 # comm operator Mem (recv)
414 # cp?
415 comm_operator_mem_byte = TENSOR_FLOAT_16 * batch_size * seq_length * hidden_size / (mp)
416 comm_operator_mem = comm_operator_mem_byte / const_from_byte_to_mb
417 # Grad/MatMul // Grad/Mul Operator Mem
418 # cp?
419 mul_operator_mem_byte = 1 * (TENSOR_FLOAT_16 * batch_size * seq_length * LLAMA_INTERMEDIATE_SIZE / (mp))
420 mul_operator_mem = mul_operator_mem_byte / const_from_byte_to_mb
421 seq_head_cost = original_head_cost - (1 - 1 / seq_split_num) * (comm_operator_mem + mul_operator_mem)
422 return seq_head_cost
424 @staticmethod
425 def compute_seq_mem_tail_cost(original_tail_cost: float,
426 extracted_training_params: dict[str, int],
427 seq_split_num: int) -> float:
428 """compute tail stage extra cost for seqpipe"""
429 batch_size = extracted_training_params['batch_size']
430 seq_length = extracted_training_params['seq_length']
431 vocab_size = extracted_training_params['vocab_size']
432 mp = extracted_training_params['model_parallel']
433 # cp = extracted_training_params['context_parallel']
434 # Memory extra introduced by loss op:
435 # cp?
436 loss_operator_mem_byte = TENSOR_FLOAT_32 * batch_size * seq_length * vocab_size / (mp)
437 loss_operator_mem = loss_operator_mem_byte / const_from_byte_to_mb
438 # New tail Cost = Old tail Cost - (3-3/k)M + (k-1)(M/k)
439 seq_tail_cost = original_tail_cost - (3 - 3 / seq_split_num) * loss_operator_mem + (
440 seq_split_num - 1) * (loss_operator_mem / seq_split_num)
441 return seq_tail_cost
443 def add_total_nb_layer_constraint(self, prob: Any, variables: Any,
444 sorted_layers: Dict[Layer.type_enum, list[Layer]]) -> Any:
445 """Enforce that the sum of assigned layers equals ``layer.nb_layer_`` per BODY layer."""
446 for layer in sorted_layers[Layer.type_enum.BODY]:
447 prob += (lpSolver.lpSum(
448 variables[layer.name_][rec] for rec in Recompute.TYPE
449 if self.recompute_considered_[rec]) == layer.nb_layer_)
450 return prob
452 def add_stage_nb_layer_constraint(self, prob: Any, variables: Any,
453 sorted_layers: Dict[Layer.type_enum, List[Layer]]) -> Any:
454 """Require each non-reserved ``(interleave, stage)`` cell to host at least one layer."""
455 layer_type_num = len(sorted_layers[Layer.type_enum.BODY])
456 reserved_positions = self._reserved_stage_positions()
457 for i in range(self.num_of_interleave_):
458 for s in range(self.num_of_stage_):
459 if (i, s) in reserved_positions:
460 continue
461 terms = []
462 for rec in Recompute.TYPE:
463 if not self.recompute_considered_[rec]:
464 continue
466 for ll in range(layer_type_num):
467 terms.append(
468 variables[
469 sorted_layers[Layer.type_enum.BODY][ll].name_
470 ][rec][i][s]
471 )
473 prob += lpSolver.lpSum(terms) >= 1
474 return prob
476 def _reserved_stage_positions(self):
477 """Return stage positions reserved for head and tail layers."""
478 if self.dual_:
479 return {(0, 0), (1, 0)}
480 return {(0, 0), (self.num_of_interleave_ - 1, self.num_of_stage_ - 1)}
482 def add_multimodal_sequence_constraint(
483 self, prob: Any, variables: Any,
484 sorted_layers: Dict[Layer.type_enum, List[Layer]]) -> Any:
485 """Enforce a stage frontier between successive BODY layer types (multimodal models)."""
486 for frontier in range(1, len(sorted_layers[Layer.type_enum.BODY])):
487 layer = sorted_layers[Layer.type_enum.BODY][frontier].name_
488 for v in range(self.num_of_interleave_):
489 for s in range(self.num_of_stage_):
490 prob = self._add_frontier_lower_bound(prob, variables, layer, frontier, v, s)
491 return self._add_frontier_upper_bounds(prob, variables, sorted_layers)
493 def _add_frontier_lower_bound(self, prob, variables, layer, frontier, interleave, stage):
494 """Add the lower bound for one multimodal frontier variable."""
495 frontier_sum = self._frontier_layer_sum(variables, layer, interleave, stage)
496 if frontier_sum is None:
497 return prob
498 prob += (
499 variables[self.LAYER_FRONTIER][frontier - 1][interleave][stage]
500 >= frontier_sum / self.BIG_M
501 )
502 return prob
504 def _frontier_layer_sum(self, variables, layer, interleave, stage):
505 """Build the layer sum used by multimodal frontier constraints."""
506 if self.dual_:
507 return self._dual_frontier_layer_sum(variables, layer, interleave, stage)
508 return self._current_layer_sum(variables, layer, interleave, range(stage)) + (
509 self._previous_layer_sum(variables, layer, interleave)
510 )
512 def _dual_frontier_layer_sum(self, variables, layer, interleave, stage):
513 """Build the layer sum for dualpipe_v multimodal frontier constraints."""
514 if interleave == 0:
515 return self._current_layer_sum(variables, layer, interleave, range(stage))
516 if interleave == 1:
517 return self._current_layer_sum(variables, layer, interleave, range(stage, self.num_of_stage_)) + (
518 self._previous_layer_sum(variables, layer, interleave)
519 )
520 return None
522 def _current_layer_sum(self, variables, layer, interleave, stage_range):
523 """Sum current interleave variables over a stage range."""
524 terms = []
525 for rec in Recompute.TYPE:
526 if not self.recompute_considered_[rec]:
527 continue
529 for stage in stage_range:
530 terms.append(variables[layer][rec][interleave][stage])
532 return lpSolver.lpSum(terms)
534 def _previous_layer_sum(self, variables, layer, interleave):
535 """Sum variables from previous interleaves."""
536 terms = []
537 for rec in Recompute.TYPE:
538 if self.recompute_considered_[rec]:
539 for prev_interleave in range(interleave):
540 for stage in range(self.num_of_stage_):
541 terms.append(variables[layer][rec][prev_interleave][stage])
542 return lpSolver.lpSum(terms)
544 def _add_frontier_upper_bounds(self, prob, variables, sorted_layers):
545 """Prevent previous body layer types after each multimodal frontier."""
546 for frontier in range(1, len(sorted_layers[Layer.type_enum.BODY])):
547 layer = sorted_layers[Layer.type_enum.BODY][frontier - 1].name_
548 for stage in range(self.num_of_stage_):
549 for interleave in range(self.num_of_interleave_):
550 prob = self._add_frontier_upper_bound(prob, variables, layer, frontier, interleave, stage)
551 return prob
553 def _add_frontier_upper_bound(self, prob, variables, layer, frontier, interleave, stage):
554 """Add one upper bound constraint for a multimodal frontier."""
555 for rec in Recompute.TYPE:
556 if self.recompute_considered_[rec]:
557 prob += variables[layer][rec][interleave][stage] <= (
558 1 - variables[self.LAYER_FRONTIER][frontier - 1][interleave][stage]
559 ) * self.BIG_M
560 return prob
562 def add_multimodal_recompute_constraint(
563 self, prob: Any, variables: Any,
564 sorted_layers: Dict[Layer.type_enum, List[Layer]]) -> Any:
565 """Keep recomputation schemes consistent across BODY layer types (MindFormer constraint)."""
567 # if (self.recompute_considered_[Recompute.TYPE.FULL] and
568 # self.recompute_considered_[Recompute.TYPE.FULL]):
570 considered = Recompute.get_used_list(self.recompute_considered_)
571 if len(considered) > 2:
572 logger.error("Careful: MindFormer does not allow a fine recomputation scheme "
573 "for heterogeneous models. Pipeline balancing is currently unable to "
574 "comply with MF constraint for more than 1 recomputation type.")
575 return prob
577 if len(considered) < 2:
578 # this constraint is unnecessary if there is no recomputation
579 return prob
581 most_rec = max(considered)
582 layer_type_num = len(sorted_layers[Layer.type_enum.BODY])
583 for v in range(self.num_of_interleave_):
584 for s in range(self.num_of_stage_):
585 for rec in Recompute.TYPE:
586 if self.recompute_considered_[rec] and rec is not Recompute.TYPE.NONE:
587 for layer_idx in range(0, layer_type_num - 1):
588 prob += variables[self.REC_FRONTIER][v][s][layer_idx] >= (
589 lpSolver.lpSum(
590 variables[sorted_layers[Layer.type_enum.BODY][next_idx].name_][most_rec][v][s]
591 for next_idx in range(layer_idx + 1, layer_type_num))) / self.BIG_M
593 least_rec = min(considered)
594 for layer_idx in range(0, layer_type_num - 1):
595 layer = sorted_layers[Layer.type_enum.BODY][layer_idx].name_
596 for v in range(0, self.num_of_interleave_):
597 for s in range(0, self.num_of_stage_):
598 prob += variables[layer][least_rec][v][s] <= (
599 1 - variables[self.REC_FRONTIER][v][s][layer_idx]
600 ) * self.BIG_M
601 return prob
603 @staticmethod
604 def find_recompute_considered(
605 layers_sorted: Dict[Layer.type_enum, List[Layer]]) -> Dict[Recompute.TYPE, bool]:
606 """Return the recomputation-considered flags copied from the first BODY layer."""
607 return layers_sorted[Layer.type_enum.BODY][0].recompute_considered_
609 def max_stage_micro_eq_stage(self, prob: Any,
610 layers_sorted: Dict[Layer.type_enum, List[Layer]]) -> Any:
611 """Apply additional VPP optimisations when ``pp == num_of_micro_batch``."""
612 last_chunk = self.num_of_interleave_ - 1
614 for i_stage in range(self.num_of_stage_):
615 for inter in range(last_chunk):
616 prob += self.variables_[self.MAX_STAGE_TIME] >= (
617 self._max_stage_bound_i_bp(layers_sorted, i_stage, inter) +
618 self._max_stage_bound_head_tail(layers_sorted, i_stage,
619 -1, inter))
621 if self.vpp_less_memory_:
622 factors = self.compute_lm_forward_in_backward(self.num_of_stage_)
623 else:
624 factors = self.compute_forward_in_backward(
625 self.num_of_stage_, self.num_of_micro_batch_)
627 for i_stage in range(self.num_of_stage_):
628 logger.debug(
629 "v=%s, s=%s: (BP + HT) + (%s / %s * FP)",
630 last_chunk,
631 i_stage,
632 factors[i_stage],
633 self.num_of_micro_batch_,
634 )
635 prob += self.variables_[self.MAX_LAST_CHUNK] >= (
636 self._max_stage_bound_i_bp(layers_sorted, i_stage, last_chunk) +
637 self._max_stage_bound_head_tail(layers_sorted, i_stage, last_chunk, last_chunk) +
638 (factors[i_stage] / self.num_of_micro_batch_) *
639 self._max_stage_bound_i_fp(layers_sorted, i_stage, last_chunk))
641 if self.optimization_level_ >= 2:
642 logger.debug("Approach 2a")
643 prob += self.variables_[self.MAX_STAGE_TIME] >= (
644 self.variables_[self.MAX_LAST_CHUNK])
646 return self.variables_[self.MAX_STAGE_TIME]
647 logger.debug("Approach 2b")
648 prob += self.variables_[self.MAX_LAST_CHUNK] >= (
649 self.variables_[self.MAX_STAGE_TIME])
651 return (self.variables_[self.MAX_STAGE_TIME] +
652 self.variables_[self.MAX_LAST_CHUNK])
654 def add_performance_constraint(self, prob: Any,
655 layers_sorted: Dict[Layer.type_enum, List[Layer]],
656 pipeline_total_time: Any) -> Any:
657 """Add the ``pipeline_total_time >= …`` performance constraints."""
658 max_stage_time = self.variables_[self.MAX_STAGE_TIME]
659 max_stage_time = self.add_max_stage_constraint(prob, layers_sorted, max_stage_time)
661 total_sum = self.variables_[self.TOTAL_SUM]
662 prob += total_sum >= self._total_sum(layers_sorted)
664 if self.optimization_level_ >= 2:
665 # approach A
666 for v in range(self.num_of_interleave_ - 1):
667 prob += self.variables_[self.PREV_DIFF][v] >= (
668 self._prev_diff_sum(layers_sorted, prob, v))
670 prob += self.variables_[self.CHUNKS_SUM][v] >= (
671 (self.num_of_interleave_ - v) / self.num_of_interleave_ *
672 self._chunks_sum(layers_sorted, v))
674 chunks_sum = lpSolver.lpSum(self.variables_[self.CHUNKS_SUM])
675 prev_diff = lpSolver.lpSum(self.variables_[self.PREV_DIFF])
677 next_diff = self.variables_[self.NEXT_DIFF]
678 prob += next_diff >= (
679 self._next_diff_sum(layers_sorted, prob))
681 prob += pipeline_total_time >= (
682 (total_sum + chunks_sum + prev_diff + next_diff)
683 / max(1, (self.num_of_interleave_ - 2))
684 + max_stage_time * (self.num_of_micro_batch_ - 2)
685 )
686 else:
687 # approach B
688 prob += pipeline_total_time >= max_stage_time
689 return prob
691 def add_max_stage_constraint(self, prob: Any,
692 layers_sorted: Dict[Layer.type_enum, List[Layer]],
693 max_stage_time: Any) -> Any:
694 """Add the ``max_stage_time`` lower-bound constraints over every ``(interleave, stage)``."""
695 if (self.num_of_interleave_ > 1 and self.optimization_level_ >= 1
696 and self.num_of_micro_batch_ == self.num_of_stage_):
697 max_stage_time = self.max_stage_micro_eq_stage(prob, layers_sorted)
698 else:
699 # Constraints on sub-main-part of a stage that it may take (for all stage)
700 for i_stage in range(self.num_of_stage_):
701 for inter_f in range(self.num_of_interleave_):
702 for inter_b in range(self.num_of_interleave_):
703 prob += max_stage_time >= (
704 self._max_stage_bound_i_fp(layers_sorted, i_stage, inter_f) +
705 self._max_stage_bound_i_bp(layers_sorted, i_stage, inter_b) +
706 self._max_stage_bound_head_tail(layers_sorted, i_stage,
707 inter_f, inter_b))
709 return max_stage_time
711 ############################################
712 # Memory Constraint #
713 ############################################
714 def stage_param_memory(self, variables: Any,
715 layers_sorted: Dict[Layer.type_enum, List[Layer]],
716 stage_id: int, num_of_stage: int,
717 num_of_interleave: int) -> Any:
718 """Return an LP expression for the parameter memory of ``stage_id``."""
719 # Add if dual to decide whether dualpipe_v is used
720 if self.dual_:
721 bound = lpSolver.LpAffineExpression()
722 for inter_id in range(num_of_interleave):
723 for layer in layers_sorted[Layer.type_enum.BODY]:
724 for rec in Recompute.TYPE:
725 if self.recompute_considered_[rec]:
726 bound += (
727 variables[layer.name_][rec][inter_id][stage_id] *
728 layer.memory_parameter_)
729 if stage_id == 0:
730 for head in layers_sorted[Layer.type_enum.HEAD]:
731 bound += head.memory_parameter_
732 for tail in layers_sorted[Layer.type_enum.TAIL]:
733 bound += tail.memory_parameter_
734 else:
735 bound = lpSolver.LpAffineExpression()
736 for inter_id in range(num_of_interleave):
737 for layer in layers_sorted[Layer.type_enum.BODY]:
738 for rec in Recompute.TYPE:
739 if self.recompute_considered_[rec]:
740 bound += (
741 variables[layer.name_][rec][inter_id][stage_id] *
742 layer.memory_parameter_)
743 if stage_id == 0:
744 for head in layers_sorted[Layer.type_enum.HEAD]:
745 bound += head.memory_parameter_
746 for tail in layers_sorted[Layer.type_enum.TAIL]:
747 bound += tail.memory_parameter_
748 if stage_id == num_of_stage - 1:
749 for tail in layers_sorted[Layer.type_enum.TAIL]:
750 bound += tail.memory_parameter_
751 return bound
753 def stage_active_memory_per_micro(
754 self, variables: Any,
755 layers_sorted: Dict[Layer.type_enum, List[Layer]],
756 stage_id: int, inter_id: int) -> Any:
757 """Return an LP expression for the activation memory of ``stage_id`` per micro-batch."""
758 bound = lpSolver.LpAffineExpression()
759 for layer in layers_sorted[Layer.type_enum.BODY]:
760 for rec in Recompute.TYPE:
761 if self.recompute_considered_[rec]:
762 bound += (variables[layer.name_][rec][inter_id][stage_id] *
763 layer.memory_activation_rec_[rec])
764 return bound
766 def stage_active_memory(self, variables: Any,
767 layers_sorted: Dict[Layer.type_enum, List[Layer]],
768 stage_id: int, num_of_interleave: int,
769 activation_nums: List[List[int]]) -> Any:
770 """Return the total activation-memory LP expression for ``stage_id``."""
771 bound = lpSolver.LpAffineExpression()
772 for inter_id in range(num_of_interleave):
773 for layer in layers_sorted[Layer.type_enum.BODY]:
774 for rec in Recompute.TYPE:
775 if self.recompute_considered_[rec]:
776 bound += (
777 variables[layer.name_][rec][inter_id][stage_id] *
778 layer.memory_activation_rec_[rec] *
779 activation_nums[inter_id][stage_id])
780 return bound
782 def init_overhead_variables(self, variables: Any, s: int) -> Any:
783 """Compute the per-stage overhead LP expression used in the VPP memory constraint."""
784 bound = lpSolver.LpAffineExpression()
785 vf = self.num_of_interleave_ - 1
786 vb = self.num_of_interleave_ - 1
787 incr_f = True
788 if self.vpp_less_memory_:
789 for _ in range(self.num_of_interleave_ - 1):
790 if incr_f:
791 vf = (vf + 1) % self.num_of_interleave_
792 factor = abs(self.num_of_stage_ - s)
793 else:
794 vb = vb - 1
795 factor = s
796 incr_f = not incr_f
798 logger.debug("%s * (act(%s,%s) - act(%s,%s)", factor, vf, s, vb, s)
799 bound += factor * (
800 self.stage_active_memory_per_micro(variables, self.layers_sorted_, s, vf)
801 - self.stage_active_memory_per_micro(variables, self.layers_sorted_, s, vb))
802 else:
803 for _ in range(self.num_of_interleave_ - 1):
804 if incr_f:
805 vf = (vf + 1) % self.num_of_interleave_
806 logger.debug(
807 "%s * (act(%s,%s) - act(%s,%s)",
808 self.num_of_stage_ - abs(self.num_of_stage_ - 2 * s - 1),
809 vf,
810 s,
811 vb,
812 s,
813 )
814 bound += (self.num_of_stage_ - abs(self.num_of_stage_ - 2 * s - 1)) * (
815 self.stage_active_memory_per_micro(variables, self.layers_sorted_, s, vf)
816 - self.stage_active_memory_per_micro(variables, self.layers_sorted_, s, vb)
817 )
818 else:
819 vb = vb - 1
820 logger.debug(
821 "%s * (act(%s,%s) - act(%s,%s)",
822 max(self.num_of_stage_ - 2 * s - 1, 0),
823 vf + 1,
824 s,
825 vb + 1,
826 s,
827 )
828 bound += max(self.num_of_stage_ - 2 * s - 1, 0) * (
829 self.stage_active_memory_per_micro(variables,
830 self.layers_sorted_, s, vf + 1)
831 - self.stage_active_memory_per_micro(variables,
832 self.layers_sorted_, s, vb + 1)
833 )
834 logger.debug(
835 "%s * (act(%s,%s) - act(%s,%s)",
836 max(-(self.num_of_stage_ - 2 * s - 1), 0),
837 vf,
838 s,
839 vb,
840 s,
841 )
842 bound += max(-(self.num_of_stage_ - 2 * s - 1), 0) * (
843 self.stage_active_memory_per_micro(variables, self.layers_sorted_, s, vf)
844 - self.stage_active_memory_per_micro(variables, self.layers_sorted_, s, vb)
845 )
846 incr_f = not incr_f
848 return bound
850 def stage_overhead_memory(self, variables: Any, stage_id: int) -> Any:
851 """Return the stage-``stage_id`` memory overhead LP expression."""
852 bound = lpSolver.LpAffineExpression()
853 for v in range(self.num_of_interleave_ - 1):
854 bound += variables[self.MEM_OVERHEAD_NAME][stage_id][v]
855 return bound
857 def add_pipeline_memory_constraint(self,
858 constraint: PipelineMemoryConstraint) -> None:
859 """Add per-stage memory upper-bound constraints to the solver problem."""
860 prob = constraint.prob
861 variables = constraint.variables
862 layers_sorted = constraint.layers_sorted
863 num_of_stage = constraint.num_of_stage
864 num_of_interleave = constraint.num_of_interleave
865 micro_batch = constraint.micro_batch
866 memory_limit = constraint.memory_limit
868 if self.vpp_less_memory_:
869 if self.seq_pipe:
870 activation_nums = self.compute_activation_seq_nums(
871 num_of_stage, num_of_interleave, self.seq_split_num_, micro_batch, True)
872 else:
873 activation_nums = self.compute_less_activation_nums(
874 num_of_stage, num_of_interleave)
875 # Add if dual to decide whether dualpipe_v is used
876 elif self.dual_:
877 activation_nums = self.compute_activation_nums_dual(
878 num_of_stage, num_of_interleave, micro_batch)
880 else:
881 if self.seq_pipe:
882 activation_nums = self.compute_activation_seq_nums(
883 num_of_stage, num_of_interleave, self.seq_split_num_, micro_batch, False)
884 else:
885 activation_nums = self.compute_activation_nums(
886 num_of_stage, num_of_interleave, micro_batch)
887 logger.info("activation nums = %s", activation_nums)
889 if self.num_of_stage_ == self.num_of_micro_batch_:
890 for s in range(num_of_stage):
891 prob += memory_limit >= (
892 self.stage_param_memory(variables, layers_sorted, s,
893 num_of_stage, num_of_interleave) +
894 self.stage_active_memory(variables, layers_sorted, s,
895 num_of_interleave, activation_nums) +
896 self.constant_memory_)
897 else:
898 for s in range(num_of_stage):
899 prob += variables[self.MEM_OVERHEAD_NAME][s] >= (
900 self.init_overhead_variables(variables, s)
901 )
902 prob += memory_limit >= (
903 self.stage_param_memory(
904 variables, layers_sorted, s, num_of_stage, num_of_interleave
905 )
906 + self.stage_active_memory(
907 variables, layers_sorted, s, num_of_interleave, activation_nums
908 )
909 + variables[self.MEM_OVERHEAD_NAME][s]
910 + self.constant_memory_
911 )
913 def get_simulator_memory_activation(self) -> list[float]:
914 """Give the activation memory per stage for simulator."""
916 memory_active = []
917 if self.has_some_memory_info():
918 for inter in range(self.num_of_interleave_):
919 memory_active.append([])
920 for stage in range(self.num_of_stage_):
921 memory_active[inter].append(0)
922 memory_activation = 0
923 for rec in Recompute.TYPE:
924 if not self.recompute_considered_[rec]:
925 continue
927 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
928 var_value = self.variables_.get(layer.name_)[rec][inter][stage].varValue
929 memory_activation += var_value * layer.memory_activation_rec_[rec]
931 memory_active[inter][stage] = memory_activation
932 return memory_active
934 def get_simulator_memory_parameter(self) -> list[float]:
935 """Give the parameter memory per stage for simulator."""
936 memory_param_stage = [0] * self.num_of_stage_
937 if self.has_some_memory_info():
938 for inter in range(self.num_of_interleave_):
939 for stage in range(self.num_of_stage_):
940 memory_param_stage[stage] += self._get_stage_parameter_memory(inter, stage)
942 for head in self.layers_sorted_[Layer.type_enum.HEAD]:
943 if head.memory_parameter_ is not None:
944 memory_param_stage[0] += head.memory_parameter_
945 for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
946 if tail.memory_parameter_ is not None:
947 memory_param_stage[self.num_of_stage_ -
948 1] += tail.memory_parameter_
949 memory_param = [memory_param_stage] * self.num_of_interleave_
950 return memory_param
952 def _get_stage_parameter_memory(self, interleave, stage):
953 """Calculate BODY-layer parameter memory for one pipeline position."""
954 total = 0
955 for rec in Recompute.TYPE:
956 if not self.recompute_considered_[rec]:
957 continue
959 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
960 if layer.memory_parameter_ is not None:
961 var_value = self.variables_.get(layer.name_)[rec][interleave][stage].varValue
962 total += var_value * layer.memory_parameter_
963 return total
965 def get_simulator_time(self) -> list[float]:
966 """Give the time per stage for simulator."""
967 time = []
968 for i in range(self.num_of_interleave_):
969 time.append([])
970 for s in range(self.num_of_stage_):
971 time[i].append(0)
972 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
973 for rec in Recompute.TYPE:
974 if self.recompute_considered_[rec]:
975 time[i][s] += self.variables_.get(
976 layer.name_)[rec][i][s].varValue * (
977 layer.forward_time_ +
978 layer.backward_time_rec_[rec])
980 for head in self.layers_sorted_[Layer.type_enum.HEAD]:
981 time[0][0] += head.time_
982 for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
983 time[self.num_of_interleave_ - 1][self.num_of_stage_ -
984 1] += tail.time_
985 return time
987 def get_simulator_forward_time(self) -> list[float]:
988 """Give the time per stage for simulator."""
989 time = []
990 for i in range(self.num_of_interleave_):
991 time.append([])
992 for s in range(self.num_of_stage_):
993 time[i].append(0)
994 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
995 for rec in Recompute.TYPE:
996 if self.recompute_considered_[rec]:
997 time[i][s] += self.variables_[layer.name_][rec][i][
998 s].varValue * (layer.forward_time_)
999 for head in self.layers_sorted_[Layer.type_enum.HEAD]:
1000 time[0][0] += head.time_
1001 for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
1002 time[self.num_of_interleave_ - 1][self.num_of_stage_ -
1003 1] += tail.time_
1004 return time
1006 def get_simulator_recompute_time(self) -> list[float]:
1007 """Give the time per stage for simulator."""
1008 time_all_rec = []
1009 time_no_rec = []
1010 for i in range(self.num_of_interleave_):
1011 time_all_rec.append([])
1012 time_no_rec.append([])
1013 for s in range(self.num_of_stage_):
1014 time_all_rec[i].append(0)
1015 time_no_rec[i].append(0)
1016 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
1017 for rec in Recompute.TYPE:
1018 if self.recompute_considered_[rec]:
1019 time_all_rec[i][s] += self.variables_[
1020 layer.name_][rec][i][s].varValue * (
1021 layer.backward_time_rec_[rec])
1022 time_no_rec[i][s] += self.variables_[
1023 layer.name_][rec][i][s].varValue * (
1024 layer.backward_time_rec_[
1025 Recompute.TYPE.NONE])
1026 return [[r - n for r, n in zip(ar, nr)]
1027 for ar, nr in zip(time_all_rec, time_no_rec)]
1029 def has_some_memory_info(self) -> bool:
1030 """Check if there is some information for memory constraint."""
1031 some_info = False
1032 for rec in Recompute.TYPE:
1033 if self.recompute_considered_[rec]:
1034 some_info = True
1035 return some_info
1037 ############################################
1038 # General Constraint #
1039 ############################################
1040 def add_optional_recompute_constraint(
1041 self, prob: Any, variables: Any,
1042 sorted_layers: Dict[Layer.type_enum, List[Layer]]) -> None:
1043 """Pin unused recomputation variables to zero in the ILP."""
1044 for layer in sorted_layers[Layer.type_enum.BODY]:
1045 for rec in Recompute.TYPE:
1046 if not self.recompute_considered_[rec]:
1047 prob += lpSolver.lpSum(variables[layer.name_][rec]) == 0
1049 def dump_problem(self, folder: Optional[str] = None) -> None:
1050 """Serialize the pulp LP model to ``<folder>/<auto-generated-name>.lp``."""
1051 dump_name = "problem_" + str(self.layers_[0].model_name_)
1052 dump_name += "_" + str(self.max_memory_)
1053 dump_name += "_" + str(self.num_of_interleave_)
1054 dump_name += "_" + str(self.num_of_stage_)
1056 logger.info("dump_problem:out folder = %s", folder)
1057 if folder is not None:
1058 dump_name = os.path.join(folder, dump_name)
1059 dump_name += ".lp"
1060 logger.info("dump problem file: %s", dump_name)
1061 self.problem_.writeLP(dump_name)
1063 def print_results(self) -> None:
1064 """Log the detailed per-layer solver assignment for the solved problem."""
1065 if self.has_some_memory_info():
1066 logger.output("For max memory %s", self.max_memory_)
1067 logger.output("==============")
1068 for body_layer in self.layers_sorted_[Layer.type_enum.BODY]:
1069 layer_name = body_layer.name_
1070 logger.output("For layer: %s", layer_name)
1071 logger.output("=========")
1072 logger.output(" Forward Prop time: %s", body_layer.forward_time_)
1073 for rec in Recompute.TYPE:
1074 if body_layer.recompute_considered_[rec]:
1075 logger.output(" Backward Prop %s time: %s",
1076 Recompute.YAML_NAME[rec], body_layer.backward_time_rec_[rec])
1077 for inter in range(self.num_of_interleave_):
1078 for stage in range(self.num_of_stage_):
1079 parts = []
1080 for rec in Recompute.TYPE:
1081 if self.recompute_considered_[rec]:
1082 value = str(int(self.variables_[layer_name][rec][inter][stage].varValue))
1083 parts.append(value if rec is Recompute.TYPE.NONE else f"+ {value} {rec.name}")
1084 chunk = f" of chunk {inter}" if self.num_of_interleave_ != 1 else ""
1085 logger.output(" Assign %s: %s%s to stage %d",
1086 layer_name, " ".join(parts), chunk, stage)
1087 for s in range(self.num_of_stage_):
1088 logger.debug(
1089 "%s[%s] =%s",
1090 self.MEM_OVERHEAD_NAME,
1091 s,
1092 self.variables_[self.MEM_OVERHEAD_NAME][s].varValue,
1093 )
1095 for v in range(self.num_of_interleave_ - 1):
1096 logger.debug(
1097 "%s[%s] = %s",
1098 self.CHUNKS_SUM,
1099 v,
1100 self.variables_[self.CHUNKS_SUM][v].varValue,
1101 )
1103 for v in range(self.num_of_interleave_ - 1):
1104 logger.debug(
1105 "%s[%s] = %s",
1106 self.PREV_DIFF,
1107 v,
1108 self.variables_[self.PREV_DIFF][v].varValue,
1109 )
1111 logger.debug("%s = %s", self.NEXT_DIFF, self.variables_[self.NEXT_DIFF].varValue)
1112 logger.debug("%s = %s", self.TOTAL_SUM, self.variables_[self.TOTAL_SUM].varValue)
1113 logger.debug("%s = %s", self.MAX_STAGE_TIME, self.variables_[self.MAX_STAGE_TIME].varValue)
1114 logger.debug("%s = %s", self.MAX_LAST_CHUNK, self.variables_[self.MAX_LAST_CHUNK].varValue)
1116 for body_layer in range(len(self.layers_sorted_[Layer.type_enum.BODY]) - 1):
1117 for v in range(self.num_of_interleave_):
1118 for s in range(self.num_of_stage_):
1119 logger.info(
1120 "%s[%s][%s][%s] = %s",
1121 self.LAYER_FRONTIER,
1122 body_layer,
1123 v,
1124 s,
1125 self.variables_[self.LAYER_FRONTIER][body_layer][v][s].varValue,
1126 )
1128 def debug_print_solver_theoretical_memory(self) -> None:
1129 """Log the solver-implied per-stage theoretical memory (debug aid)."""
1130 logger.info("%s Solver Theoretical Memory Analysis %s", "=" * 20, "=" * 20)
1132 if self.vpp_less_memory_:
1133 if self.seq_pipe:
1134 activation_nums = self.compute_activation_seq_nums(
1135 self.num_of_stage_, self.num_of_interleave_, self.seq_split_num_, self.num_of_micro_batch_, True)
1136 else:
1137 activation_nums = self.compute_less_activation_nums(
1138 self.num_of_stage_, self.num_of_interleave_)
1139 else:
1140 if self.seq_pipe:
1141 activation_nums = self.compute_activation_seq_nums(
1142 self.num_of_stage_, self.num_of_interleave_, self.seq_split_num_, self.num_of_micro_batch_, False)
1143 else:
1144 activation_nums = self.compute_activation_nums(
1145 self.num_of_stage_, self.num_of_interleave_, self.num_of_micro_batch_)
1147 # compute theoretical value for each stage
1148 for s in range(self.num_of_stage_):
1149 param_mem = self.stage_param_memory(
1150 self.variables_,
1151 self.layers_sorted_,
1152 s,
1153 self.num_of_stage_,
1154 self.num_of_interleave_
1155 ).value()
1157 act_mem = self.stage_active_memory(
1158 self.variables_,
1159 self.layers_sorted_,
1160 s,
1161 self.num_of_interleave_,
1162 activation_nums
1163 ).value()
1165 # overhead = self.variables_[self.MEM_OVERHEAD_NAME][s].varValue * overhead_factors[s]
1166 overhead = 0
1167 total = param_mem + act_mem + overhead + self.constant_memory_
1169 logger.info("Stage %d Solver Memory Analysis:", s)
1170 logger.info("Parameter Memory: %.2f", param_mem)
1171 logger.info("Activation Memory: %.2f", act_mem)
1172 logger.info("Memory Overhead: %.2f", overhead)
1173 logger.info("Constant Memory: %.2f", self.constant_memory_)
1174 logger.info("Total Theoretical Memory: %.2f", total)
1177 def solve(self, time_limit: int = 90, dump_folder: Optional[str] = None) -> None:
1178 """Solve the ILP problem using PuLP's bundled CBC backend.
1180 Args:
1181 time_limit: Upper bound on solver wall-clock time in seconds.
1182 dump_folder: Directory to write the LP model to; ``None`` skips the dump.
1183 """
1184 logger.info("solve:out folder = %s", dump_folder)
1185 self.dump_problem(dump_folder)
1186 solver = lpSolver.getSolver("PULP_CBC_CMD", timeLimit=time_limit)
1187 self.problem_.solve(solver)
1189 self.print_results()
1191 self.debug_print_solver_theoretical_memory()
1193 for name, result in self.result().items():
1194 logger.output("%s %s %s", name, result, "\n")
1196 def result(self) -> dict[str, list[list[str]]]:
1197 """return schedule distribution for each layer (in the form of a dict)"""
1198 r = {}
1199 for layer in self.layers_sorted_[Layer.type_enum.BODY]:
1200 layer_name = layer.name_
1201 inter = []
1202 for i in range(self.num_of_interleave_):
1203 stage = []
1204 for s in range(self.num_of_stage_):
1205 for rec in Recompute.TYPE:
1206 if self.recompute_considered_[rec]:
1207 stage.append(
1208 str(
1209 self.variables_.get(layer_name)[rec][i]
1210 [s].varValue) + " + ")
1211 inter.append(stage)
1212 r[layer_name] = inter
1213 return r
1215 def _create_problem_(self, description: str) -> lpSolver.LpProblem:
1216 """create the problem"""
1217 prob = lpSolver.LpProblem(description, lpSolver.LpMinimize)
1218 layers_sorted = self.layers_sorted_
1219 num_of_stage = self.num_of_stage_
1220 num_of_interleave = self.num_of_interleave_
1221 num_of_micro_batch = self.num_of_micro_batch_
1222 max_memory = self.max_memory_
1223 # Local variable declaration
1224 # max time that a "main" stage have to take (var to minimize)
1225 pipeline_total_time = lpSolver.LpVariable("pipeline_total_time", 0,
1226 None, lpSolver.LpContinuous)
1228 # Var to Minimize
1229 prob += pipeline_total_time
1231 result = self.add_total_nb_layer_constraint(prob, self.variables_, layers_sorted)
1232 if result is None:
1233 raise RuntimeError("add_total_nb_layer_constraint() returned None.")
1234 # Add if dual to the original layer order constraint
1235 try:
1236 prob = self.add_stage_nb_layer_constraint(
1237 prob, self.variables_, layers_sorted
1238 )
1239 except Exception:
1240 logger.exception("Failed to add stage number layer constraint.")
1241 raise
1242 try:
1243 result = self.add_multimodal_sequence_constraint(prob, self.variables_, layers_sorted)
1244 except Exception:
1245 logger.exception("Failed to add multimodal sequence constraint.")
1246 raise
1248 #self.add_stage_nb_layer_constraint_dual(prob, self.variables_, layers_sorted)
1249 #self.add_multimodal_sequence_constraint_dual(prob, self.variables_, layers_sorted)
1250 try:
1251 result = self.add_multimodal_recompute_constraint(prob, self.variables_, layers_sorted)
1252 if result is None:
1253 raise RuntimeError("add_multimodal_recompute_constraint() returned None.")
1254 except Exception:
1255 logger.exception("Failed to add multimodal recompute constraint.")
1256 raise
1258 try:
1259 result = self.add_performance_constraint(prob, layers_sorted, pipeline_total_time)
1260 if result is None:
1261 raise RuntimeError("add_performance_constraint() returned None.")
1262 prob = result
1263 except Exception:
1264 logger.exception("Failed to add performance constraint.")
1265 raise
1267 constraint = PipelineMemoryConstraint(
1268 prob=prob,
1269 variables=self.variables_,
1270 layers_sorted=layers_sorted,
1271 num_of_stage=num_of_stage,
1272 num_of_interleave=num_of_interleave,
1273 micro_batch=num_of_micro_batch,
1274 memory_limit=max_memory,
1275 )
1276 if self.has_some_memory_info():
1277 self.add_pipeline_memory_constraint(constraint)
1278 return prob
1280 def _create_variables_to_solve_(
1281 self,
1282 num_of_stage: int,
1283 num_of_interleave: int,
1284 layers: dict[Layer.type_enum, list[Layer]],
1285 ):
1286 """create variables to solve"""
1287 variables = {}
1289 variables[self.TOTAL_SUM] = lpSolver.LpVariable(
1290 self.TOTAL_SUM, 0, None, lpSolver.LpContinuous)
1292 chunks_sum_dict = lpSolver.LpVariable.dicts(
1293 name=self.CHUNKS_SUM,
1294 indices=(range(0, self.num_of_interleave_ - 1)),
1295 lowBound=0,
1296 upBound=None,
1297 cat=lpSolver.LpContinuous
1298 )
1299 chunks_sum_list = list(chunks_sum_dict.values())
1300 variables[self.CHUNKS_SUM] = chunks_sum_list
1302 prev_diff_dict = lpSolver.LpVariable.dicts(
1303 name=self.PREV_DIFF,
1304 indices=(range(0, self.num_of_interleave_ - 1)),
1305 lowBound=0,
1306 upBound=None,
1307 cat=lpSolver.LpContinuous
1308 )
1309 prev_diff_list = list(prev_diff_dict.values())
1310 variables[self.PREV_DIFF] = prev_diff_list
1312 layer_frontier_dict = lpSolver.LpVariable.dicts(
1313 name=self.LAYER_FRONTIER,
1314 indices=(
1315 range(1, len(self.layers_sorted_[Layer.type_enum.BODY])),
1316 range(0, self.num_of_interleave_),
1317 range(0, self.num_of_stage_)),
1318 lowBound=0,
1319 upBound=1,
1320 cat=lpSolver.LpBinary
1321 )
1322 layer_frontier_list = list(layer_frontier_dict.values())
1323 variables[self.LAYER_FRONTIER] = layer_frontier_list
1325 rec_frontier_dict = lpSolver.LpVariable.dicts(
1326 name=self.REC_FRONTIER,
1327 indices=(
1328 range(0, self.num_of_interleave_),
1329 range(0, self.num_of_stage_),
1330 range(0, len(self.layers_sorted_[Layer.type_enum.BODY])-1)),
1331 lowBound=0,
1332 upBound=1,
1333 cat=lpSolver.LpBinary
1334 )
1335 rec_frontier_list = list(rec_frontier_dict.values())
1336 variables[self.REC_FRONTIER] = rec_frontier_list
1338 variables[self.NEXT_DIFF] = lpSolver.LpVariable(
1339 self.NEXT_DIFF, 0, None, lpSolver.LpContinuous)
1341 variables[self.MAX_STAGE_TIME] = lpSolver.LpVariable(
1342 self.MAX_STAGE_TIME, 0, None, lpSolver.LpContinuous)
1344 variables[self.MAX_LAST_CHUNK] = lpSolver.LpVariable(
1345 self.MAX_LAST_CHUNK, 0, None, lpSolver.LpContinuous)
1347 lp_variable_dict = lpSolver.LpVariable.dicts(
1348 name=self.MEM_OVERHEAD_NAME,
1349 indices=(range(0, self.num_of_stage_)),
1350 lowBound=0,
1351 upBound=None,
1352 cat=lpSolver.LpInteger,
1353 )
1354 variables_list = list(lp_variable_dict.values())
1355 variables[self.MEM_OVERHEAD_NAME] = variables_list
1357 for layer in layers[Layer.type_enum.BODY]:
1358 variable_dict = lpSolver.LpVariable.dicts(
1359 name=layer.name_,
1360 indices=(
1361 range(0, len(Recompute.TYPE)),
1362 range(0, num_of_interleave),
1363 range(0, num_of_stage),
1364 ),
1365 lowBound=0,
1366 upBound=None,
1367 cat=lpSolver.LpInteger,
1368 )
1369 variable_values = list(variable_dict.values())
1370 interleave_values = []
1371 for interleave in variable_values:
1372 interleave_value = list(interleave.values())
1373 interleave_values.append(interleave_value)
1374 variables[layer.name_] = interleave_values
1376 return variables
1378 ############################################
1379 # Time Constraint #
1380 ############################################
1381 def _max_stage_bound_i_fp(self, layers_sorted, stage_id, inter_f):
1382 bound = lpSolver.LpAffineExpression()
1383 for layer in layers_sorted[Layer.type_enum.BODY]:
1384 for rec in Recompute.TYPE:
1385 if self.recompute_considered_[rec]:
1386 bound += (self.variables_[layer.name_][rec][inter_f][stage_id] *
1387 layer.forward_time_)
1388 return bound
1390 def _max_stage_bound_i_bp(self, layers_sorted, stage_id, inter_b):
1391 bound = lpSolver.LpAffineExpression()
1392 for layer in layers_sorted[Layer.type_enum.BODY]:
1393 for rec in Recompute.TYPE:
1394 if self.recompute_considered_[rec]:
1395 bound += (self.variables_[layer.name_][rec][inter_b][stage_id] *
1396 layer.backward_time_rec_[rec])
1397 return bound
1399 def _max_stage_bound_head_tail(self, layers_sorted, stage_id, inter_f,
1400 inter_b):
1401 """maximize the stage bound of head and tail"""
1402 bound = lpSolver.LpAffineExpression()
1403 if stage_id == 0:
1404 if inter_f == 0:
1405 for head in layers_sorted[Layer.type_enum.HEAD]:
1406 bound += head.time_
1407 if inter_b == 0:
1408 for head in layers_sorted[Layer.type_enum.HEAD]:
1409 bound += head.time_ * 2
1410 if stage_id == self.num_of_stage_ - 1:
1411 if inter_f == self.num_of_interleave_ - 1:
1412 for tail in layers_sorted[Layer.type_enum.TAIL]:
1413 bound += tail.time_
1414 if inter_b == self.num_of_interleave_ - 1:
1415 for tail in layers_sorted[Layer.type_enum.TAIL]:
1416 bound += tail.time_ * 2
1417 return bound
1419 def _total_sum(self, layers_sorted):
1420 """sum up the layer time"""
1421 bound = lpSolver.LpAffineExpression()
1422 for layer in layers_sorted[Layer.type_enum.BODY]:
1423 for rec in Recompute.TYPE:
1424 if self.recompute_considered_[rec]:
1425 for inter in range(self.num_of_interleave_):
1426 for stage in range(self.num_of_stage_):
1427 bound += self.variables_[layer.name_][rec][inter][stage] * (
1428 layer.forward_time_ +
1429 layer.backward_time_rec_[rec])
1430 return bound
1432 def body_layer_time(self, prop: "SappSolver.PROP_PHASE", layer: Layer,
1433 inter: int, stage: int) -> Any:
1434 """Return a forward or backward time LP expression for ``layer`` at ``(inter, stage)``."""
1435 if prop == self.PROP_PHASE.FW:
1436 bound = lpSolver.lpSum(
1437 self.variables_[layer.name_][rec][inter][stage] * layer.forward_time_
1438 for rec in Recompute.TYPE if self.recompute_considered_[rec])
1439 else:
1440 bound = lpSolver.lpSum(
1441 self.variables_[layer.name_][rec][inter][stage] * layer.backward_time_rec_[rec]
1442 for rec in Recompute.TYPE if self.recompute_considered_[rec])
1444 return bound
1446 def micro_batch_time(self, prop: "SappSolver.PROP_PHASE",
1447 layers_sorted: Dict[Layer.type_enum, List[Layer]],
1448 inter: int, stage: int) -> Any:
1449 """Return the total micro-batch time LP expression at ``(inter, stage)``."""
1450 bound = lpSolver.LpAffineExpression()
1451 if prop == self.PROP_PHASE.FW:
1452 for layer in layers_sorted[Layer.type_enum.BODY]:
1453 bound = self.body_layer_time(prop, layer, inter, stage)
1454 if stage == 0 and inter == 0:
1455 for head in layers_sorted[Layer.type_enum.HEAD]:
1456 bound += head.time_
1457 if stage == self.num_of_stage_ - 1 and inter == self.num_of_interleave_ - 1:
1458 for tail in layers_sorted[Layer.type_enum.TAIL]:
1459 bound += tail.time_
1460 else:
1461 for layer in layers_sorted[Layer.type_enum.BODY]:
1462 bound = self.body_layer_time(prop, layer, inter, stage)
1463 if stage == 0 and inter == 0:
1464 for head in layers_sorted[Layer.type_enum.HEAD]:
1465 bound += head.time_ * 2
1466 if stage == self.num_of_stage_ - 1 and inter == self.num_of_interleave_ - 1:
1467 for tail in layers_sorted[Layer.type_enum.TAIL]:
1468 bound += tail.time_ * 2
1469 return bound
1471 def _chunks_sum(self, layers_sorted, v):
1472 """sum up the warm-up and cool-down time of a given chunk"""
1473 bound = lpSolver.LpAffineExpression()
1474 for stage in range(self.num_of_stage_):
1475 bound += self.micro_batch_time(self.PROP_PHASE.FW, layers_sorted, v, stage)
1476 bound += self.micro_batch_time(self.PROP_PHASE.BW, layers_sorted, v, stage)
1477 # normalize
1478 bound = bound / self.num_of_stage_
1479 return bound
1481 def _prev_diff_sum(self, layers_sorted, prob, v):
1482 """models bubble time for the first diagonal (forward, interleave 0)"""
1483 max_prev_stages = lpSolver.LpVariable.dicts(
1484 name="max_prev_stages_" + str(v),
1485 indices=(range(self.num_of_stage_)),
1486 lowBound=0,
1487 upBound=None,
1488 cat=lpSolver.LpContinuous,
1489 )
1491 diff_with_prev_stages = lpSolver.LpVariable.dicts(
1492 name="diff_with_prev_stages_" + str(v),
1493 indices=(range(self.num_of_stage_)),
1494 lowBound=0,
1495 upBound=None,
1496 cat=lpSolver.LpContinuous,
1497 )
1499 bound = lpSolver.LpAffineExpression()
1501 head_time = 0
1502 for head in layers_sorted[Layer.type_enum.HEAD]:
1503 head_time = head.time_
1505 prob += max_prev_stages[0] >= (self.micro_batch_time(
1506 self.PROP_PHASE.FW, layers_sorted, v, 0)) - head_time
1508 for stage in range(1, self.num_of_stage_):
1509 prob += max_prev_stages[stage] >= max_prev_stages[stage - 1]
1510 prob += max_prev_stages[stage] >= (self.micro_batch_time(
1511 self.PROP_PHASE.FW, layers_sorted, v, stage))
1514 prob += diff_with_prev_stages[stage] >= (
1515 max_prev_stages[stage - 1] - self.micro_batch_time(
1516 self.PROP_PHASE.FW, layers_sorted, v, stage))
1518 bound += self.num_of_micro_batch_ * lpSolver.lpSum(
1519 diff_with_prev_stages[s] for s in range(1, self.num_of_stage_))
1520 return bound
1522 def _next_diff_sum(self, layers_sorted, prob):
1523 """models bubble time for the last diagonal (forward, last chunk)"""
1524 last_chunk = self.num_of_interleave_ - 1
1525 max_next_stages = lpSolver.LpVariable.dicts(
1526 name="max_next_stages",
1527 indices=(range(self.num_of_stage_)),
1528 lowBound=0,
1529 upBound=None,
1530 cat=lpSolver.LpContinuous,
1531 )
1533 diff_with_next_stages = lpSolver.LpVariable.dicts(
1534 name="diff_with_next_stages",
1535 indices=(range(self.num_of_stage_)),
1536 lowBound=0,
1537 upBound=None,
1538 cat=lpSolver.LpContinuous,
1539 )
1541 bound = lpSolver.LpAffineExpression()
1543 prob += max_next_stages[self.num_of_stage_ -
1544 1] >= (self.micro_batch_time(
1545 self.PROP_PHASE.FW, layers_sorted, last_chunk,
1546 self.num_of_stage_ - 1))
1548 for stage in reversed(range(0, self.num_of_stage_ - 1)):
1549 prob += max_next_stages[stage] >= max_next_stages[stage + 1]
1550 prob += max_next_stages[stage] >= (self.micro_batch_time(
1551 self.PROP_PHASE.FW, layers_sorted, last_chunk, stage))
1553 prob += diff_with_next_stages[stage] >= (
1554 max_next_stages[stage + 1] - self.micro_batch_time(
1555 self.PROP_PHASE.FW, layers_sorted, last_chunk, stage))
1557 bound += self.num_of_micro_batch_ * lpSolver.lpSum(
1558 diff_with_next_stages[s] for s in range(self.num_of_stage_ - 1))
1559 return bound