已合并
feat: add mhc recompute #3619
wuweiqiang24创建于 7月5日
feat: add mhc recompute #3619
已合并
共 14 个文件变更+617-394
| @@ -6,12 +6,14 @@ from torch import Tensor | |||
| 6 | 6 | ||
| 7 | from megatron.core import tensor_parallel, parallel_state, mpu | 7 | from megatron.core import tensor_parallel, parallel_state, mpu |
| 8 | from megatron.core.packed_seq_params import PackedSeqParams | 8 | from megatron.core.packed_seq_params import PackedSeqParams |
| 9 | +from mindspeed.core.memory.recompute.recompute_common import get_recompute_priority | ||
| 9 | from mindspeed.core.memory.adaptive_memory.adaptive_memory_swap_manager import SwapManager as AdaptiveMemorySwapManager | 10 | from mindspeed.core.memory.adaptive_memory.adaptive_memory_swap_manager import SwapManager as AdaptiveMemorySwapManager |
| 10 | from mindspeed.core.memory.adaptive_recomputing.swap_manager import SwapManager as AdaptiveRecomputingSwapManager | 11 | from mindspeed.core.memory.adaptive_recomputing.swap_manager import SwapManager as AdaptiveRecomputingSwapManager |
| 11 | 12 | ||
| 12 | 13 | ||
| 13 | def swap_out_by_size(size): | 14 | def swap_out_by_size(size): |
| 14 | from megatron.training import get_args | 15 | from megatron.training import get_args |
| 16 | + | ||
| 15 | args = get_args() | 17 | args = get_args() |
| 16 | if args.adaptive_memory_optimization: | 18 | if args.adaptive_memory_optimization: |
| 17 | return AdaptiveMemorySwapManager().swap_out_by_size(size) | 19 | return AdaptiveMemorySwapManager().swap_out_by_size(size) |
| @@ -21,26 +23,28 @@ def swap_out_by_size(size): | |||
| 21 | 23 | ||
| 22 | def linear_forward_main_grad_wrapper(forward_func): | 24 | def linear_forward_main_grad_wrapper(forward_func): |
| 23 | 25 | ||
| 24 | - def linear_forward_main_grad(ctx, | 26 | + def linear_forward_main_grad( |
| 25 | - inputs, | 27 | + ctx, |
| 26 | - weight, | 28 | + inputs, |
| 27 | - bias, | 29 | + weight, |
| 28 | - gradient_accumulation_fusion, | 30 | + bias, |
| 29 | - allreduce_dgrad, | 31 | + gradient_accumulation_fusion, |
| 30 | - sequence_parallel, | 32 | + allreduce_dgrad, |
| 31 | - grad_output_buffer, | 33 | + sequence_parallel, |
| 32 | - wgrad_deferral_limit, | 34 | + grad_output_buffer, |
| 33 | - ): | 35 | + wgrad_deferral_limit, |
| 34 | - output = forward_func(ctx, | 36 | + ): |
| 35 | - inputs, | 37 | + output = forward_func( |
| 36 | - weight, | 38 | + ctx, |
| 37 | - bias, | 39 | + inputs, |
| 38 | - gradient_accumulation_fusion, | 40 | + weight, |
| 39 | - allreduce_dgrad, | 41 | + bias, |
| 40 | - sequence_parallel, | 42 | + gradient_accumulation_fusion, |
| 41 | - grad_output_buffer, | 43 | + allreduce_dgrad, |
| 42 | - wgrad_deferral_limit, | 44 | + sequence_parallel, |
| 43 | - ) | 45 | + grad_output_buffer, |
| 46 | + wgrad_deferral_limit, | ||
| 47 | + ) | ||
| 44 | ctx.weight = weight | 48 | ctx.weight = weight |
| 45 | return output | 49 | return output |
| 46 | 50 | ||
| @@ -52,6 +56,7 @@ def linear_backward_main_grad_wrapper(backward_func): | |||
| 52 | def linear_backward_main_grad(ctx, grad_output): | 56 | def linear_backward_main_grad(ctx, grad_output): |
| 53 | class NewCtx: | 57 | class NewCtx: |
| 54 | pass | 58 | pass |
| 59 | + | ||
| 55 | new_ctx = NewCtx() | 60 | new_ctx = NewCtx() |
| 56 | inputs, _ = ctx.saved_tensors | 61 | inputs, _ = ctx.saved_tensors |
| 57 | for key in dir(ctx): | 62 | for key in dir(ctx): |
| @@ -71,34 +76,33 @@ def linear_backward_main_grad_wrapper(backward_func): | |||
| 71 | 76 | ||
| 72 | 77 | ||
| 73 | def transformer_block_checkpointed_forward( | 78 | def transformer_block_checkpointed_forward( |
| 74 | - self, | 79 | + self, |
| 75 | - hidden_states: Tensor, | 80 | + hidden_states: Tensor, |
| 76 | - attention_mask: Tensor, | 81 | + attention_mask: Tensor, |
| 77 | - context: Tensor, | 82 | + context: Tensor, |
| 78 | - context_mask: Tensor, | 83 | + context_mask: Tensor, |
| 79 | - rotary_pos_emb: Tensor, | 84 | + rotary_pos_emb: Tensor, |
| 80 | - attention_bias: Tensor, | 85 | + attention_bias: Tensor, |
| 81 | - packed_seq_params: PackedSeqParams, | 86 | + packed_seq_params: PackedSeqParams, |
| 82 | - use_inner_fp8_context | 87 | + use_inner_fp8_context, |
| 83 | ): | 88 | ): |
| 84 | """Forward method with activation checkpointing.""" | 89 | """Forward method with activation checkpointing.""" |
| 85 | 90 | ||
| 86 | def custom(start: int, end: int): | 91 | def custom(start: int, end: int): |
| 87 | def custom_forward( | 92 | def custom_forward( |
| 88 | - hidden_states, | 93 | + hidden_states, |
| 89 | - attention_mask, | 94 | + attention_mask, |
| 90 | - context, | 95 | + context, |
| 91 | - context_mask, | 96 | + context_mask, |
| 92 | - rotary_pos_emb, | 97 | + rotary_pos_emb, |
| 93 | ): | 98 | ): |
| 94 | from megatron.core.fp8_utils import get_fp8_context | 99 | from megatron.core.fp8_utils import get_fp8_context |
| 95 | from contextlib import nullcontext | 100 | from contextlib import nullcontext |
| 101 | + | ||
| 96 | for index in range(start, end): | 102 | for index in range(start, end): |
| 97 | layer = self._get_layer(index) | 103 | layer = self._get_layer(index) |
| 98 | inner_fp8_context = ( | 104 | inner_fp8_context = ( |
| 99 | - get_fp8_context(self.config, layer.layer_number - 1) | 105 | + get_fp8_context(self.config, layer.layer_number - 1) if use_inner_fp8_context else nullcontext() |
| 100 | - if use_inner_fp8_context | ||
| 101 | - else nullcontext() | ||
| 102 | ) | 106 | ) |
| 103 | with inner_fp8_context: | 107 | with inner_fp8_context: |
| 104 | hidden_states, context = layer( | 108 | hidden_states, context = layer( |
| @@ -164,12 +168,6 @@ def transformer_block_checkpointed_forward( | |||
| 164 | rotary_pos_emb, | 168 | rotary_pos_emb, |
| 165 | ) | 169 | ) |
| 166 | elif self.config.recompute_method == 'block': | 170 | elif self.config.recompute_method == 'block': |
| 167 | - vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() | ||
| 168 | - vpp_size = self.config.virtual_pipeline_model_parallel_size | ||
| 169 | - if vpp_rank is None or not getattr(self.config, 'enable_recompute_layers_per_pp_rank', False): | ||
| 170 | - vpp_rank = 0 | ||
| 171 | - if vpp_size is None or not getattr(self.config, 'enable_recompute_layers_per_pp_rank', False): | ||
| 172 | - vpp_size = 1 | ||
| 173 | for layer_idx in range(self.num_layers_per_pipeline_rank): | 171 | for layer_idx in range(self.num_layers_per_pipeline_rank): |
| 174 | # The number of layers each pipeline rank recomputes is self.recompute_num_layers. | 172 | # The number of layers each pipeline rank recomputes is self.recompute_num_layers. |
| 175 | # If self.recompute_num_layers cannot divide exactly the number of layers in each pp rank, | 173 | # If self.recompute_num_layers cannot divide exactly the number of layers in each pp rank, |
| @@ -180,16 +178,14 @@ def transformer_block_checkpointed_forward( | |||
| 180 | # Stage 1: [2, 3] [6, 7] | 178 | # Stage 1: [2, 3] [6, 7] |
| 181 | # With self.recompute_num_layers = 2, we will recompute layers 0,4 for stage 0, and 2,6 for stage 1. | 179 | # With self.recompute_num_layers = 2, we will recompute layers 0,4 for stage 0, and 2,6 for stage 1. |
| 182 | # With self.recompute_num_layers = 3, we will recompute layers 0,1,4 for stage 0, and 2,3,6 for stage 1. | 180 | # With self.recompute_num_layers = 3, we will recompute layers 0,1,4 for stage 0, and 2,3,6 for stage 1. |
| 183 | - def should_recompute(): | 181 | + layer = self._get_layer(layer_idx) |
| 184 | - if getattr(self.config, 'reduce_recompute_for_last_chunk', False): | 182 | + recompute_priority = get_recompute_priority(self.config, layer.layer_number) |
| 185 | - def is_last_layer(): | 183 | + should_recompute_layer = recompute_priority < self.config.recompute_num_layers |
| 186 | - return (layer_idx == self.num_layers_per_pipeline_rank - 1) and mpu.is_pipeline_last_stage() | 184 | + if getattr(self.config, 'reduce_recompute_for_last_chunk', False): |
| 185 | + is_last_layer = layer_idx == self.num_layers_per_pipeline_rank - 1 and mpu.is_pipeline_last_stage() | ||
| 186 | + should_recompute_layer = should_recompute_layer and not is_last_layer | ||
| 187 | 187 | ||
| 188 | - return ((layer_idx * vpp_size + vpp_rank) < self.config.recompute_num_layers) and not is_last_layer() | 188 | + if should_recompute_layer and not getattr(self.config, 'swap_attention', False): |
| 189 | - else: | ||
| 190 | - return (layer_idx * vpp_size + vpp_rank) < self.config.recompute_num_layers | ||
| 191 | - | ||
| 192 | - if should_recompute() and not getattr(self.config, 'swap_attention', False): | ||
| 193 | hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1)) | 189 | hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1)) |
| 194 | else: | 190 | else: |
| 195 | hidden_states, context = custom(layer_idx, layer_idx + 1)( | 191 | hidden_states, context = custom(layer_idx, layer_idx + 1)( |
| @@ -35,18 +35,16 @@ class CheckpointWithoutOutput: | |||
| 35 | self.fwd_cuda_rng_state = None | 35 | self.fwd_cuda_rng_state = None |
| 36 | self.fwd_cuda_rng_state_tracker = None | 36 | self.fwd_cuda_rng_state_tracker = None |
| 37 | self.outputs = None | 37 | self.outputs = None |
| 38 | + self.ctx = None | ||
| 38 | self.get_cuda_rng_tracker = get_cuda_rng_tracker_func | 39 | self.get_cuda_rng_tracker = get_cuda_rng_tracker_func |
| 39 | 40 | ||
| 40 | def checkpoint(self, run_function, distribute_saved_activations, *args): | 41 | def checkpoint(self, run_function, distribute_saved_activations, *args): |
| 41 | self.run_function = run_function | 42 | self.run_function = run_function |
| 42 | 43 | ||
| 43 | if distribute_saved_activations: | 44 | if distribute_saved_activations: |
| 44 | - raise RuntimeError( | 45 | + raise RuntimeError("CheckpointFunctionWithoutOutput does not support distribute_saved_activations") |
| 45 | - "CheckpointFunctionWithoutOutput does not support " | ||
| 46 | - "distribute_saved_activations" | ||
| 47 | - ) | ||
| 48 | 46 | ||
| 49 | - #Copy the rng states. | 47 | + # Copy the rng states. |
| 50 | self.fwd_cpu_rng_state = torch.get_rng_state() | 48 | self.fwd_cpu_rng_state = torch.get_rng_state() |
| 51 | self.fwd_cuda_rng_state = torch.cuda.get_rng_state() | 49 | self.fwd_cuda_rng_state = torch.cuda.get_rng_state() |
| 52 | self.fwd_cuda_rng_state_tracker = self.get_cuda_rng_tracker().get_states() | 50 | self.fwd_cuda_rng_state_tracker = self.get_cuda_rng_tracker().get_states() |
| @@ -64,10 +62,7 @@ class CheckpointWithoutOutput: | |||
| 64 | 62 | ||
| 65 | def recompute(self, _): | 63 | def recompute(self, _): |
| 66 | if not torch.autograd._is_checkpoint_valid(): | 64 | if not torch.autograd._is_checkpoint_valid(): |
| 67 | - raise RuntimeError( | 65 | + raise RuntimeError("Checkpointing is not compatible with .grad(), please use .backward() if possible") |
| 68 | - "Checkpointing is not compatible with .grad(), " | ||
| 69 | - "please use .backward() if possible" | ||
| 70 | - ) | ||
| 71 | 66 | ||
| 72 | # Store the current states. | 67 | # Store the current states. |
| 73 | cur_cpu_rng_state = torch.get_rng_state() | 68 | cur_cpu_rng_state = torch.get_rng_state() |
| @@ -105,25 +100,112 @@ class CheckpointWithoutOutput: | |||
| 105 | self.ctx = None | 100 | self.ctx = None |
| 106 | 101 | ||
| 107 | 102 | ||
| 108 | -def should_recompute(config, layer_number, num_recompute): | 103 | +def _get_pipeline_model_parallel_layout(config): |
| 104 | + layout = getattr(config, 'pipeline_model_parallel_layout', None) | ||
| 105 | + if layout is None: | ||
| 106 | + return None | ||
| 107 | + | ||
| 108 | + from mindspeed.core.pipeline_parallel.pipeline_model_parallel_layout.layout import ( | ||
| 109 | + PipelineParallelLayerLayout, | ||
| 110 | + ) | ||
| 111 | + | ||
| 112 | + if isinstance(layout, PipelineParallelLayerLayout): | ||
| 113 | + return layout | ||
| 114 | + | ||
| 115 | + pp_size = getattr(config, 'pipeline_model_parallel_size', None) | ||
| 116 | + if pp_size is None: | ||
| 117 | + pp_size = mpu.get_pipeline_model_parallel_world_size() | ||
| 118 | + if isinstance(layout, str): | ||
| 119 | + return PipelineParallelLayerLayout.from_str(layout, pp_size) | ||
| 120 | + if isinstance(layout, list): | ||
| 121 | + return PipelineParallelLayerLayout(layout, pp_size) | ||
| 122 | + raise TypeError( | ||
| 123 | + f'pipeline_model_parallel_layout must be a str, list, or PipelineParallelLayerLayout, but got {type(layout)}' | ||
| 124 | + ) | ||
| 125 | + | ||
| 126 | + | ||
| 127 | +def _get_layout_recompute_priority(layout, layer_number, enable_per_pp_rank): | ||
| 128 | + from mindspeed.core.pipeline_parallel.pipeline_model_parallel_layout.layout import LayerType | ||
| 129 | + | ||
| 130 | + pp_rank = mpu.get_pipeline_model_parallel_rank() | ||
| 131 | + vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() or 0 | ||
| 132 | + vpp_size = layout.virtual_pipeline_model_parallel_size | ||
| 133 | + if not 0 <= vpp_rank < vpp_size: | ||
| 134 | + raise ValueError(f'Invalid virtual pipeline rank {vpp_rank} for layout with {vpp_size} virtual stages.') | ||
| 135 | + | ||
| 136 | + decoder_layer_id = layer_number - 1 | ||
| 137 | + layer_ids_per_vpp_rank = [[] for _ in range(vpp_size)] | ||
| 138 | + decoder_layer_offset = 0 | ||
| 139 | + for current_vpp_rank in range(vpp_size): | ||
| 140 | + for current_pp_rank in range(layout.pipeline_model_parallel_size): | ||
| 141 | + num_decoder_layers = layout.layout[current_pp_rank][current_vpp_rank].count(LayerType.decoder) | ||
| 142 | + layer_ids = list(range(decoder_layer_offset, decoder_layer_offset + num_decoder_layers)) | ||
| 143 | + if current_pp_rank == pp_rank: | ||
| 144 | + layer_ids_per_vpp_rank[current_vpp_rank] = layer_ids | ||
| 145 | + decoder_layer_offset += num_decoder_layers | ||
| 146 | + | ||
| 147 | + current_chunk_layer_ids = layer_ids_per_vpp_rank[vpp_rank] | ||
| 148 | + if decoder_layer_id not in current_chunk_layer_ids: | ||
| 149 | + raise ValueError( | ||
| 150 | + f'Decoder layer {layer_number} is not present in pipeline layout for ' | ||
| 151 | + f'pp_rank={pp_rank}, vpp_rank={vpp_rank}. layout={layout}' | ||
| 152 | + ) | ||
| 153 | + | ||
| 154 | + local_layer_index = current_chunk_layer_ids.index(decoder_layer_id) | ||
| 155 | + if not enable_per_pp_rank: | ||
| 156 | + return local_layer_index | ||
| 157 | + | ||
| 158 | + recompute_priority = 0 | ||
| 159 | + max_chunk_size = max((len(layer_ids) for layer_ids in layer_ids_per_vpp_rank), default=0) | ||
| 160 | + for layer_index in range(max_chunk_size): | ||
| 161 | + for layer_ids in layer_ids_per_vpp_rank: | ||
| 162 | + if layer_index >= len(layer_ids): | ||
| 163 | + continue | ||
| 164 | + if layer_ids[layer_index] == decoder_layer_id: | ||
| 165 | + return recompute_priority | ||
| 166 | + recompute_priority += 1 | ||
| 167 | + | ||
| 168 | + raise RuntimeError(f'Failed to calculate recompute priority for decoder layer {layer_number}.') | ||
| 169 | + | ||
| 170 | + | ||
| 171 | +def get_recompute_priority(config, layer_number, enable_per_pp_rank=None): | ||
| 172 | + """Return the layer recompute priority for uniform or custom pipeline layouts.""" | ||
| 173 | + if layer_number is None: | ||
| 174 | + raise ValueError('layer_number must not be None when calculating recompute priority.') | ||
| 175 | + | ||
| 176 | + if enable_per_pp_rank is None: | ||
| 177 | + enable_per_pp_rank = getattr(config, 'enable_recompute_layers_per_pp_rank', False) | ||
| 178 | + | ||
| 179 | + layout = _get_pipeline_model_parallel_layout(config) | ||
| 180 | + if layout is not None: | ||
| 181 | + return _get_layout_recompute_priority(layout, layer_number, enable_per_pp_rank) | ||
| 182 | + | ||
| 109 | vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() | 183 | vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() |
| 110 | - vpp_size = config.virtual_pipeline_model_parallel_size | 184 | + vpp_size = getattr(config, 'virtual_pipeline_model_parallel_size', None) |
| 111 | - pp_size = config.pipeline_model_parallel_size | 185 | + pp_size = getattr(config, 'pipeline_model_parallel_size', None) |
| 112 | 186 | ||
| 113 | if vpp_size is not None: | 187 | if vpp_size is not None: |
| 114 | - layer_per_chunk = config.num_layers_per_virtual_pipeline_stage | 188 | + layer_per_chunk = getattr(config, 'num_layers_per_virtual_pipeline_stage', None) |
| 189 | + if layer_per_chunk is None: | ||
| 190 | + layer_per_chunk = config.num_layers // pp_size // vpp_size | ||
| 115 | elif pp_size is not None: | 191 | elif pp_size is not None: |
| 116 | layer_per_chunk = config.num_layers // pp_size | 192 | layer_per_chunk = config.num_layers // pp_size |
| 117 | else: | 193 | else: |
| 118 | layer_per_chunk = config.num_layers | 194 | layer_per_chunk = config.num_layers |
| 119 | 195 | ||
| 120 | - if vpp_rank is None or not getattr(config, 'enable_recompute_layers_per_pp_rank', False): | 196 | + if vpp_rank is None or not enable_per_pp_rank: |
| 121 | vpp_rank = 0 | 197 | vpp_rank = 0 |
| 122 | - if vpp_size is None or not getattr(config, 'enable_recompute_layers_per_pp_rank', False): | 198 | + if vpp_size is None or not enable_per_pp_rank: |
| 123 | vpp_size = 1 | 199 | vpp_size = 1 |
| 200 | + return ((layer_number - 1) % layer_per_chunk) * vpp_size + vpp_rank | ||
| 124 | 201 | ||
| 125 | - recompute_priority = ((layer_number - 1) % layer_per_chunk) * vpp_size + vpp_rank | 202 | + |
| 203 | +def should_recompute(config, layer_number, num_recompute): | ||
| 126 | full_recompute_layers = config.recompute_num_layers | 204 | full_recompute_layers = config.recompute_num_layers |
| 205 | + if not full_recompute_layers and num_recompute is None: | ||
| 206 | + return True | ||
| 207 | + | ||
| 208 | + recompute_priority = get_recompute_priority(config, layer_number) | ||
| 127 | 209 | ||
| 128 | if full_recompute_layers: | 210 | if full_recompute_layers: |
| 129 | if recompute_priority < full_recompute_layers: | 211 | if recompute_priority < full_recompute_layers: |
| @@ -4,7 +4,53 @@ from functools import wraps | |||
| 4 | import torch | 4 | import torch |
| 5 | 5 | ||
| 6 | from mindspeed.core.memory.swap_layer_input.swap_layer_input_manager import SwapLayerInputManager | 6 | from mindspeed.core.memory.swap_layer_input.swap_layer_input_manager import SwapLayerInputManager |
| 7 | -from mindspeed.core.transformer.moe.moe_feature.fb_overlap.modules.utils import LayerGraph | 7 | + |
| 8 | + | ||
| 9 | +def _get_hidden_states(args, kwargs): | ||
| 10 | + hidden_states = kwargs.get('hidden_states') | ||
| 11 | + if hidden_states is not None: | ||
| 12 | + return hidden_states | ||
| 13 | + if args and isinstance(args[0], torch.Tensor): | ||
| 14 | + return args[0] | ||
| 15 | + return None | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +def _get_backward_layer_graph(args, kwargs, name, position): | ||
| 19 | + layer_graph = kwargs.get(name) | ||
| 20 | + if layer_graph is not None and hasattr(layer_graph, 'layer') and hasattr(layer_graph, 'layer_input'): | ||
| 21 | + return layer_graph | ||
| 22 | + if len(args) > position: | ||
| 23 | + layer_graph = args[position] | ||
| 24 | + if layer_graph is not None and hasattr(layer_graph, 'layer') and hasattr(layer_graph, 'layer_input'): | ||
| 25 | + return layer_graph | ||
| 26 | + return None | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +def _swap_out_layer_input(layer, hidden_states): | ||
| 30 | + if ( | ||
| 31 | + layer is None | ||
| 32 | + or getattr(layer, 'is_mtp', False) | ||
| 33 | + or not hasattr(layer, 'swap_manager') | ||
| 34 | + or not isinstance(hidden_states, torch.Tensor) | ||
| 35 | + ): | ||
| 36 | + return None | ||
| 37 | + | ||
| 38 | + hidden_states.swap_this_tensor = True | ||
| 39 | + return layer.swap_manager.swap_out_tensors([hidden_states]) | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +def _restore_swap_entry(layer_graph): | ||
| 43 | + """Restore the exact input associated with a forward layer graph.""" | ||
| 44 | + if layer_graph is None or not hasattr(layer_graph, 'layer') or not hasattr(layer_graph, 'layer_input'): | ||
| 45 | + return | ||
| 46 | + | ||
| 47 | + swap_entry = getattr(layer_graph, 'swap_layer_input_entry', None) | ||
| 48 | + manager = getattr(layer_graph.layer, 'swap_manager', None) | ||
| 49 | + if swap_entry is None or manager is None: | ||
| 50 | + return | ||
| 51 | + | ||
| 52 | + manager.restore_swap_entry(swap_entry) | ||
| 53 | + layer_graph.swap_layer_input_entry = None | ||
| 8 | 54 | ||
| 9 | 55 | ||
| 10 | def swap_layer_input_init_wrapper(fn): | 56 | def swap_layer_input_init_wrapper(fn): |
| @@ -16,7 +62,9 @@ def swap_layer_input_init_wrapper(fn): | |||
| 16 | 62 | ||
| 17 | # MindSpeed-LLM MTP calls TransformerLayer twice, and the first call needs to be removed. | 63 | # MindSpeed-LLM MTP calls TransformerLayer twice, and the first call needs to be removed. |
| 18 | if getattr(self, 'is_mtp', False): | 64 | if getattr(self, 'is_mtp', False): |
| 19 | - SwapLayerInputManager.manager_map["default"].pop(-1) | 65 | + managers = SwapLayerInputManager.manager_map.get('default', []) |
| 66 | + if managers: | ||
| 67 | + managers.pop() | ||
| 20 | 68 | ||
| 21 | return result | 69 | return result |
| 22 | 70 | ||
| @@ -26,21 +74,19 @@ def swap_layer_input_init_wrapper(fn): | |||
| 26 | def swap_layer_input_forward_wrapper(fn): | 74 | def swap_layer_input_forward_wrapper(fn): |
| 27 | 75 | ||
| 28 | def wrapper(self, *args, **kwargs): | 76 | def wrapper(self, *args, **kwargs): |
| 29 | - if not hasattr(self, 'swap_manager'): | 77 | + if getattr(self, 'is_mtp', False) or not hasattr(self, 'swap_manager'): |
| 30 | return fn(self, *args, **kwargs) | 78 | return fn(self, *args, **kwargs) |
| 31 | 79 | ||
| 32 | - hidden_states = None | 80 | + hidden_states = _get_hidden_states(args, kwargs) |
| 33 | - if 'hidden_states' in kwargs: | 81 | + if isinstance(hidden_states, torch.Tensor): |
| 34 | - hidden_states = kwargs['hidden_states'] | 82 | + hidden_states.swap_this_tensor = True |
| 35 | - elif len(args) > 0 and isinstance(args[0], torch.Tensor): | ||
| 36 | - hidden_states = args[0] | ||
| 37 | - hidden_states.swap_this_tensor = True | ||
| 38 | 83 | ||
| 39 | - if not torch.is_grad_enabled(): | 84 | + if not torch.is_grad_enabled() and isinstance(hidden_states, torch.Tensor): |
| 40 | self.swap_manager.swap_out_tensors([hidden_states]) | 85 | self.swap_manager.swap_out_tensors([hidden_states]) |
| 41 | self.swap_manager.forward_hook() | 86 | self.swap_manager.forward_hook() |
| 87 | + | ||
| 42 | result = fn(self, *args, **kwargs) | 88 | result = fn(self, *args, **kwargs) |
| 43 | - if result[0].requires_grad: | 89 | + if isinstance(result, tuple) and result and isinstance(result[0], torch.Tensor) and result[0].requires_grad: |
| 44 | result[0].register_hook(self.swap_manager.backward_hook) | 90 | result[0].register_hook(self.swap_manager.backward_hook) |
| 45 | return result | 91 | return result |
| 46 | 92 | ||
| @@ -50,23 +96,17 @@ def swap_layer_input_forward_wrapper(fn): | |||
| 50 | def swap_layer_input_fboverlap_forward_wrapper(fn): | 96 | def swap_layer_input_fboverlap_forward_wrapper(fn): |
| 51 | 97 | ||
| 52 | def wrapper(self, *args, **kwargs): | 98 | def wrapper(self, *args, **kwargs): |
| 53 | - if not hasattr(self, 'swap_manager'): | 99 | + swap_entry = _swap_out_layer_input(self, _get_hidden_states(args, kwargs)) |
| 54 | - return fn(self, *args, **kwargs) | 100 | + try: |
| 55 | - | ||
| 56 | - hidden_states = None | ||
| 57 | - if 'hidden_states' in kwargs: | ||
| 58 | - hidden_states = kwargs['hidden_states'] | ||
| 59 | - elif len(args) > 0 and isinstance(args[0], torch.Tensor): | ||
| 60 | - hidden_states = args[0] | ||
| 61 | - | ||
| 62 | - if kwargs.get('checkpoint', False): | ||
| 63 | - hidden_states.swap_this_tensor = True | ||
| 64 | - self.swap_manager.swap_out_tensors([hidden_states]) | ||
| 65 | - result = fn(self, *args, **kwargs) | ||
| 66 | - self.swap_manager.wait_swap_out() | ||
| 67 | - else: | ||
| 68 | result = fn(self, *args, **kwargs) | 101 | result = fn(self, *args, **kwargs) |
| 102 | + finally: | ||
| 103 | + if swap_entry is not None: | ||
| 104 | + self.swap_manager.wait_swap_out(swap_entry) | ||
| 69 | 105 | ||
| 106 | + if isinstance(result, tuple) and len(result) > 2 and swap_entry is not None: | ||
| 107 | + layer_graph = result[2] | ||
| 108 | + if layer_graph is not None and hasattr(layer_graph, 'layer') and hasattr(layer_graph, 'layer_input'): | ||
| 109 | + layer_graph.swap_layer_input_entry = swap_entry | ||
| 70 | return result | 110 | return result |
| 71 | 111 | ||
| 72 | return wrapper | 112 | return wrapper |
| @@ -75,31 +115,20 @@ def swap_layer_input_fboverlap_forward_wrapper(fn): | |||
| 75 | def swap_layer_input_fboverlap_1f1b_wrapper(fn): | 115 | def swap_layer_input_fboverlap_1f1b_wrapper(fn): |
| 76 | 116 | ||
| 77 | def wrapper(self, *args, **kwargs): | 117 | def wrapper(self, *args, **kwargs): |
| 78 | - if not hasattr(self, 'swap_manager'): | 118 | + _restore_swap_entry(_get_backward_layer_graph(args, kwargs, 'bwd_layer_graph', 3)) |
| 79 | - return fn(self, *args, **kwargs) | 119 | + _restore_swap_entry(_get_backward_layer_graph(args, kwargs, 'next_bwd_layer_graph', 5)) |
| 80 | 120 | ||
| 81 | - hidden_states = None | 121 | + swap_entry = _swap_out_layer_input(self, _get_hidden_states(args, kwargs)) |
| 82 | - if 'hidden_states' in kwargs: | 122 | + try: |
| 83 | - hidden_states = kwargs['hidden_states'] | 123 | + result = fn(self, *args, **kwargs) |
| 84 | - elif len(args) > 0 and isinstance(args[0], torch.Tensor): | 124 | + finally: |
| 85 | - hidden_states = args[0] | 125 | + if swap_entry is not None: |
| 86 | - | 126 | + self.swap_manager.wait_swap_out(swap_entry) |
| 87 | - backward_layer = None | ||
| 88 | - if 'bwd_layer_graph' in kwargs: | ||
| 89 | - backward_layer = kwargs['bwd_layer_graph'].layer | ||
| 90 | - elif len(args) > 3 and isinstance(args[3], LayerGraph): | ||
| 91 | - backward_layer = args[3].layer | ||
| 92 | - | ||
| 93 | - hidden_states.swap_this_tensor = True | ||
| 94 | - | ||
| 95 | - self.swap_manager.swap_out_tensors([hidden_states]) | ||
| 96 | - backward_layer.swap_manager.swap_in_prev_layer() | ||
| 97 | - | ||
| 98 | - result = fn(self, *args, **kwargs) | ||
| 99 | - | ||
| 100 | - self.swap_manager.wait_swap_out() | ||
| 101 | - backward_layer.swap_manager.wait_swap_in() | ||
| 102 | 127 | ||
| 128 | + if isinstance(result, tuple) and len(result) > 2 and swap_entry is not None: | ||
| 129 | + layer_graph = result[2] | ||
| 130 | + if layer_graph is not None and hasattr(layer_graph, 'layer') and hasattr(layer_graph, 'layer_input'): | ||
| 131 | + layer_graph.swap_layer_input_entry = swap_entry | ||
| 103 | return result | 132 | return result |
| 104 | 133 | ||
| 105 | return wrapper | 134 | return wrapper |
| @@ -108,13 +137,8 @@ def swap_layer_input_fboverlap_1f1b_wrapper(fn): | |||
| 108 | def swap_layer_input_fboverlap_backward_wrapper(fn): | 137 | def swap_layer_input_fboverlap_backward_wrapper(fn): |
| 109 | 138 | ||
| 110 | def wrapper(*args, **kwargs): | 139 | def wrapper(*args, **kwargs): |
| 111 | - self = args[1].layer | 140 | + layer_graph = args[1] if len(args) > 1 else kwargs.get('layer_graph') |
| 112 | - if not hasattr(self, 'swap_manager'): | 141 | + _restore_swap_entry(layer_graph) |
| 113 | - return fn(*args, **kwargs) | 142 | + return fn(*args, **kwargs) |
| 114 | - | ||
| 115 | - self.swap_manager.swap_in_prev_layer() | ||
| 116 | - result = fn(*args, **kwargs) | ||
| 117 | - self.swap_manager.wait_swap_in() | ||
| 118 | - return result | ||
| 119 | 143 | ||
| 120 | return wrapper | 144 | return wrapper |
| @@ -1,6 +1,6 @@ | |||
| 1 | # Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved. | 1 | # Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved. |
| 2 | 2 | ||
| 3 | -from typing import Optional, Callable, List | 3 | +from typing import Callable, List, Optional |
| 4 | 4 | ||
| 5 | import torch | 5 | import torch |
| 6 | import torch_npu | 6 | import torch_npu |
| @@ -210,16 +210,48 @@ class SwapLayerInputManager: | |||
| 210 | if cls._d2h_stream is None: | 210 | if cls._d2h_stream is None: |
| 211 | cls._d2h_stream = torch_npu.npu.Stream(device=torch.npu.current_device()) | 211 | cls._d2h_stream = torch_npu.npu.Stream(device=torch.npu.current_device()) |
| 212 | 212 | ||
| 213 | - def swap_out_tensors(self, tensors: List[torch.Tensor]): | 213 | + def swap_out_tensors(self, tensors: List[torch.Tensor]) -> Optional[SwapTensors]: |
| 214 | - """Add a tensor to the current micro-batch's swap collection.""" | 214 | + """Swap valid tensors to host and return their tracked swap entry.""" |
| 215 | - if self.layer_idx + 1 < self.num_layers: | 215 | + if self.layer_idx + 1 >= self.num_layers: |
| 216 | - swap_tensors = [t for t in tensors if is_valid_for_swap(t, self.custom_check_fn)] | 216 | + return None |
| 217 | - self.batch_stack.append(SwapTensors(swap_tensors)) | ||
| 218 | - self.batch_stack[-1].swap_to_host(SwapLayerInputManager._d2h_stream, async_op=True) | ||
| 219 | 217 | ||
| 220 | - def wait_swap_out(self): | 218 | + swap_tensors = [tensor for tensor in tensors if is_valid_for_swap(tensor, self.custom_check_fn)] |
| 221 | - if len(self.batch_stack) > 0: | 219 | + if not swap_tensors: |
| 222 | - self.batch_stack[-1].wait_d2h_and_release_device_tensor() | 220 | + return None |
| 221 | + | ||
| 222 | + swap_entry = SwapTensors(swap_tensors) | ||
| 223 | + self.batch_stack.append(swap_entry) | ||
| 224 | + swap_entry.swap_to_host(SwapLayerInputManager._d2h_stream, async_op=True) | ||
| 225 | + return swap_entry | ||
| 226 | + | ||
| 227 | + def wait_swap_out(self, swap_entry: Optional[SwapTensors] = None): | ||
| 228 | + """Wait for a specific swap-out, or the latest one when omitted.""" | ||
| 229 | + if swap_entry is None and self.batch_stack: | ||
| 230 | + swap_entry = self.batch_stack[-1] | ||
| 231 | + if swap_entry is not None: | ||
| 232 | + swap_entry.wait_d2h_and_release_device_tensor() | ||
| 233 | + | ||
| 234 | + def restore_swap_entry(self, swap_entry: Optional[SwapTensors]): | ||
| 235 | + """Restore one tracked entry to device and remove it from this manager.""" | ||
| 236 | + if swap_entry is None: | ||
| 237 | + return | ||
| 238 | + | ||
| 239 | + self._ensure_streams() | ||
| 240 | + if swap_entry.stat == "d2h": | ||
| 241 | + swap_entry.wait_d2h_and_release_device_tensor() | ||
| 242 | + if swap_entry.stat == "host": | ||
| 243 | + swap_entry.swap_to_device(self._h2d_stream, async_op=True) | ||
| 244 | + if swap_entry.stat == "h2d": | ||
| 245 | + swap_entry.wait_h2d_and_release_cpu_tensor() | ||
| 246 | + if swap_entry.stat == "device": | ||
| 247 | + self.remove_swap_entry(swap_entry) | ||
| 248 | + | ||
| 249 | + def remove_swap_entry(self, swap_entry: SwapTensors): | ||
| 250 | + """Remove an entry by identity without disturbing pipeline ordering.""" | ||
| 251 | + for index, tracked_entry in enumerate(self.batch_stack): | ||
| 252 | + if tracked_entry is swap_entry: | ||
| 253 | + self.batch_stack.pop(index) | ||
| 254 | + return | ||
| 223 | 255 | ||
| 224 | def swap_in_prev_layer(self): | 256 | def swap_in_prev_layer(self): |
| 225 | prev_manager = self._get_prev_layer_manager() | 257 | prev_manager = self._get_prev_layer_manager() |
| @@ -248,7 +280,7 @@ class SwapLayerInputManager: | |||
| 248 | swap_entry = prev_manager.batch_stack[-1] | 280 | swap_entry = prev_manager.batch_stack[-1] |
| 249 | swap_entry.wait_d2h_and_release_device_tensor() | 281 | swap_entry.wait_d2h_and_release_device_tensor() |
| 250 | 282 | ||
| 251 | - if self.layer_idx + 1 < self.num_layers: | 283 | + if self.layer_idx + 1 < self.num_layers and self.batch_stack: |
| 252 | self.batch_stack[-1].swap_to_host(SwapLayerInputManager._d2h_stream, async_op=True) | 284 | self.batch_stack[-1].swap_to_host(SwapLayerInputManager._d2h_stream, async_op=True) |
| 253 | 285 | ||
| 254 | def backward_hook(self, _): | 286 | def backward_hook(self, _): |
| @@ -3,7 +3,6 @@ | |||
| 3 | import torch | 3 | import torch |
| 4 | 4 | ||
| 5 | from megatron.training import get_args | 5 | from megatron.training import get_args |
| 6 | -from mindspeed.core.transformer.moe.comm_utils import async_all_to_all | ||
| 7 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | 6 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput |
| 8 | 7 | ||
| 9 | 8 | ||
| @@ -20,8 +19,6 @@ def get_async_alltoall_outputs(): | |||
| 20 | 19 | ||
| 21 | 20 | ||
| 22 | def launch_async_all2all(): | 21 | def launch_async_all2all(): |
| 23 | - global AsyncAll2All_INPUT | ||
| 24 | - global AsyncAll2All_OUTPUT | ||
| 25 | if len(AsyncAll2All_INPUT) > 0: | 22 | if len(AsyncAll2All_INPUT) > 0: |
| 26 | comm_fn, args, kwargs = AsyncAll2All_INPUT.pop(0) | 23 | comm_fn, args, kwargs = AsyncAll2All_INPUT.pop(0) |
| 27 | output, a2a_handle = comm_fn(*args, **kwargs) | 24 | output, a2a_handle = comm_fn(*args, **kwargs) |
| @@ -32,6 +29,71 @@ def launch_async_all2all_hook(_): | |||
| 32 | launch_async_all2all() | 29 | launch_async_all2all() |
| 33 | 30 | ||
| 34 | 31 | ||
| 32 | +def _discard_mhc_pre_recompute_output(mhc_module, hook_tensor): | ||
| 33 | + if not hasattr(mhc_module, 'discard_mhc_pre_ascend_output'): | ||
| 34 | + return | ||
| 35 | + if isinstance(hook_tensor, torch.Tensor) and hook_tensor.requires_grad: | ||
| 36 | + mhc_module.discard_mhc_pre_ascend_output(hook_tensor) | ||
| 37 | + | ||
| 38 | + | ||
| 39 | +def discard_mlp_mhc_pre_recompute_output(layer, hook_tensor): | ||
| 40 | + _discard_mhc_pre_recompute_output(getattr(layer, 'mlp_mhc', None), hook_tensor) | ||
| 41 | + | ||
| 42 | + | ||
| 43 | +def _should_defer_attention_recompute_for_mhc_post(layer): | ||
| 44 | + args = get_args() | ||
| 45 | + if not getattr(args, 'recompute_csa_attention', False): | ||
| 46 | + return False | ||
| 47 | + attn_mhc = getattr(layer, 'attn_mhc', None) | ||
| 48 | + should_recompute_post = getattr(attn_mhc, '_should_recompute_mhc_post_ascend', None) | ||
| 49 | + if callable(should_recompute_post): | ||
| 50 | + return should_recompute_post() | ||
| 51 | + return ( | ||
| 52 | + getattr(args, 'enable_mhc', False) | ||
| 53 | + and getattr(args, 'mhc_recompute', False) | ||
| 54 | + and getattr(args, 'use_fused_mhc', False) | ||
| 55 | + and getattr(layer, 'training', False) | ||
| 56 | + and not getattr(attn_mhc, 'is_mtp_layer', False) | ||
| 57 | + ) | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +def discard_attention_recompute_outputs_for_mhc_post(layer, hook_tensor): | ||
| 61 | + hook_tensor_has_grad = isinstance(hook_tensor, torch.Tensor) and hook_tensor.requires_grad | ||
| 62 | + attn_mhc = getattr(layer, 'attn_mhc', None) | ||
| 63 | + _discard_mhc_pre_recompute_output(attn_mhc, hook_tensor) | ||
| 64 | + | ||
| 65 | + defer_attention_recompute = getattr(layer, 'defer_attention_recompute_for_mhc_post', False) | ||
| 66 | + if defer_attention_recompute: | ||
| 67 | + norm_checkpoint = getattr(layer, 'norm_ckpt1', None) | ||
| 68 | + if getattr(layer, 'defer_attention_recompute_norm', False) and norm_checkpoint is not None: | ||
| 69 | + if hook_tensor_has_grad: | ||
| 70 | + norm_checkpoint.discard_output() | ||
| 71 | + hook_tensor.register_hook(norm_checkpoint.recompute) | ||
| 72 | + layer.norm_ckpt1 = None | ||
| 73 | + | ||
| 74 | + self_attention = getattr(layer, 'self_attention', None) | ||
| 75 | + if hook_tensor_has_grad and hasattr(self_attention, 'discard_csa_attention_output'): | ||
| 76 | + self_attention.discard_csa_attention_output(hook_tensor) | ||
| 77 | + | ||
| 78 | + bda_checkpoint = getattr(layer, 'self_attn_bda_checkpoint', None) | ||
| 79 | + if bda_checkpoint is not None: | ||
| 80 | + if hook_tensor_has_grad: | ||
| 81 | + bda_checkpoint.discard_output() | ||
| 82 | + hook_tensor.register_hook(bda_checkpoint.recompute) | ||
| 83 | + layer.self_attn_bda_checkpoint = None | ||
| 84 | + | ||
| 85 | + if hook_tensor_has_grad and hasattr(attn_mhc, 'discard_mhc_post_ascend_output'): | ||
| 86 | + attn_mhc.discard_mhc_post_ascend_output(hook_tensor) | ||
| 87 | + | ||
| 88 | + if defer_attention_recompute: | ||
| 89 | + layer.defer_attention_recompute_for_mhc_post = False | ||
| 90 | + layer.defer_attention_recompute_norm = False | ||
| 91 | + | ||
| 92 | + | ||
| 93 | +def _self_attn_bda_is_output_only(bda_module, attention_bias): | ||
| 94 | + return attention_bias is None and bda_module.__class__.__name__ == 'AddOpWithBias' | ||
| 95 | + | ||
| 96 | + | ||
| 35 | def attention_forward( | 97 | def attention_forward( |
| 36 | self, | 98 | self, |
| 37 | hidden_states, | 99 | hidden_states, |
| @@ -43,9 +105,13 @@ def attention_forward( | |||
| 43 | rotary_pos_sin=None, | 105 | rotary_pos_sin=None, |
| 44 | attention_bias=None, | 106 | attention_bias=None, |
| 45 | packed_seq_params=None, | 107 | packed_seq_params=None, |
| 46 | - recompute_norm=False | 108 | + recompute_norm=False, |
| 47 | ): | 109 | ): |
| 48 | args = get_args() | 110 | args = get_args() |
| 111 | + defer_attention_recompute = _should_defer_attention_recompute_for_mhc_post(self) | ||
| 112 | + self.self_attn_bda_checkpoint = None | ||
| 113 | + self.defer_attention_recompute_for_mhc_post = defer_attention_recompute | ||
| 114 | + self.defer_attention_recompute_norm = recompute_norm if defer_attention_recompute else False | ||
| 49 | if getattr(args, 'enable_mhc', False): | 115 | if getattr(args, 'enable_mhc', False): |
| 50 | # attn mHC pre | 116 | # attn mHC pre |
| 51 | post, comb = None, None | 117 | post, comb = None, None |
| @@ -79,24 +145,45 @@ def attention_forward( | |||
| 79 | packed_seq_params=packed_seq_params, | 145 | packed_seq_params=packed_seq_params, |
| 80 | ) | 146 | ) |
| 81 | 147 | ||
| 82 | - # TODO: could we move `bias_dropout_add_exec_handler` itself | 148 | + # NOTE: `bias_dropout_add_exec_handler` could potentially be moved |
| 83 | # inside the module provided in the `bias_dropout_add_spec` module? | 149 | # inside the module provided in the `bias_dropout_add_spec` module? |
| 84 | - with self.bias_dropout_add_exec_handler(): | 150 | + bda_module = self.self_attn_bda(self.training, self.config.bias_dropout_fusion) |
| 85 | - hidden_states = self.self_attn_bda(self.training, self.config.bias_dropout_fusion)( | 151 | + if ( |
| 86 | - attention_output_with_bias, residual, self.hidden_dropout | 152 | + defer_attention_recompute |
| 87 | - ) | 153 | + and isinstance(attention_output_with_bias, tuple) |
| 154 | + and len(attention_output_with_bias) == 2 | ||
| 155 | + and isinstance(attention_output_with_bias[0], torch.Tensor) | ||
| 156 | + and not _self_attn_bda_is_output_only(bda_module, attention_output_with_bias[1]) | ||
| 157 | + ): | ||
| 158 | + attention_output, attention_bias = attention_output_with_bias | ||
| 88 | 159 | ||
| 89 | - if recompute_norm: | 160 | + def self_attn_bda_without_bias(attention_output, residual): |
| 161 | + with self.bias_dropout_add_exec_handler(): | ||
| 162 | + return bda_module((attention_output, None), residual, self.hidden_dropout) | ||
| 163 | + | ||
| 164 | + def self_attn_bda_with_bias(attention_output, attention_bias, residual): | ||
| 165 | + with self.bias_dropout_add_exec_handler(): | ||
| 166 | + return bda_module((attention_output, attention_bias), residual, self.hidden_dropout) | ||
| 167 | + | ||
| 168 | + self.self_attn_bda_checkpoint = CheckpointWithoutOutput() | ||
| 169 | + if attention_bias is None: | ||
| 170 | + hidden_states = self.self_attn_bda_checkpoint.checkpoint( | ||
| 171 | + self_attn_bda_without_bias, False, attention_output, residual | ||
| 172 | + ) | ||
| 173 | + else: | ||
| 174 | + hidden_states = self.self_attn_bda_checkpoint.checkpoint( | ||
| 175 | + self_attn_bda_with_bias, False, attention_output, attention_bias, residual | ||
| 176 | + ) | ||
| 177 | + else: | ||
| 178 | + with self.bias_dropout_add_exec_handler(): | ||
| 179 | + hidden_states = bda_module(attention_output_with_bias, residual, self.hidden_dropout) | ||
| 180 | + | ||
| 181 | + if recompute_norm and not defer_attention_recompute: | ||
| 90 | self.norm_ckpt1.discard_output() | 182 | self.norm_ckpt1.discard_output() |
| 91 | hidden_states.register_hook(self.norm_ckpt1.recompute) | 183 | hidden_states.register_hook(self.norm_ckpt1.recompute) |
| 92 | 184 | ||
| 93 | if getattr(args, 'enable_mhc', False): | 185 | if getattr(args, 'enable_mhc', False): |
| 94 | # attn mHC post | 186 | # attn mHC post |
| 95 | - hidden_states = self.attn_mhc(hidden_states, | 187 | + hidden_states = self.attn_mhc(hidden_states, mhc_stage='post', residual=residual, post=post, comb=comb) |
| 96 | - mhc_stage='post', | ||
| 97 | - residual=residual, | ||
| 98 | - post=post, | ||
| 99 | - comb=comb | ||
| 100 | - ) | ||
| 101 | 188 | ||
| 102 | return hidden_states | 189 | return hidden_states |
| @@ -72,6 +72,7 @@ class NoopLayerGraph: | |||
| 72 | self.checkpointed = checkpointed | 72 | self.checkpointed = checkpointed |
| 73 | self.is_moe_layer = False | 73 | self.is_moe_layer = False |
| 74 | self.layer = layer | 74 | self.layer = layer |
| 75 | + self.swap_layer_input_entry = None | ||
| 75 | 76 | ||
| 76 | def record_layer_inputs(self, *args): | 77 | def record_layer_inputs(self, *args): |
| 77 | self.layer_inputs = args # pylint: disable=attribute-defined-outside-init | 78 | self.layer_inputs = args # pylint: disable=attribute-defined-outside-init |
| @@ -134,6 +135,8 @@ class LayerGraph: | |||
| 134 | # For selective recompute | 135 | # For selective recompute |
| 135 | self.act_ckpt_manager = None | 136 | self.act_ckpt_manager = None |
| 136 | self.remote_hot_act_ckpt_manager = None | 137 | self.remote_hot_act_ckpt_manager = None |
| 138 | + # Exact layer-input swap entry associated with this forward graph. | ||
| 139 | + self.swap_layer_input_entry = None | ||
| 137 | 140 | ||
| 138 | def record_layer_inputs(self, *args): | 141 | def record_layer_inputs(self, *args): |
| 139 | self.layer_inputs = args # pylint: disable=attribute-defined-outside-init | 142 | self.layer_inputs = args # pylint: disable=attribute-defined-outside-init |
| @@ -10,7 +10,11 @@ from mindspeed.core.memory.recompute.activation.activation_recompute_forward imp | |||
| 10 | DENSE_FC1_MEMORY_CONTEXT, | 10 | DENSE_FC1_MEMORY_CONTEXT, |
| 11 | ) | 11 | ) |
| 12 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | 12 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput |
| 13 | -from ..modules.attention import attention_forward | 13 | +from ..modules.attention import ( |
| 14 | + attention_forward, | ||
| 15 | + discard_attention_recompute_outputs_for_mhc_post, | ||
| 16 | + discard_mlp_mhc_pre_recompute_output, | ||
| 17 | +) | ||
| 14 | from ..modules.utils import ( | 18 | from ..modules.utils import ( |
| 15 | detach_tensor, | 19 | detach_tensor, |
| 16 | NoopLayerGraph, | 20 | NoopLayerGraph, |
| @@ -243,6 +247,8 @@ def transformer_layer_forward_moe( | |||
| 243 | post=detached_mlp_mhc_post, | 247 | post=detached_mlp_mhc_post, |
| 244 | comb=detached_mlp_mhc_comb, | 248 | comb=detached_mlp_mhc_comb, |
| 245 | ) | 249 | ) |
| 250 | + discard_attention_recompute_outputs_for_mhc_post(self, mlp_mhc_output) | ||
| 251 | + discard_mlp_mhc_pre_recompute_output(self, mlp_mhc_output) | ||
| 246 | 252 | ||
| 247 | saved_tensors = [ | 253 | saved_tensors = [ |
| 248 | (attention_out, detached_attention_out), | 254 | (attention_out, detached_attention_out), |
| @@ -380,6 +386,7 @@ def transformer_layer_forward_dense( | |||
| 380 | # p2p_communication), it serves to document the origin of this | 386 | # p2p_communication), it serves to document the origin of this |
| 381 | # 'view' tensor. | 387 | # 'view' tensor. |
| 382 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) | 388 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) |
| 389 | + discard_attention_recompute_outputs_for_mhc_post(self, output) | ||
| 383 | 390 | ||
| 384 | saved_tensors = ( | 391 | saved_tensors = ( |
| 385 | (attention_graph, detached_attention_out), | 392 | (attention_graph, detached_attention_out), |
| @@ -13,7 +13,13 @@ from mindspeed.args_utils import get_full_args | |||
| 13 | from mindspeed.core.transformer.moe.comm_utils import async_all_to_all | 13 | from mindspeed.core.transformer.moe.comm_utils import async_all_to_all |
| 14 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | 14 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput |
| 15 | from ..modules.weight_grad_store import WeightGradStore | 15 | from ..modules.weight_grad_store import WeightGradStore |
| 16 | -from ..modules.attention import attention_forward, set_async_alltoall_inputs, get_async_alltoall_outputs | 16 | +from ..modules.attention import ( |
| 17 | + attention_forward, | ||
| 18 | + set_async_alltoall_inputs, | ||
| 19 | + get_async_alltoall_outputs, | ||
| 20 | + discard_attention_recompute_outputs_for_mhc_post, | ||
| 21 | + discard_mlp_mhc_pre_recompute_output, | ||
| 22 | +) | ||
| 17 | from ..modules.utils import ( | 23 | from ..modules.utils import ( |
| 18 | detach_tensor, | 24 | detach_tensor, |
| 19 | run_graph_backward, | 25 | run_graph_backward, |
| @@ -322,6 +328,7 @@ def transformer_layer_forward_dense_backward_moe_overlaping( | |||
| 322 | # p2p_communication), it serves to document the origin of this | 328 | # p2p_communication), it serves to document the origin of this |
| 323 | # 'view' tensor. | 329 | # 'view' tensor. |
| 324 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) | 330 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) |
| 331 | + discard_attention_recompute_outputs_for_mhc_post(fwd_layer, output) | ||
| 325 | 332 | ||
| 326 | # handle fwd p2p communication | 333 | # handle fwd p2p communication |
| 327 | next_iter_input_tensor, fwd_p2p_handles = None, None | 334 | next_iter_input_tensor, fwd_p2p_handles = None, None |
| @@ -622,6 +629,7 @@ def transformer_layer_forward_moe_backward_dense_overlaping( | |||
| 622 | # p2p_communication), it serves to document the origin of this | 629 | # p2p_communication), it serves to document the origin of this |
| 623 | # 'view' tensor. | 630 | # 'view' tensor. |
| 624 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) | 631 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) |
| 632 | + discard_attention_recompute_outputs_for_mhc_post(fwd_layer, output) | ||
| 625 | 633 | ||
| 626 | # handle fwd p2p communication | 634 | # handle fwd p2p communication |
| 627 | next_iter_input_tensor, fwd_p2p_handles = None, None | 635 | next_iter_input_tensor, fwd_p2p_handles = None, None |
| @@ -770,6 +778,7 @@ def transformer_layer_forward_dense_backward_dense_overlaping( | |||
| 770 | # p2p_communication), it serves to document the origin of this | 778 | # p2p_communication), it serves to document the origin of this |
| 771 | # 'view' tensor. | 779 | # 'view' tensor. |
| 772 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) | 780 | output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) |
| 781 | + discard_attention_recompute_outputs_for_mhc_post(fwd_layer, output) | ||
| 773 | 782 | ||
| 774 | # handle fwd p2p communication | 783 | # handle fwd p2p communication |
| 775 | next_iter_input_tensor, fwd_p2p_handles = None, None | 784 | next_iter_input_tensor, fwd_p2p_handles = None, None |
| @@ -1308,6 +1317,8 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 1308 | post=detached_mlp_mhc_post, | 1317 | post=detached_mlp_mhc_post, |
| 1309 | comb=detached_mlp_mhc_comb, | 1318 | comb=detached_mlp_mhc_comb, |
| 1310 | ) | 1319 | ) |
| 1320 | + discard_attention_recompute_outputs_for_mhc_post(fwd_layer, mlp_mhc_output) | ||
| 1321 | + discard_mlp_mhc_pre_recompute_output(fwd_layer, mlp_mhc_output) | ||
| 1311 | 1322 | ||
| 1312 | # handle fwd p2p communication | 1323 | # handle fwd p2p communication |
| 1313 | next_iter_input_tensor, fwd_p2p_handles = None, None | 1324 | next_iter_input_tensor, fwd_p2p_handles = None, None |
| @@ -2,13 +2,13 @@ | |||
| 2 | # Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved. | 2 | # Copyright (c) 2025, Huawei Technologies Co., Ltd. All rights reserved. |
| 3 | 3 | ||
| 4 | from functools import wraps | 4 | from functools import wraps |
| 5 | -from typing import Optional | ||
| 6 | import torch | 5 | import torch |
| 7 | -import torch_npu | ||
| 8 | import torch.nn.functional as F | 6 | import torch.nn.functional as F |
| 9 | -from mindspeed.core.transformer.moe.moe_feature import ( | 7 | +import torch_npu |
| 10 | - parallel_state, MLP, build_module, TransformerConfig, MLPSubmodules, TransformerConfig) | 8 | +from typing import Optional |
| 9 | +from mindspeed.core.transformer.moe.moe_feature import MLP, build_module, MLPSubmodules, TransformerConfig | ||
| 11 | from mindspeed.model.transformer import should_recompute_activation | 10 | from mindspeed.model.transformer import should_recompute_activation |
| 11 | +from mindspeed.core.memory.recompute.recompute_common import get_recompute_priority | ||
| 12 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | 12 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput |
| 13 | from mindspeed.core.fusions.fused_bias_swiglu import fused_swiglu | 13 | from mindspeed.core.fusions.fused_bias_swiglu import fused_swiglu |
| 14 | 14 | ||
| @@ -29,11 +29,11 @@ def mlp_init( | |||
| 29 | submodules: MLPSubmodules, | 29 | submodules: MLPSubmodules, |
| 30 | is_expert: bool = False, | 30 | is_expert: bool = False, |
| 31 | input_size: int = None, | 31 | input_size: int = None, |
| 32 | - with_shared_expert=False | 32 | + with_shared_expert=False, |
| 33 | ): | 33 | ): |
| 34 | """ | 34 | """ |
| 35 | Shared expert MLP init with Moe_overlap. | 35 | Shared expert MLP init with Moe_overlap. |
| 36 | - In 0.10.0, the definition of shared_experts has conflict. | 36 | + In 0.10.0, the definition of shared_experts has conflict. |
| 37 | Rename the MindSpeed version to 'with_shared_expert'. | 37 | Rename the MindSpeed version to 'with_shared_expert'. |
| 38 | """ | 38 | """ |
| 39 | super(MLP, self).__init__(config=config) | 39 | super(MLP, self).__init__(config=config) |
| @@ -57,7 +57,7 @@ def mlp_init( | |||
| 57 | skip_bias_add=True, | 57 | skip_bias_add=True, |
| 58 | is_expert=is_expert, | 58 | is_expert=is_expert, |
| 59 | tp_comm_buffer_name='fc1', | 59 | tp_comm_buffer_name='fc1', |
| 60 | - with_shared_expert=with_shared_expert | 60 | + with_shared_expert=with_shared_expert, |
| 61 | ) | 61 | ) |
| 62 | else: | 62 | else: |
| 63 | self.linear_fc1 = build_module( | 63 | self.linear_fc1 = build_module( |
| @@ -70,7 +70,7 @@ def mlp_init( | |||
| 70 | bias=self.config.add_bias_linear, | 70 | bias=self.config.add_bias_linear, |
| 71 | skip_bias_add=True, | 71 | skip_bias_add=True, |
| 72 | is_expert=is_expert, | 72 | is_expert=is_expert, |
| 73 | - tp_comm_buffer_name='fc1' | 73 | + tp_comm_buffer_name='fc1', |
| 74 | ) | 74 | ) |
| 75 | 75 | ||
| 76 | self.activation_func = self.config.activation_func | 76 | self.activation_func = self.config.activation_func |
| @@ -87,7 +87,7 @@ def mlp_init( | |||
| 87 | skip_bias_add=True, | 87 | skip_bias_add=True, |
| 88 | is_expert=is_expert, | 88 | is_expert=is_expert, |
| 89 | tp_comm_buffer_name='fc2', | 89 | tp_comm_buffer_name='fc2', |
| 90 | - with_shared_expert=with_shared_expert | 90 | + with_shared_expert=with_shared_expert, |
| 91 | ) | 91 | ) |
| 92 | else: | 92 | else: |
| 93 | self.linear_fc2 = build_module( | 93 | self.linear_fc2 = build_module( |
| @@ -100,7 +100,7 @@ def mlp_init( | |||
| 100 | input_is_parallel=True, | 100 | input_is_parallel=True, |
| 101 | skip_bias_add=True, | 101 | skip_bias_add=True, |
| 102 | is_expert=is_expert, | 102 | is_expert=is_expert, |
| 103 | - tp_comm_buffer_name='fc2' | 103 | + tp_comm_buffer_name='fc2', |
| 104 | ) | 104 | ) |
| 105 | 105 | ||
| 106 | self.with_shared_expert = with_shared_expert | 106 | self.with_shared_expert = with_shared_expert |
| @@ -110,6 +110,7 @@ def core_mlp_forward_wrapper(fn): | |||
| 110 | """ | 110 | """ |
| 111 | A wrapper about setting args for zero_memory&recompute in MLP. | 111 | A wrapper about setting args for zero_memory&recompute in MLP. |
| 112 | """ | 112 | """ |
| 113 | + | ||
| 113 | 114 | ||
| 114 | def wrapper(self, *args, **kwargs): | 115 | def wrapper(self, *args, **kwargs): |
| 115 | if isinstance(args, tuple): | 116 | if isinstance(args, tuple): |
| @@ -117,12 +118,14 @@ def core_mlp_forward_wrapper(fn): | |||
| 117 | 118 | ||
| 118 | if getattr(self.config, 'profile', False) and not self.config.num_experts: | 119 | if getattr(self.config, 'profile', False) and not self.config.num_experts: |
| 119 | from mindspeed.auto_settings.module.black.patch.hccl_operator import MOEOrMLPStartOp, MOEOrMLPEndOp | 120 | from mindspeed.auto_settings.module.black.patch.hccl_operator import MOEOrMLPStartOp, MOEOrMLPEndOp |
| 121 | + | ||
| 120 | args[0] = MOEOrMLPStartOp.apply(args[0]) | 122 | args[0] = MOEOrMLPStartOp.apply(args[0]) |
| 121 | activation_func_1 = torch.nn.Softplus() | 123 | activation_func_1 = torch.nn.Softplus() |
| 122 | args[0] = activation_func_1(args[0]) | 124 | args[0] = activation_func_1(args[0]) |
| 123 | 125 | ||
| 124 | self.layer_number = getattr(self, "layer_number", None) | 126 | self.layer_number = getattr(self, "layer_number", None) |
| 125 | is_recompute_activation = should_recompute_activation(self.layer_number) | 127 | is_recompute_activation = should_recompute_activation(self.layer_number) |
| 128 | + moe_ctx = None | ||
| 126 | if self.config.moe_alltoall_overlap_comm and not isinstance(args[-1], torch.Tensor): | 129 | if self.config.moe_alltoall_overlap_comm and not isinstance(args[-1], torch.Tensor): |
| 127 | moe_ctx = args[-1] | 130 | moe_ctx = args[-1] |
| 128 | args = args[:-1] | 131 | args = args[:-1] |
| @@ -132,7 +135,7 @@ def core_mlp_forward_wrapper(fn): | |||
| 132 | if bias is not None: | 135 | if bias is not None: |
| 133 | intermediate = intermediate + bias | 136 | intermediate = intermediate + bias |
| 134 | if self.config.gated_linear_unit: | 137 | if self.config.gated_linear_unit: |
| 135 | - assert (self.config.activation_func == F.silu), 'Activation function must be silu when using fused_swiglu' | 138 | + assert self.config.activation_func == F.silu, 'Activation function must be silu when using fused_swiglu' |
| 136 | if not hasattr(self, 'origin_activation_func'): | 139 | if not hasattr(self, 'origin_activation_func'): |
| 137 | self.origin_activation_func = self.activation_func | 140 | self.origin_activation_func = self.activation_func |
| 138 | self.activation_func = fused_swiglu | 141 | self.activation_func = fused_swiglu |
| @@ -150,6 +153,8 @@ def core_mlp_forward_wrapper(fn): | |||
| 150 | elif moe_zero_memory == "level1" and not only_recompute_activation(self.config, layer_number=self.layer_number): | 153 | elif moe_zero_memory == "level1" and not only_recompute_activation(self.config, layer_number=self.layer_number): |
| 151 | # Only for zm1 in alltoall_seq dispatcher. | 154 | # Only for zm1 in alltoall_seq dispatcher. |
| 152 | if self.with_shared_expert: | 155 | if self.with_shared_expert: |
| 156 | + if moe_ctx is None: | ||
| 157 | + raise RuntimeError("MoE context is required when shared-expert zero-memory recompute is enabled.") | ||
| 153 | self.activation_function = activation_function | 158 | self.activation_function = activation_function |
| 154 | hidden_states = args[0] | 159 | hidden_states = args[0] |
| 155 | fc1_out_parallel, bias_parallel = self.linear_fc1(hidden_states) | 160 | fc1_out_parallel, bias_parallel = self.linear_fc1(hidden_states) |
| @@ -165,10 +170,9 @@ def core_mlp_forward_wrapper(fn): | |||
| 165 | hidden_states = args[0] | 170 | hidden_states = args[0] |
| 166 | intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states) | 171 | intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states) |
| 167 | self.activation_checkpoint_manager = CheckpointWithoutOutput() | 172 | self.activation_checkpoint_manager = CheckpointWithoutOutput() |
| 168 | - intermediate_parallel = self.activation_checkpoint_manager.checkpoint(activation_function, | 173 | + intermediate_parallel = self.activation_checkpoint_manager.checkpoint( |
| 169 | - False, | 174 | + activation_function, False, intermediate_parallel, bias_parallel |
| 170 | - intermediate_parallel, | 175 | + ) |
| 171 | - bias_parallel) | ||
| 172 | # [s, b, h] | 176 | # [s, b, h] |
| 173 | output, output_bias = self.linear_fc2(intermediate_parallel) | 177 | output, output_bias = self.linear_fc2(intermediate_parallel) |
| 174 | 178 | ||
| @@ -187,6 +191,7 @@ def core_mlp_forward_wrapper(fn): | |||
| 187 | output = MOEOrMLPEndOp.apply(output) | 191 | output = MOEOrMLPEndOp.apply(output) |
| 188 | 192 | ||
| 189 | return output, output_bias | 193 | return output, output_bias |
| 194 | + | ||
| 190 | return wrapper | 195 | return wrapper |
| 191 | 196 | ||
| 192 | 197 | ||
| @@ -195,6 +200,7 @@ def parallel_transformer_layer_init_wrapper(fn): | |||
| 195 | def wrapper(self, *args, **kwargs): | 200 | def wrapper(self, *args, **kwargs): |
| 196 | fn(self, *args, **kwargs) | 201 | fn(self, *args, **kwargs) |
| 197 | from megatron.core.transformer.moe.moe_layer import MoELayer | 202 | from megatron.core.transformer.moe.moe_layer import MoELayer |
| 203 | + | ||
| 198 | if self.config.moe_alltoall_overlap_comm or self.config.moe_allgather_overlap_comm: | 204 | if self.config.moe_alltoall_overlap_comm or self.config.moe_allgather_overlap_comm: |
| 199 | if self.mlp.__class__ is MoELayer: | 205 | if self.mlp.__class__ is MoELayer: |
| 200 | self.mlp.experts.layer_number = self.layer_number | 206 | self.mlp.experts.layer_number = self.layer_number |
| @@ -202,6 +208,7 @@ def parallel_transformer_layer_init_wrapper(fn): | |||
| 202 | self.mlp.shared_experts.layer_number = self.layer_number | 208 | self.mlp.shared_experts.layer_number = self.layer_number |
| 203 | else: | 209 | else: |
| 204 | self.mlp.layer_number = self.layer_number | 210 | self.mlp.layer_number = self.layer_number |
| 211 | + | ||
| 205 | return wrapper | 212 | return wrapper |
| 206 | 213 | ||
| 207 | 214 | ||
| @@ -304,22 +311,7 @@ def get_all2all_experts_output(): | |||
| 304 | 311 | ||
| 305 | 312 | ||
| 306 | def only_recompute_activation(config, layer_number): | 313 | def only_recompute_activation(config, layer_number): |
| 307 | - | 314 | + recompute_priority = get_recompute_priority(config, layer_number, enable_per_pp_rank=True) |
| 308 | - vpp_rank = parallel_state.get_virtual_pipeline_model_parallel_rank() | ||
| 309 | - vpp_size = config.virtual_pipeline_model_parallel_size | ||
| 310 | - pp_size = config.pipeline_model_parallel_size | ||
| 311 | - | ||
| 312 | - if vpp_size is not None: | ||
| 313 | - layer_per_chunk = config.num_layers_per_virtual_pipeline_stage | ||
| 314 | - elif pp_size is not None: | ||
| 315 | - layer_per_chunk = config.num_layers // pp_size | ||
| 316 | - else: | ||
| 317 | - layer_per_chunk = config.num_layers | ||
| 318 | - | ||
| 319 | - vpp_rank = vpp_rank or 0 | ||
| 320 | - vpp_size = vpp_size or 1 | ||
| 321 | - | ||
| 322 | - recompute_priority = ((layer_number - 1) % layer_per_chunk) * vpp_size + vpp_rank | ||
| 323 | moe_zero_memory_num_layers = config.moe_zero_memory_num_layers | 315 | moe_zero_memory_num_layers = config.moe_zero_memory_num_layers |
| 324 | 316 | ||
| 325 | if moe_zero_memory_num_layers: | 317 | if moe_zero_memory_num_layers: |
| @@ -328,7 +320,7 @@ def only_recompute_activation(config, layer_number): | |||
| 328 | else: | 320 | else: |
| 329 | return True | 321 | return True |
| 330 | else: | 322 | else: |
| 331 | - return False | 323 | + return False |
| 332 | 324 | ||
| 333 | 325 | ||
| 334 | def forward_func(func, inputs): | 326 | def forward_func(func, inputs): |
| @@ -375,19 +367,19 @@ def backward_func(func_tensor, gradinputs): | |||
| 375 | 367 | ||
| 376 | 368 | ||
| 377 | def async_comm_sort_chunks_by_idxs( | 369 | def async_comm_sort_chunks_by_idxs( |
| 378 | - input: torch.Tensor, | 370 | + input_tensor: torch.Tensor, |
| 379 | split_sizes: torch.Tensor, | 371 | split_sizes: torch.Tensor, |
| 380 | sorted_idxs: torch.Tensor, | 372 | sorted_idxs: torch.Tensor, |
| 381 | probs: Optional[torch.Tensor] = None, | 373 | probs: Optional[torch.Tensor] = None, |
| 382 | fused: bool = False, | 374 | fused: bool = False, |
| 383 | - prob_handle = None | 375 | + prob_handle=None, |
| 384 | ): | 376 | ): |
| 385 | """Split and sort the input tensor based on the split_sizes and sorted indices.""" | 377 | """Split and sort the input tensor based on the split_sizes and sorted indices.""" |
| 386 | if fused: | 378 | if fused: |
| 387 | raise AssertionError('async sort_chunks_by_idxs not support fused now.') | 379 | raise AssertionError('async sort_chunks_by_idxs not support fused now.') |
| 388 | 380 | ||
| 389 | - input = torch.split(input, split_sizes.tolist(), dim=0) | 381 | + input_tensor = torch.split(input_tensor, split_sizes.tolist(), dim=0) |
| 390 | - output = torch.cat([input[i] for i in sorted_idxs.tolist()], dim=0) | 382 | + output = torch.cat([input_tensor[i] for i in sorted_idxs.tolist()], dim=0) |
| 391 | if probs is not None: | 383 | if probs is not None: |
| 392 | if prob_handle: | 384 | if prob_handle: |
| 393 | prob_handle.wait() | 385 | prob_handle.wait() |
| @@ -395,4 +387,4 @@ def async_comm_sort_chunks_by_idxs( | |||
| 395 | permuted_probs = torch.cat([probs[i] for i in sorted_idxs.tolist()], dim=0) | 387 | permuted_probs = torch.cat([probs[i] for i in sorted_idxs.tolist()], dim=0) |
| 396 | else: | 388 | else: |
| 397 | permuted_probs = None | 389 | permuted_probs = None |
| 398 | - return output, permuted_probs | 390 | + return output, permuted_probs |
| @@ -5,6 +5,7 @@ import torch | |||
| 5 | import torch_npu | 5 | import torch_npu |
| 6 | 6 | ||
| 7 | from mindspeed.args_utils import get_full_args as get_args | 7 | from mindspeed.args_utils import get_full_args as get_args |
| 8 | +from mindspeed.core.memory.recompute.recompute_common import get_recompute_priority | ||
| 8 | from megatron.core import parallel_state | 9 | from megatron.core import parallel_state |
| 9 | from megatron.core.transformer.moe.moe_utils import ( | 10 | from megatron.core.transformer.moe.moe_utils import ( |
| 10 | reduce_aux_losses_tracker_across_ranks, | 11 | reduce_aux_losses_tracker_across_ranks, |
| @@ -125,22 +126,7 @@ def get_all2all_experts_output(): | |||
| 125 | 126 | ||
| 126 | def only_recompute_activation(layer_number): | 127 | def only_recompute_activation(layer_number): |
| 127 | args = get_args() | 128 | args = get_args() |
| 128 | - vpp_rank = parallel_state.get_virtual_pipeline_model_parallel_rank() | 129 | + recompute_priority = get_recompute_priority(args, layer_number, enable_per_pp_rank=True) |
| 129 | - vpp_size = args.virtual_pipeline_model_parallel_size | ||
| 130 | - pp_size = args.pipeline_model_parallel_size | ||
| 131 | - | ||
| 132 | - if vpp_size is not None: | ||
| 133 | - layer_per_chunk = args.num_layers_per_virtual_pipeline_stage | ||
| 134 | - elif pp_size is not None: | ||
| 135 | - layer_per_chunk = args.num_layers // pp_size | ||
| 136 | - else: | ||
| 137 | - layer_per_chunk = args.num_layers | ||
| 138 | - | ||
| 139 | - if vpp_rank is None: | ||
| 140 | - vpp_rank = 0 | ||
| 141 | - if vpp_size is None: | ||
| 142 | - vpp_size = 1 | ||
| 143 | - recompute_priority = ((layer_number - 1) % layer_per_chunk) * vpp_size + vpp_rank | ||
| 144 | moe_zero_memory_num_layers = args.moe_zero_memory_num_layers | 130 | moe_zero_memory_num_layers = args.moe_zero_memory_num_layers |
| 145 | 131 | ||
| 146 | if moe_zero_memory_num_layers: | 132 | if moe_zero_memory_num_layers: |
| @@ -1,24 +1,17 @@ | |||
| 1 | # Copyright (c) 2023, NVIDIA CORPORATION. All rights reversed. | 1 | # Copyright (c) 2023, NVIDIA CORPORATION. All rights reversed. |
| 2 | # Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. | 2 | # Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved. |
| 3 | -import enum | ||
| 4 | -import os | ||
| 5 | from functools import wraps | 3 | from functools import wraps |
| 6 | from typing import Optional | 4 | from typing import Optional |
| 7 | 5 | ||
| 8 | -from contextlib import nullcontext | ||
| 9 | import torch | 6 | import torch |
| 10 | -import torch_npu | ||
| 11 | import torch.nn.functional as F | 7 | import torch.nn.functional as F |
| 12 | 8 | ||
| 13 | -from megatron import core | ||
| 14 | from megatron.training import get_args | 9 | from megatron.training import get_args |
| 15 | -from megatron.core.num_microbatches_calculator import get_num_microbatches | 10 | +from megatron.core import tensor_parallel, mpu |
| 16 | -from megatron.core import tensor_parallel, parallel_state, mpu | ||
| 17 | from megatron.core.utils import make_viewless_tensor | 11 | from megatron.core.utils import make_viewless_tensor |
| 18 | from megatron.core.transformer.identity_op import IdentityOp | 12 | from megatron.core.transformer.identity_op import IdentityOp |
| 19 | -from megatron.legacy.model.transformer import bias_dropout_add_fused_train, get_bias_dropout_add, bias_dropout_add_fused_inference | ||
| 20 | -from megatron.legacy.model.enums import AttnMaskType, LayerType, AttnType | ||
| 21 | from mindspeed.model.transformer import should_recompute_activation | 13 | from mindspeed.model.transformer import should_recompute_activation |
| 14 | +from mindspeed.core.memory.recompute.recompute_common import get_recompute_priority | ||
| 22 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | 15 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput |
| 23 | from mindspeed.core.fusions.fused_bias_swiglu import fused_swiglu | 16 | from mindspeed.core.fusions.fused_bias_swiglu import fused_swiglu |
| 24 | from mindspeed.core.transformer.moe.moe_utils import only_recompute_activation | 17 | from mindspeed.core.transformer.moe.moe_utils import only_recompute_activation |
| @@ -30,6 +23,7 @@ def parallel_transformer_layer_init_wrapper(fn): | |||
| 30 | def wrapper(self, *args, **kwargs): | 23 | def wrapper(self, *args, **kwargs): |
| 31 | from megatron.core.transformer.moe.moe_layer import MoELayer | 24 | from megatron.core.transformer.moe.moe_layer import MoELayer |
| 32 | from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP | 25 | from megatron.core.transformer.moe.experts import GroupedMLP, SequentialMLP |
| 26 | + | ||
| 33 | fn(self, *args, **kwargs) | 27 | fn(self, *args, **kwargs) |
| 34 | if self.mlp.__class__ is MoELayer: | 28 | if self.mlp.__class__ is MoELayer: |
| 35 | if self.mlp.experts.__class__ is GroupedMLP: | 29 | if self.mlp.experts.__class__ is GroupedMLP: |
| @@ -59,9 +53,9 @@ def parallel_transformer_checkpointed_forward_wrapper(forward_func): | |||
| 59 | return row_parallel_forward | 53 | return row_parallel_forward |
| 60 | 54 | ||
| 61 | 55 | ||
| 62 | -def parallel_transformer_checkpointed_forward(self, hidden_states, attention_mask, | 56 | +def parallel_transformer_checkpointed_forward( |
| 63 | - encoder_output, enc_dec_attn_mask, | 57 | + self, hidden_states, attention_mask, encoder_output, enc_dec_attn_mask, rotary_pos_emb, is_first_microbatch |
| 64 | - rotary_pos_emb, is_first_microbatch): | 58 | +): |
| 65 | """Forward method with activation checkpointing.""" | 59 | """Forward method with activation checkpointing.""" |
| 66 | 60 | ||
| 67 | def custom(start, end): | 61 | def custom(start, end): |
| @@ -81,33 +75,41 @@ def parallel_transformer_checkpointed_forward(self, hidden_states, attention_mas | |||
| 81 | # checkpoint the input activation of each divided chunk. | 75 | # checkpoint the input activation of each divided chunk. |
| 82 | # A method to further reduce memory usage reducing checkpoints. | 76 | # A method to further reduce memory usage reducing checkpoints. |
| 83 | if not global_args.swap_attention: | 77 | if not global_args.swap_attention: |
| 84 | - l = 0 | 78 | + layer_idx = 0 |
| 85 | - while l < num_layers_per_pipeline_rank: | 79 | + while layer_idx < num_layers_per_pipeline_rank: |
| 86 | hidden_states = tensor_parallel.checkpoint( | 80 | hidden_states = tensor_parallel.checkpoint( |
| 87 | - custom(l, l + self.recompute_num_layers), | 81 | + custom(layer_idx, layer_idx + self.recompute_num_layers), |
| 88 | self.distribute_saved_activations, | 82 | self.distribute_saved_activations, |
| 89 | - hidden_states, attention_mask, | 83 | + hidden_states, |
| 90 | - encoder_output, enc_dec_attn_mask, | 84 | + attention_mask, |
| 91 | - None, None, None, None, rotary_pos_emb) | 85 | + encoder_output, |
| 86 | + enc_dec_attn_mask, | ||
| 87 | + None, | ||
| 88 | + None, | ||
| 89 | + None, | ||
| 90 | + None, | ||
| 91 | + rotary_pos_emb, | ||
| 92 | + ) | ||
| 92 | 93 | ||
| 93 | - l += self.recompute_num_layers | 94 | + layer_idx += self.recompute_num_layers |
| 94 | else: | 95 | else: |
| 95 | - for l in range(num_layers_per_pipeline_rank): | 96 | + for layer_idx in range(num_layers_per_pipeline_rank): |
| 96 | - hidden_states = custom(l, l + 1)( | 97 | + hidden_states = custom(layer_idx, layer_idx + 1)( |
| 97 | - hidden_states, attention_mask, | 98 | + hidden_states, |
| 98 | - encoder_output, enc_dec_attn_mask, | 99 | + attention_mask, |
| 99 | - None, None, None, None, rotary_pos_emb) | 100 | + encoder_output, |
| 101 | + enc_dec_attn_mask, | ||
| 102 | + None, | ||
| 103 | + None, | ||
| 104 | + None, | ||
| 105 | + None, | ||
| 106 | + rotary_pos_emb, | ||
| 107 | + ) | ||
| 100 | elif self.recompute_method == 'block': | 108 | elif self.recompute_method == 'block': |
| 101 | # Checkpoint the input activation of only a set number of individual | 109 | # Checkpoint the input activation of only a set number of individual |
| 102 | # Transformer layers and skip the rest. | 110 | # Transformer layers and skip the rest. |
| 103 | # A method fully use the device memory removing redundant re-computation. | 111 | # A method fully use the device memory removing redundant re-computation. |
| 104 | - vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() | 112 | + for layer_idx in range(self.num_layers): |
| 105 | - vpp_size = global_args.virtual_pipeline_model_parallel_size | ||
| 106 | - if vpp_rank is None or not global_args.enable_recompute_layers_per_pp_rank: | ||
| 107 | - vpp_rank = 0 | ||
| 108 | - if vpp_size is None or not global_args.enable_recompute_layers_per_pp_rank: | ||
| 109 | - vpp_size = 1 | ||
| 110 | - for l in range(self.num_layers): | ||
| 111 | # The number of layers each pipeline rank recomputes is self.recompute_num_layers. | 113 | # The number of layers each pipeline rank recomputes is self.recompute_num_layers. |
| 112 | # If self.recompute_num_layers cannot divide exactly the number of layers in each pp rank, | 114 | # If self.recompute_num_layers cannot divide exactly the number of layers in each pp rank, |
| 113 | # we try to balance the number of recomputed layers in each model chunk. | 115 | # we try to balance the number of recomputed layers in each model chunk. |
| @@ -117,27 +119,39 @@ def parallel_transformer_checkpointed_forward(self, hidden_states, attention_mas | |||
| 117 | # Stage 1: [2, 3] [6, 7] | 119 | # Stage 1: [2, 3] [6, 7] |
| 118 | # With self.recompute_num_layers = 2, we will recompute layers 0,4 for stage 0, and 2,6 for stage 1. | 120 | # With self.recompute_num_layers = 2, we will recompute layers 0,4 for stage 0, and 2,6 for stage 1. |
| 119 | # With self.recompute_num_layers = 3, we will recompute layers 0,1,4 for stage 0, and 2,3,6 for stage 1. | 121 | # With self.recompute_num_layers = 3, we will recompute layers 0,1,4 for stage 0, and 2,3,6 for stage 1. |
| 120 | - def should_recompute(): | 122 | + layer = self._get_layer(layer_idx) |
| 121 | - if global_args.reduce_recompute_for_last_chunk: | 123 | + recompute_priority = get_recompute_priority(global_args, layer.layer_number) |
| 122 | - def is_last_layer(): | 124 | + should_recompute_layer = recompute_priority < self.recompute_num_layers |
| 123 | - return (l == self.num_layers - 1) and mpu.is_pipeline_last_stage() | 125 | + if global_args.reduce_recompute_for_last_chunk: |
| 126 | + is_last_layer = layer_idx == self.num_layers - 1 and mpu.is_pipeline_last_stage() | ||
| 127 | + should_recompute_layer = should_recompute_layer and not is_last_layer | ||
| 124 | 128 | ||
| 125 | - return ((l * vpp_size + vpp_rank) < self.recompute_num_layers) and not is_last_layer() | 129 | + if should_recompute_layer and not global_args.swap_attention: |
| 126 | - else: | ||
| 127 | - return (l * vpp_size + vpp_rank) < self.recompute_num_layers | ||
| 128 | - | ||
| 129 | - if should_recompute() and not global_args.swap_attention: | ||
| 130 | hidden_states = tensor_parallel.checkpoint( | 130 | hidden_states = tensor_parallel.checkpoint( |
| 131 | - custom(l, l + 1), | 131 | + custom(layer_idx, layer_idx + 1), |
| 132 | self.distribute_saved_activations, | 132 | self.distribute_saved_activations, |
| 133 | - hidden_states, attention_mask, | 133 | + hidden_states, |
| 134 | - encoder_output, enc_dec_attn_mask, | 134 | + attention_mask, |
| 135 | - None, None, None, None, rotary_pos_emb) | 135 | + encoder_output, |
| 136 | + enc_dec_attn_mask, | ||
| 137 | + None, | ||
| 138 | + None, | ||
| 139 | + None, | ||
| 140 | + None, | ||
| 141 | + rotary_pos_emb, | ||
| 142 | + ) | ||
| 136 | else: | 143 | else: |
| 137 | - hidden_states = custom(l, l + 1)( | 144 | + hidden_states = custom(layer_idx, layer_idx + 1)( |
| 138 | - hidden_states, attention_mask, | 145 | + hidden_states, |
| 139 | - encoder_output, enc_dec_attn_mask, | 146 | + attention_mask, |
| 140 | - None, None, None, None, rotary_pos_emb) | 147 | + encoder_output, |
| 148 | + enc_dec_attn_mask, | ||
| 149 | + None, | ||
| 150 | + None, | ||
| 151 | + None, | ||
| 152 | + None, | ||
| 153 | + rotary_pos_emb, | ||
| 154 | + ) | ||
| 141 | else: | 155 | else: |
| 142 | raise ValueError("Invalid activation recompute method.") | 156 | raise ValueError("Invalid activation recompute method.") |
| 143 | 157 | ||
| @@ -152,12 +166,14 @@ def core_mlp_forward_wrapper(fn): | |||
| 152 | 166 | ||
| 153 | if get_args().prof_file and not get_args().num_experts: | 167 | if get_args().prof_file and not get_args().num_experts: |
| 154 | from mindspeed.auto_settings.module.black.patch.hccl_operator import MOEOrMLPStartOp, MOEOrMLPEndOp | 168 | from mindspeed.auto_settings.module.black.patch.hccl_operator import MOEOrMLPStartOp, MOEOrMLPEndOp |
| 169 | + | ||
| 155 | args[0] = MOEOrMLPStartOp.apply(args[0]) | 170 | args[0] = MOEOrMLPStartOp.apply(args[0]) |
| 156 | activation_func_1 = torch.nn.Softplus() | 171 | activation_func_1 = torch.nn.Softplus() |
| 157 | args[0] = activation_func_1(args[0]) | 172 | args[0] = activation_func_1(args[0]) |
| 158 | 173 | ||
| 159 | self.layer_number = getattr(self, "layer_number", None) | 174 | self.layer_number = getattr(self, "layer_number", None) |
| 160 | is_recompute_activation = should_recompute_activation(self.layer_number) | 175 | is_recompute_activation = should_recompute_activation(self.layer_number) |
| 176 | + moe_ctx = None | ||
| 161 | if get_args().moe_alltoall_overlap_comm and not isinstance(args[-1], torch.Tensor): | 177 | if get_args().moe_alltoall_overlap_comm and not isinstance(args[-1], torch.Tensor): |
| 162 | moe_ctx = args[-1] | 178 | moe_ctx = args[-1] |
| 163 | args = args[:-1] | 179 | args = args[:-1] |
| @@ -167,7 +183,7 @@ def core_mlp_forward_wrapper(fn): | |||
| 167 | if bias is not None: | 183 | if bias is not None: |
| 168 | intermediate = intermediate + bias | 184 | intermediate = intermediate + bias |
| 169 | if self.config.gated_linear_unit: | 185 | if self.config.gated_linear_unit: |
| 170 | - assert (self.config.activation_func == F.silu), 'Activation function must be silu when using fused_swiglu' | 186 | + assert self.config.activation_func == F.silu, 'Activation function must be silu when using fused_swiglu' |
| 171 | if not hasattr(self, 'origin_activation_func'): | 187 | if not hasattr(self, 'origin_activation_func'): |
| 172 | self.origin_activation_func = self.activation_func | 188 | self.origin_activation_func = self.activation_func |
| 173 | self.activation_func = fused_swiglu | 189 | self.activation_func = fused_swiglu |
| @@ -182,8 +198,14 @@ def core_mlp_forward_wrapper(fn): | |||
| 182 | if hasattr(self, 'origin_activation_func'): | 198 | if hasattr(self, 'origin_activation_func'): |
| 183 | self.activation_func = self.origin_activation_func | 199 | self.activation_func = self.origin_activation_func |
| 184 | output, output_bias = fn(self, *args, **kwargs) | 200 | output, output_bias = fn(self, *args, **kwargs) |
| 185 | - elif moe_zero_memory == "level1" and not get_args().moe_fb_overlap and not only_recompute_activation(self.layer_number): | 201 | + elif ( |
| 202 | + moe_zero_memory == "level1" | ||
| 203 | + and not get_args().moe_fb_overlap | ||
| 204 | + and not only_recompute_activation(self.layer_number) | ||
| 205 | + ): | ||
| 186 | if self.with_shared_expert: | 206 | if self.with_shared_expert: |
| 207 | + if moe_ctx is None: | ||
| 208 | + raise RuntimeError("MoE context is required when shared-expert zero-memory recompute is enabled.") | ||
| 187 | self.activation_function = activation_function | 209 | self.activation_function = activation_function |
| 188 | hidden_states = args[0] | 210 | hidden_states = args[0] |
| 189 | fc1_out_parallel, bias_parallel = self.linear_fc1(hidden_states) | 211 | fc1_out_parallel, bias_parallel = self.linear_fc1(hidden_states) |
| @@ -199,10 +221,9 @@ def core_mlp_forward_wrapper(fn): | |||
| 199 | hidden_states = args[0] | 221 | hidden_states = args[0] |
| 200 | intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states) | 222 | intermediate_parallel, bias_parallel = self.linear_fc1(hidden_states) |
| 201 | self.activation_checkpoint_manager = CheckpointWithoutOutput() | 223 | self.activation_checkpoint_manager = CheckpointWithoutOutput() |
| 202 | - intermediate_parallel = self.activation_checkpoint_manager.checkpoint(activation_function, | 224 | + intermediate_parallel = self.activation_checkpoint_manager.checkpoint( |
| 203 | - False, | 225 | + activation_function, False, intermediate_parallel, bias_parallel |
| 204 | - intermediate_parallel, | 226 | + ) |
| 205 | - bias_parallel) | ||
| 206 | # [s, b, h] | 227 | # [s, b, h] |
| 207 | output, output_bias = self.linear_fc2(intermediate_parallel) | 228 | output, output_bias = self.linear_fc2(intermediate_parallel) |
| 208 | 229 | ||
| @@ -221,18 +242,17 @@ def core_mlp_forward_wrapper(fn): | |||
| 221 | output = MOEOrMLPEndOp.apply(output) | 242 | output = MOEOrMLPEndOp.apply(output) |
| 222 | 243 | ||
| 223 | return output, output_bias | 244 | return output, output_bias |
| 245 | + | ||
| 224 | return wrapper | 246 | return wrapper |
| 225 | 247 | ||
| 226 | 248 | ||
| 227 | def enable_recompute_norm_checkpoint( | 249 | def enable_recompute_norm_checkpoint( |
| 228 | - layer, | 250 | + layer, norm_ckpt, submodule_name: Optional[str] = None, support_module_type=MindSpeedTELayerNormColumnParallelLinear |
| 229 | - norm_ckpt, | ||
| 230 | - submodule_name: Optional[str] = None, | ||
| 231 | - support_module_type=MindSpeedTELayerNormColumnParallelLinear | ||
| 232 | ): | 251 | ): |
| 233 | if layer is None or norm_ckpt is None: | 252 | if layer is None or norm_ckpt is None: |
| 234 | raise ValueError("Please check your input!!!") | 253 | raise ValueError("Please check your input!!!") |
| 235 | 254 | ||
| 255 | + target_layer = None | ||
| 236 | if submodule_name is not None: | 256 | if submodule_name is not None: |
| 237 | target_layer = getattr(layer, submodule_name, None) | 257 | target_layer = getattr(layer, submodule_name, None) |
| 238 | 258 | ||
| @@ -340,8 +360,6 @@ def norm_recompute_forward( | |||
| 340 | mlp_output_with_bias, residual, self.hidden_dropout | 360 | mlp_output_with_bias, residual, self.hidden_dropout |
| 341 | ) | 361 | ) |
| 342 | 362 | ||
| 343 | - output = make_viewless_tensor( | 363 | + output = make_viewless_tensor(inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True) |
| 344 | - inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True | ||
| 345 | - ) | ||
| 346 | 364 | ||
| 347 | return output, context | 365 | return output, context |
| @@ -8,10 +8,12 @@ from megatron.core import tensor_parallel, parallel_state, mpu | |||
| 8 | from megatron.core.packed_seq_params import PackedSeqParams | 8 | from megatron.core.packed_seq_params import PackedSeqParams |
| 9 | from megatron.core.transformer.module import MegatronModule | 9 | from megatron.core.transformer.module import MegatronModule |
| 10 | from megatron.core.transformer.spec_utils import build_module | 10 | from megatron.core.transformer.spec_utils import build_module |
| 11 | -from megatron.core.extensions.transformer_engine import TENorm | ||
| 12 | from mindspeed.args_utils import get_full_args as get_args | 11 | from mindspeed.args_utils import get_full_args as get_args |
| 13 | -from mindspeed.core.tensor_parallel.comm_autograd_function import auto_grad_sync_gather_along_last_dim, \ | 12 | +from mindspeed.core.memory.recompute.recompute_common import get_recompute_priority |
| 14 | - auto_grad_sync_gather_along_first_dim | 13 | +from mindspeed.core.tensor_parallel.comm_autograd_function import ( |
| 14 | + auto_grad_sync_gather_along_last_dim, | ||
| 15 | + auto_grad_sync_gather_along_first_dim, | ||
| 16 | +) | ||
| 15 | from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPYCollectiveComm | 17 | from mindspeed.core.tensor_parallel.comm_group_api import TPXCollectiveComm, TPYCollectiveComm |
| 16 | 18 | ||
| 17 | from mindspeed.deprecate import Deprecated, MEGATRON_ADAPTOR_DEPRECATED_TIME | 19 | from mindspeed.deprecate import Deprecated, MEGATRON_ADAPTOR_DEPRECATED_TIME |
| @@ -101,14 +103,14 @@ def transformer_block_checkpointed_forward( | |||
| 101 | # checkpoint the input activation of each divided chunk. | 103 | # checkpoint the input activation of each divided chunk. |
| 102 | # A method to further reduce memory usage reducing checkpoints. | 104 | # A method to further reduce memory usage reducing checkpoints. |
| 103 | if not global_args.swap_attention: | 105 | if not global_args.swap_attention: |
| 104 | - l = 0 | 106 | + layer_idx = 0 |
| 105 | - while l < self.num_layers_per_pipeline_rank: | 107 | + while layer_idx < self.num_layers_per_pipeline_rank: |
| 106 | - hidden_states = checkpoint_handler(custom(l, l + 1)) | 108 | + hidden_states = checkpoint_handler(custom(layer_idx, layer_idx + 1)) |
| 107 | 109 | ||
| 108 | - l += self.config.recompute_num_layers | 110 | + layer_idx += self.config.recompute_num_layers |
| 109 | else: | 111 | else: |
| 110 | - for l in range(self.num_layers_per_pipeline_rank): | 112 | + for layer_idx in range(self.num_layers_per_pipeline_rank): |
| 111 | - hidden_states, context = custom(l, l + 1)( | 113 | + hidden_states, context = custom(layer_idx, layer_idx + 1)( |
| 112 | hidden_states, | 114 | hidden_states, |
| 113 | attention_mask, | 115 | attention_mask, |
| 114 | context, | 116 | context, |
| @@ -116,13 +118,7 @@ def transformer_block_checkpointed_forward( | |||
| 116 | rotary_pos_emb, | 118 | rotary_pos_emb, |
| 117 | ) | 119 | ) |
| 118 | elif self.config.recompute_method == 'block': | 120 | elif self.config.recompute_method == 'block': |
| 119 | - vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() | 121 | + for layer_idx in range(self.num_layers_per_pipeline_rank): |
| 120 | - vpp_size = self.config.virtual_pipeline_model_parallel_size | ||
| 121 | - if vpp_rank is None or not global_args.enable_recompute_layers_per_pp_rank: | ||
| 122 | - vpp_rank = 0 | ||
| 123 | - if vpp_size is None or not global_args.enable_recompute_layers_per_pp_rank: | ||
| 124 | - vpp_size = 1 | ||
| 125 | - for l in range(self.num_layers_per_pipeline_rank): | ||
| 126 | # The number of layers each pipeline rank recomputes is self.recompute_num_layers. | 122 | # The number of layers each pipeline rank recomputes is self.recompute_num_layers. |
| 127 | # If self.recompute_num_layers cannot divide exactly the number of layers in each pp rank, | 123 | # If self.recompute_num_layers cannot divide exactly the number of layers in each pp rank, |
| 128 | # we try to balance the number of recomputed layers in each model chunk. | 124 | # we try to balance the number of recomputed layers in each model chunk. |
| @@ -132,19 +128,21 @@ def transformer_block_checkpointed_forward( | |||
| 132 | # Stage 1: [2, 3] [6, 7] | 128 | # Stage 1: [2, 3] [6, 7] |
| 133 | # With self.recompute_num_layers = 2, we will recompute layers 0,4 for stage 0, and 2,6 for stage 1. | 129 | # With self.recompute_num_layers = 2, we will recompute layers 0,4 for stage 0, and 2,6 for stage 1. |
| 134 | # With self.recompute_num_layers = 3, we will recompute layers 0,1,4 for stage 0, and 2,3,6 for stage 1. | 130 | # With self.recompute_num_layers = 3, we will recompute layers 0,1,4 for stage 0, and 2,3,6 for stage 1. |
| 135 | - def should_recompute(): | 131 | + layer = self._get_layer(layer_idx) |
| 136 | - if getattr(global_args, 'reduce_recompute_for_last_chunk', False): | 132 | + recompute_priority = get_recompute_priority( |
| 137 | - def is_last_layer(): | 133 | + self.config, |
| 138 | - return (l == self.num_layers_per_pipeline_rank - 1) and mpu.is_pipeline_last_stage() | 134 | + layer.layer_number, |
| 135 | + global_args.enable_recompute_layers_per_pp_rank, | ||
| 136 | + ) | ||
| 137 | + should_recompute_layer = recompute_priority < self.config.recompute_num_layers | ||
| 138 | + if getattr(global_args, 'reduce_recompute_for_last_chunk', False): | ||
| 139 | + is_last_layer = layer_idx == self.num_layers_per_pipeline_rank - 1 and mpu.is_pipeline_last_stage() | ||
| 140 | + should_recompute_layer = should_recompute_layer and not is_last_layer | ||
| 139 | 141 | ||
| 140 | - return ((l * vpp_size + vpp_rank) < self.config.recompute_num_layers) and not is_last_layer() | 142 | + if should_recompute_layer and not global_args.swap_attention: |
| 141 | - else: | 143 | + hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1)) |
| 142 | - return (l * vpp_size + vpp_rank) < self.config.recompute_num_layers | ||
| 143 | - | ||
| 144 | - if should_recompute() and not global_args.swap_attention: | ||
| 145 | - hidden_states, context = checkpoint_handler(custom(l, l + 1)) | ||
| 146 | else: | 144 | else: |
| 147 | - hidden_states, context = custom(l, l + 1)( | 145 | + hidden_states, context = custom(layer_idx, layer_idx + 1)( |
| 148 | hidden_states, | 146 | hidden_states, |
| 149 | attention_mask, | 147 | attention_mask, |
| 150 | context, | 148 | context, |
| @@ -160,9 +158,21 @@ class NoopTransformerLayer(MegatronModule): | |||
| 160 | super().__init__(None) | 158 | super().__init__(None) |
| 161 | self.layer_number = layer_number | 159 | self.layer_number = layer_number |
| 162 | 160 | ||
| 163 | - def forward(self, hidden_states, attention_mask, context, context_mask, rotary_pos_emb, rotary_pos_cos=None, | 161 | + def forward( |
| 164 | - rotary_pos_sin=None, inference_params=None, attention_bias=None, inference_context=None, | 162 | + self, |
| 165 | - packed_seq_params=None, sequence_len_offset=None): | 163 | + hidden_states, |
| 164 | + attention_mask, | ||
| 165 | + context, | ||
| 166 | + context_mask, | ||
| 167 | + rotary_pos_emb, | ||
| 168 | + rotary_pos_cos=None, | ||
| 169 | + rotary_pos_sin=None, | ||
| 170 | + inference_params=None, | ||
| 171 | + attention_bias=None, | ||
| 172 | + inference_context=None, | ||
| 173 | + packed_seq_params=None, | ||
| 174 | + sequence_len_offset=None, | ||
| 175 | + ): | ||
| 166 | return hidden_states.clone(), context | 176 | return hidden_states.clone(), context |
| 167 | 177 | ||
| 168 | 178 | ||
| @@ -170,9 +180,7 @@ def _get_layer_offset(args): | |||
| 170 | num_layers = args.num_layers | 180 | num_layers = args.num_layers |
| 171 | pipeline_rank = parallel_state.get_pipeline_model_parallel_rank() | 181 | pipeline_rank = parallel_state.get_pipeline_model_parallel_rank() |
| 172 | 182 | ||
| 173 | - num_layers_per_pipeline_rank = ( | 183 | + num_layers_per_pipeline_rank = num_layers // parallel_state.get_pipeline_model_parallel_world_size() |
| 174 | - num_layers // parallel_state.get_pipeline_model_parallel_world_size() | ||
| 175 | - ) | ||
| 176 | 184 | ||
| 177 | if args.schedules_method == "dualpipev": | 185 | if args.schedules_method == "dualpipev": |
| 178 | num_layers_per_dualpipe_chunk = num_layers_per_pipeline_rank // 2 | 186 | num_layers_per_dualpipe_chunk = num_layers_per_pipeline_rank // 2 |
| @@ -211,16 +219,20 @@ def _build_layers(self): | |||
| 211 | 219 | ||
| 212 | def build_layer(layer_spec, layer_number): | 220 | def build_layer(layer_spec, layer_number): |
| 213 | global_layer_number = _get_layer_offset(args) + layer_number | 221 | global_layer_number = _get_layer_offset(args) + layer_number |
| 214 | - if (hasattr(args, 'noop_layers') and isinstance(args.noop_layers, set) | 222 | + if ( |
| 215 | - and global_layer_number - 1 in args.noop_layers): | 223 | + hasattr(args, 'noop_layers') |
| 224 | + and isinstance(args.noop_layers, set) | ||
| 225 | + and global_layer_number - 1 in args.noop_layers | ||
| 226 | + ): | ||
| 216 | return NoopTransformerLayer(global_layer_number) | 227 | return NoopTransformerLayer(global_layer_number) |
| 217 | - return build_module(layer_spec, config=self.config, layer_number=layer_number, ) | 228 | + return build_module( |
| 229 | + layer_spec, | ||
| 230 | + config=self.config, | ||
| 231 | + layer_number=layer_number, | ||
| 232 | + ) | ||
| 218 | 233 | ||
| 219 | self.layers = torch.nn.ModuleList( | 234 | self.layers = torch.nn.ModuleList( |
| 220 | - [ | 235 | + [build_layer(layer_spec, i + 1) for i, layer_spec in enumerate(self.submodules.layer_specs)] |
| 221 | - build_layer(layer_spec, i + 1) | ||
| 222 | - for i, layer_spec in enumerate(self.submodules.layer_specs) | ||
| 223 | - ] | ||
| 224 | ) | 236 | ) |
| 225 | 237 | ||
| 226 | if self.submodules.layer_norm and self.post_process and self.post_layer_norm: | 238 | if self.submodules.layer_norm and self.post_process and self.post_layer_norm: |
| @@ -18,47 +18,48 @@ class SwapLayerInputFeature(MindSpeedFeature): | |||
| 18 | ) | 18 | ) |
| 19 | 19 | ||
| 20 | def register_patches(self, patch_manager, args): | 20 | def register_patches(self, patch_manager, args): |
| 21 | - if getattr(args, self.feature_name, None): | 21 | + if not getattr(args, self.feature_name, False): |
| 22 | - from mindspeed.core.memory.swap_layer_input.swap_layer_input import ( | 22 | + return |
| 23 | - swap_layer_input_init_wrapper, | ||
| 24 | - swap_layer_input_forward_wrapper, | ||
| 25 | - swap_layer_input_fboverlap_forward_wrapper, | ||
| 26 | - swap_layer_input_fboverlap_1f1b_wrapper, | ||
| 27 | - swap_layer_input_fboverlap_backward_wrapper, | ||
| 28 | - ) | ||
| 29 | 23 | ||
| 30 | - patch_manager.register_patch( | 24 | + from mindspeed.core.memory.swap_layer_input.swap_layer_input import ( |
| 31 | - 'megatron.core.transformer.transformer_layer.TransformerLayer.__init__', swap_layer_input_init_wrapper | 25 | + swap_layer_input_fboverlap_1f1b_wrapper, |
| 32 | - ) | 26 | + swap_layer_input_fboverlap_backward_wrapper, |
| 33 | - patch_manager.register_patch( | 27 | + swap_layer_input_fboverlap_forward_wrapper, |
| 34 | - 'megatron.core.transformer.transformer_layer.TransformerLayer.forward', swap_layer_input_forward_wrapper | 28 | + swap_layer_input_forward_wrapper, |
| 35 | - ) | 29 | + swap_layer_input_init_wrapper, |
| 36 | - if getattr(args, 'moe_fb_overlap', None): | 30 | + ) |
| 37 | - patch_manager.register_patch( | 31 | + |
| 38 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwd.transformer_layer_forward_moe', | 32 | + patch_manager.register_patch( |
| 39 | - swap_layer_input_fboverlap_forward_wrapper, | 33 | + 'megatron.core.transformer.transformer_layer.TransformerLayer.__init__', |
| 40 | - ) | 34 | + swap_layer_input_init_wrapper, |
| 41 | - patch_manager.register_patch( | 35 | + ) |
| 42 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwd.transformer_layer_forward_dense', | 36 | + patch_manager.register_patch( |
| 43 | - swap_layer_input_fboverlap_forward_wrapper, | 37 | + 'megatron.core.transformer.transformer_layer.TransformerLayer.forward', |
| 44 | - ) | 38 | + swap_layer_input_forward_wrapper, |
| 45 | - patch_manager.register_patch( | 39 | + ) |
| 46 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.transformer_layer_forward_dense_backward_dense_overlaping', | 40 | + |
| 47 | - swap_layer_input_fboverlap_1f1b_wrapper, | 41 | + if not getattr(args, 'moe_fb_overlap', False): |
| 48 | - ) | 42 | + return |
| 49 | - patch_manager.register_patch( | 43 | + |
| 50 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.transformer_layer_forward_moe_backward_dense_overlaping', | 44 | + for target in ( |
| 51 | - swap_layer_input_fboverlap_1f1b_wrapper, | 45 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwd.transformer_layer_forward_moe', |
| 52 | - ) | 46 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwd.transformer_layer_forward_dense', |
| 53 | - patch_manager.register_patch( | 47 | + ): |
| 54 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.transformer_layer_forward_dense_backward_moe_overlaping', | 48 | + patch_manager.register_patch(target, swap_layer_input_fboverlap_forward_wrapper) |
| 55 | - swap_layer_input_fboverlap_1f1b_wrapper, | 49 | + |
| 56 | - ) | 50 | + for target in ( |
| 57 | - patch_manager.register_patch( | 51 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.' |
| 58 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.transformer_layer_forward_moe_backward_moe_overlaping', | 52 | + 'transformer_layer_forward_dense_backward_dense_overlaping', |
| 59 | - swap_layer_input_fboverlap_1f1b_wrapper, | 53 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.' |
| 60 | - ) | 54 | + 'transformer_layer_forward_moe_backward_dense_overlaping', |
| 61 | - patch_manager.register_patch( | 55 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.' |
| 62 | - 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.transformer_layer.transformer_layer_backward', | 56 | + 'transformer_layer_forward_dense_backward_moe_overlaping', |
| 63 | - swap_layer_input_fboverlap_backward_wrapper, | 57 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.overlap_funcs.fwdbwd.' |
| 64 | - ) | 58 | + 'transformer_layer_forward_moe_backward_moe_overlaping', |
| 59 | + ): | ||
| 60 | + patch_manager.register_patch(target, swap_layer_input_fboverlap_1f1b_wrapper) | ||
| 61 | + | ||
| 62 | + patch_manager.register_patch( | ||
| 63 | + 'mindspeed.core.transformer.moe.moe_feature.fb_overlap.transformer_layer.transformer_layer_backward', | ||
| 64 | + swap_layer_input_fboverlap_backward_wrapper, | ||
| 65 | + ) | ||
| @@ -17,9 +17,9 @@ | |||
| 17 | import warnings | 17 | import warnings |
| 18 | 18 | ||
| 19 | import torch | 19 | import torch |
| 20 | -from megatron.core import parallel_state, mpu, tensor_parallel | ||
| 21 | from megatron.core.transformer.module import MegatronModule | 20 | from megatron.core.transformer.module import MegatronModule |
| 22 | from mindspeed.args_utils import get_full_args as get_args | 21 | from mindspeed.args_utils import get_full_args as get_args |
| 22 | +from mindspeed.core.memory.recompute.recompute_common import should_recompute as should_recompute_common | ||
| 23 | from mindspeed.core.tensor_parallel.checkpoint_manager import get_pipeline_checkpoint_manager | 23 | from mindspeed.core.tensor_parallel.checkpoint_manager import get_pipeline_checkpoint_manager |
| 24 | 24 | ||
| 25 | try: | 25 | try: |
| @@ -49,27 +49,30 @@ def generate_attention_mask(compress, device): | |||
| 49 | args = get_args() | 49 | args = get_args() |
| 50 | if not args.use_flash_attn: | 50 | if not args.use_flash_attn: |
| 51 | warnings.warn("Flash Attention is highly recommended") | 51 | warnings.warn("Flash Attention is highly recommended") |
| 52 | - _GLOBAL_ATTN_MASK = (torch.tril(torch.ones([args.micro_batch_size, 1, args.seq_length, args.seq_length], dtype=bool, device=device), diagonal=-(args.pre_tockens + 1)) \ | 52 | + _GLOBAL_ATTN_MASK = torch.tril( |
| 53 | - + torch.triu(torch.ones([args.micro_batch_size, 1, args.seq_length, args.seq_length], dtype=bool, device=device), diagonal=args.next_tockens + 1)) | 53 | + torch.ones([args.micro_batch_size, 1, args.seq_length, args.seq_length], dtype=bool, device=device), |
| 54 | + diagonal=-(args.pre_tockens + 1), | ||
| 55 | + ) + torch.triu( | ||
| 56 | + torch.ones([args.micro_batch_size, 1, args.seq_length, args.seq_length], dtype=bool, device=device), | ||
| 57 | + diagonal=args.next_tockens + 1, | ||
| 58 | + ) | ||
| 54 | return | 59 | return |
| 55 | 60 | ||
| 56 | if compress: | 61 | if compress: |
| 57 | seq_len = 2048 | 62 | seq_len = 2048 |
| 58 | else: | 63 | else: |
| 59 | seq_len = args.seq_length | 64 | seq_len = args.seq_length |
| 60 | - | 65 | + |
| 61 | - _GLOBAL_ATTN_MASK = torch.triu( | 66 | + _GLOBAL_ATTN_MASK = torch.triu(torch.ones((seq_len, seq_len), device=device, dtype=torch.bool), diagonal=1) |
| 62 | - torch.ones((seq_len, seq_len), | ||
| 63 | - device=device, dtype=torch.bool), diagonal=1) | ||
| 64 | 67 | ||
| 65 | 68 | ||
| 66 | def get_attention_mask(): | 69 | def get_attention_mask(): |
| 67 | - global _GLOBAL_ATTN_MASK | ||
| 68 | if _GLOBAL_ATTN_MASK is not None: | 70 | if _GLOBAL_ATTN_MASK is not None: |
| 69 | return _GLOBAL_ATTN_MASK | 71 | return _GLOBAL_ATTN_MASK |
| 70 | 72 | ||
| 71 | args = get_args() | 73 | args = get_args() |
| 72 | should_generate_mask = False | 74 | should_generate_mask = False |
| 75 | + compress = False | ||
| 73 | device = 'npu' | 76 | device = 'npu' |
| 74 | 77 | ||
| 75 | if args.attention_mask_type == 'causal': | 78 | if args.attention_mask_type == 'causal': |
| @@ -93,7 +96,6 @@ def get_attention_mask(): | |||
| 93 | should_generate_mask = True | 96 | should_generate_mask = True |
| 94 | compress = True | 97 | compress = True |
| 95 | 98 | ||
| 96 | - | ||
| 97 | if getattr(args, 'attention_mask_on_cpu', False): | 99 | if getattr(args, 'attention_mask_on_cpu', False): |
| 98 | device = 'cpu' | 100 | device = 'cpu' |
| 99 | 101 | ||
| @@ -104,39 +106,7 @@ def get_attention_mask(): | |||
| 104 | 106 | ||
| 105 | 107 | ||
| 106 | def should_recompute(args, layer_number, num_recompute): | 108 | def should_recompute(args, layer_number, num_recompute): |
| 107 | - vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() | 109 | + return should_recompute_common(args, layer_number, num_recompute) |
| 108 | - vpp_size = args.virtual_pipeline_model_parallel_size | ||
| 109 | - pp_size = args.pipeline_model_parallel_size | ||
| 110 | - | ||
| 111 | - if vpp_size is not None: | ||
| 112 | - layer_per_chunk = args.num_layers_per_virtual_pipeline_stage | ||
| 113 | - elif pp_size is not None: | ||
| 114 | - layer_per_chunk = args.num_layers // pp_size | ||
| 115 | - else: | ||
| 116 | - layer_per_chunk = args.num_layers | ||
| 117 | - | ||
| 118 | - if vpp_rank is None or not args.enable_recompute_layers_per_pp_rank: | ||
| 119 | - vpp_rank = 0 | ||
| 120 | - if vpp_size is None or not args.enable_recompute_layers_per_pp_rank: | ||
| 121 | - vpp_size = 1 | ||
| 122 | - recompute_priority = ((layer_number - 1) % layer_per_chunk) * vpp_size + vpp_rank | ||
| 123 | - full_recompute_layers = args.recompute_num_layers | ||
| 124 | - | ||
| 125 | - if full_recompute_layers: | ||
| 126 | - if recompute_priority < full_recompute_layers: | ||
| 127 | - # Do full recomputation | ||
| 128 | - return False | ||
| 129 | - elif num_recompute is None: | ||
| 130 | - return True | ||
| 131 | - elif recompute_priority < full_recompute_layers + num_recompute: | ||
| 132 | - return True | ||
| 133 | - else: | ||
| 134 | - return False | ||
| 135 | - | ||
| 136 | - if num_recompute is None: | ||
| 137 | - return True | ||
| 138 | - else: | ||
| 139 | - return recompute_priority < num_recompute | ||
| 140 | 110 | ||
| 141 | 111 | ||
| 142 | def should_recompute_activation(layer_number): | 112 | def should_recompute_activation(layer_number): |
| @@ -155,7 +125,9 @@ def should_recompute_activation(layer_number): | |||
| 155 | if args.recompute_activation_function_num_layers < 0: | 125 | if args.recompute_activation_function_num_layers < 0: |
| 156 | raise AssertionError('--recompute-activation-function-num-layers cannot be less than 0.') | 126 | raise AssertionError('--recompute-activation-function-num-layers cannot be less than 0.') |
| 157 | elif args.recompute_activation_function_num_layers > args.num_layers: | 127 | elif args.recompute_activation_function_num_layers > args.num_layers: |
| 158 | - raise AssertionError('--recompute-activation-function-num-layers cannot be greater than the number of layers.') | 128 | + raise AssertionError( |
| 129 | + '--recompute-activation-function-num-layers cannot be greater than the number of layers.' | ||
| 130 | + ) | ||
| 159 | return should_recompute(args, layer_number, args.recompute_activation_function_num_layers) | 131 | return should_recompute(args, layer_number, args.recompute_activation_function_num_layers) |
| 160 | 132 | ||
| 161 | 133 | ||