已合并
fboverlap supports swap unperm2 activation #2212
AtomGit-Bot创建于 2025年4月18日
fboverlap supports swap unperm2 activation #2212
已合并
从refs/pull/2212/head合入到core_r0.8.0
共 8 个文件变更+585-21
| @@ -8,7 +8,7 @@ from megatron.core.transformer.moe.moe_utils import permute, unpermute | |||
| 8 | from megatron.core.tensor_parallel.mappings import _gather_along_first_dim_expert_parallel | 8 | from megatron.core.tensor_parallel.mappings import _gather_along_first_dim_expert_parallel |
| 9 | from megatron.core.utils import make_viewless_tensor | 9 | from megatron.core.utils import make_viewless_tensor |
| 10 | from megatron.training import get_args | 10 | from megatron.training import get_args |
| 11 | -from mindspeed.core.transformer.moe.unpermute_without_activation import UnpermuteWithoutActivation | 11 | +from mindspeed.core.pipeline_parallel.fb_overlap.modules.utils import TensorSwapManager |
| 12 | 12 | ||
| 13 | 13 | ||
| 14 | def preprocess(self, indices: torch.Tensor) -> torch.Tensor: | 14 | def preprocess(self, indices: torch.Tensor) -> torch.Tensor: |
| @@ -269,19 +269,38 @@ def alltoall_token_unperm1( | |||
| 269 | def alltoall_token_unperm2(self, permutated_local_input_tokens, probs=None): | 269 | def alltoall_token_unperm2(self, permutated_local_input_tokens, probs=None): |
| 270 | # Unpermutation 1: AlltoAll output to output | 270 | # Unpermutation 1: AlltoAll output to output |
| 271 | 271 | ||
| 272 | + args = get_args() | ||
| 272 | probs = probs if probs is not None else self.probs | 273 | probs = probs if probs is not None else self.probs |
| 273 | - output = unpermute( | 274 | + if args.moe_unperm2_mem_optim_swap: |
| 275 | + probs = None | ||
| 276 | + output_unperm = unpermute( | ||
| 274 | permutated_local_input_tokens, | 277 | permutated_local_input_tokens, |
| 275 | self.reversed_local_input_permutation_mapping, | 278 | self.reversed_local_input_permutation_mapping, |
| 276 | probs=probs, | 279 | probs=probs, |
| 277 | padded_mode=self.drop_and_pad, | 280 | padded_mode=self.drop_and_pad, |
| 278 | restore_shape=self.hiddden_shape_before_permute, | 281 | restore_shape=self.hiddden_shape_before_permute, |
| 279 | ) | 282 | ) |
| 280 | - if get_args().moe_unperm2_mem_optim: | 283 | + output_swap_manager = None |
| 281 | - output = output.view(-1, self.router_topk, output.shape[-1]) | 284 | + if args.moe_unperm2_mem_optim: |
| 282 | - output = output.sum(dim=1) | 285 | + output_unperm = output_unperm.view(-1, self.router_topk, output_unperm.shape[-1]) |
| 283 | - | 286 | + output = output_unperm.sum(dim=1) |
| 287 | + elif args.moe_unperm2_mem_optim_swap: | ||
| 288 | + permutated_local_input_tokens.untyped_storage().resize_(0) | ||
| 289 | + self.probs_detached = self.probs.detach() | ||
| 290 | + assert not self.probs_detached.requires_grad | ||
| 284 | 291 | ||
| 292 | + # [s, k, h] * [s, k, 1] -> [s, k, h] | ||
| 293 | + # forward pass here is broadcast and mul | ||
| 294 | + # backward pass is mul and reducesum | ||
| 295 | + weighted_output = output_unperm.view(-1, self.router_topk, output_unperm.shape[-1]) * self.probs_detached.unsqueeze(-1) | ||
| 296 | + # swap output | ||
| 297 | + TensorSwapManager.wait_all_swap_out('unperm2') | ||
| 298 | + output_swap_manager = TensorSwapManager(output_unperm, 'unperm2') | ||
| 299 | + output_swap_manager.async_swap_out(wait_stream=torch.npu.current_stream()) | ||
| 300 | + output = weighted_output.sum(dim=1).to(permutated_local_input_tokens.dtype) # [s, h] | ||
| 301 | + else: | ||
| 302 | + # cast for unfused unpermute | ||
| 303 | + output = output_unperm.to(permutated_local_input_tokens.dtype) | ||
| 285 | # Reshape the output tensor | 304 | # Reshape the output tensor |
| 286 | output = output.view(self.hidden_shape) | 305 | output = output.view(self.hidden_shape) |
| 287 | 306 | ||
| @@ -289,5 +308,4 @@ def alltoall_token_unperm2(self, permutated_local_input_tokens, probs=None): | |||
| 289 | inp=output, requires_grad=output.requires_grad, keep_graph=True | 308 | inp=output, requires_grad=output.requires_grad, keep_graph=True |
| 290 | ) | 309 | ) |
| 291 | 310 | ||
| 292 | - | 311 | + return output, output_swap_manager |
| 293 | - return output, None | ||
| @@ -1,3 +1,7 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved. | ||
| 2 | + | ||
| 3 | +from typing import Optional, Dict, List | ||
A | |||
| 4 | + | ||
| 1 | import torch | 5 | import torch |
| 2 | from torch.autograd.variable import Variable | 6 | from torch.autograd.variable import Variable |
| 3 | from megatron.core.pipeline_parallel import p2p_communication | 7 | from megatron.core.pipeline_parallel import p2p_communication |
| @@ -141,3 +145,164 @@ def p2p_comm_helper(comm_params: P2PCommParams, tensor_tosend): | |||
| 141 | else: | 145 | else: |
| 142 | return None, p2p_handles | 146 | return None, p2p_handles |
| 143 | 147 | ||
| 148 | + | ||
| 149 | +class TensorSwapManager: | ||
| 150 | + """Manager for asynchronous tensor swapping between NPU and CPU memory.""" | ||
| 151 | + | ||
| 152 | + _SWAP_OUT_STREAM = None | ||
| 153 | + _SWAP_IN_STREAM = None | ||
| 154 | + _ALL_SWAP_OUT_QUEUES = {} # Dictionary of swap groups for coordinated operations | ||
| 155 | + | ||
| 156 | + | ||
| 157 | + def _get_swap_out_stream(cls): | ||
| 158 | + """Get or create the swap-out stream.""" | ||
| 159 | + if cls._SWAP_OUT_STREAM is None: | ||
| 160 | + cls._SWAP_OUT_STREAM = torch.npu.Stream(device=torch.npu.current_device()) | ||
| 161 | + return cls._SWAP_OUT_STREAM | ||
| 162 | + | ||
| 163 | + | ||
| 164 | + def _get_swap_in_stream(cls): | ||
| 165 | + """Get or create the swap-in stream.""" | ||
| 166 | + if cls._SWAP_IN_STREAM is None: | ||
| 167 | + cls._SWAP_IN_STREAM = torch.npu.Stream(device=torch.npu.current_device()) | ||
| 168 | + return cls._SWAP_IN_STREAM | ||
| 169 | + | ||
| 170 | + def __init__(self, tensor, swap_group_name=None): | ||
| 171 | + """ | ||
| 172 | + Initialize a tensor swap manager. | ||
| 173 | + | ||
| 174 | + Args: | ||
| 175 | + tensor: The NPU tensor to swap | ||
| 176 | + swap_group_name: Optional group name for coordinated swap operations | ||
| 177 | + """ | ||
| 178 | + # Check if tensor is a slice (requires special handling, not supported yet) | ||
| 179 | + if tensor.storage().size() != tensor.numel(): | ||
| 180 | + raise AssertionError('TensorSwapManager cannot handle sliced tensors') | ||
| 181 | + | ||
| 182 | + self.npu_tensor = tensor | ||
| 183 | + self.cpu_tensor = None | ||
| 184 | + self.swap_out_event = None | ||
| 185 | + self.swap_in_event = None | ||
| 186 | + | ||
| 187 | + self.swap_group_name = swap_group_name # Group for coordinated operations | ||
| 188 | + self.under_swap_in = False | ||
| 189 | + | ||
| 190 | + # Initialize queue for this swap group if it doesn't exist | ||
| 191 | + if swap_group_name and swap_group_name not in self._ALL_SWAP_OUT_QUEUES: | ||
| 192 | + self._ALL_SWAP_OUT_QUEUES[swap_group_name] = [] | ||
| 193 | + | ||
| 194 | + def async_swap_out(self, wait_event=None, wait_stream=None): | ||
| 195 | + """ | ||
| 196 | + Asynchronously copy tensor from NPU to CPU memory. | ||
| 197 | + | ||
| 198 | + Args: | ||
| 199 | + event: Optional event to wait for before swapping | ||
| 200 | + stream: Optional stream to synchronize with before swapping | ||
| 201 | + """ | ||
| 202 | + swap_stream = self._get_swap_out_stream() | ||
| 203 | + # Allocate pinned CPU memory (enables faster async transfers) | ||
| 204 | + self.cpu_tensor = torch.empty_like(self.npu_tensor, | ||
| 205 | + pin_memory=True, | ||
| 206 | + device='cpu') | ||
| 207 | + | ||
| 208 | + with torch.npu.stream(swap_stream): | ||
| 209 | + if wait_event: | ||
| 210 | + swap_stream.wait_event(wait_event) | ||
| 211 | + if wait_stream: | ||
| 212 | + swap_stream.wait_stream(wait_stream) | ||
| 213 | + | ||
| 214 | + self.cpu_tensor.untyped_storage().copy_( | ||
| 215 | + self.npu_tensor.untyped_storage(), | ||
| 216 | + non_blocking=True) | ||
| 217 | + | ||
| 218 | + self.swap_out_event = torch.npu.Event() | ||
| 219 | + self.swap_out_event.record() | ||
| 220 | + | ||
| 221 | + # Add to swap group if specified | ||
| 222 | + if self.swap_group_name: | ||
| 223 | + self._ALL_SWAP_OUT_QUEUES[self.swap_group_name].append(self) | ||
| 224 | + | ||
| 225 | + def wait_swap_out(self): | ||
| 226 | + """Wait for swap-out to complete and release NPU memory.""" | ||
| 227 | + if self.swap_out_event: | ||
| 228 | + torch.npu.current_stream().wait_event(self.swap_out_event) | ||
| 229 | + # Release NPU memory (but keep storage object alive) | ||
| 230 | + if not self.under_swap_in: | ||
| 231 | + self.npu_tensor.untyped_storage().resize_(0) | ||
| 232 | + | ||
| 233 | + def async_swap_in(self, wait_event=None, wait_stream=None): | ||
| 234 | + """ | ||
| 235 | + Asynchronously copy tensor from CPU back to NPU memory. | ||
| 236 | + | ||
| 237 | + Args: | ||
| 238 | + event: Optional event to wait for before swapping | ||
| 239 | + stream: Optional stream to synchronize with before swapping | ||
| 240 | + """ | ||
| 241 | + self.under_swap_in = True | ||
| 242 | + if self.npu_tensor.untyped_storage().size() != 0: | ||
| 243 | + return | ||
| 244 | + swap_stream = self._get_swap_in_stream() | ||
| 245 | + # Ensure NPU storage is properly sized | ||
| 246 | + self.npu_tensor.untyped_storage().resize_( | ||
| 247 | + self.cpu_tensor.untyped_storage().size()) | ||
| 248 | + # Wait for previous swap-out to complete | ||
| 249 | + torch.npu.current_stream().wait_event(self.swap_out_event) | ||
| 250 | + | ||
| 251 | + with torch.npu.stream(swap_stream): | ||
| 252 | + if wait_event: | ||
| 253 | + swap_stream.wait_event(wait_event) | ||
| 254 | + if wait_stream: | ||
| 255 | + swap_stream.wait_stream(wait_stream) | ||
| 256 | + | ||
| 257 | + self.npu_tensor.untyped_storage().copy_( | ||
| 258 | + self.cpu_tensor.untyped_storage(), | ||
| 259 | + non_blocking=True) | ||
| 260 | + | ||
| 261 | + self.swap_in_event = torch.npu.Event() | ||
| 262 | + self.swap_in_event.record() | ||
| 263 | + | ||
| 264 | + def wait_swap_in(self): | ||
| 265 | + """Wait for swap-in to complete and release CPU memory.""" | ||
| 266 | + if self.swap_in_event: | ||
| 267 | + torch.npu.current_stream().wait_event(self.swap_in_event) | ||
| 268 | + # Release CPU memory | ||
| 269 | + self.cpu_tensor = None | ||
| 270 | + self.under_swap_in = False | ||
| 271 | + | ||
| 272 | + | ||
| 273 | + def wait_all_swap_out(cls, group_name): | ||
| 274 | + """ | ||
| 275 | + Wait for all swap-out operations in a group to complete. | ||
| 276 | + | ||
| 277 | + Args: | ||
| 278 | + group_name: Name of the swap group to synchronize | ||
| 279 | + """ | ||
| 280 | + if group_name in cls._ALL_SWAP_OUT_QUEUES: | ||
| 281 | + for manager in cls._ALL_SWAP_OUT_QUEUES[group_name]: | ||
| 282 | + manager.wait_swap_out() | ||
| 283 | + # Clear the group while maintaining the list object | ||
| 284 | + cls._ALL_SWAP_OUT_QUEUES[group_name].clear() | ||
| 285 | + | ||
| 286 | + | ||
| 287 | +def make_wait_swap_in_hook(swap_manager): | ||
| 288 | + """ | ||
| 289 | + Create a hook that waits for a swap-in operation to complete. | ||
| 290 | + | ||
| 291 | + Returns: | ||
| 292 | + A callable hook function that waits for swap-in completion | ||
| 293 | + """ | ||
| 294 | + return lambda *_: swap_manager.wait_swap_in() | ||
| 295 | + | ||
| 296 | + | ||
| 297 | +def make_async_swap_in_hook(swap_managers): | ||
| 298 | + """ | ||
| 299 | + Create a hook that initiates async swap-in for multiple tensors. | ||
| 300 | + | ||
| 301 | + Args: | ||
| 302 | + swap_managers: List of TensorSwapManager instances | ||
| 303 | + | ||
| 304 | + Returns: | ||
| 305 | + A callable hook function that triggers swap-in for all managers | ||
| 306 | + """ | ||
| 307 | + return lambda *_: [m.async_swap_in(wait_stream=torch.npu.current_stream()) | ||
| 308 | + for m in swap_managers] | ||
| @@ -22,6 +22,11 @@ def transformer_layer_backward_moe( | |||
| 22 | dispached_input, probs, indices, global_input_tokens_local_experts_indices = self.recompute_needed_tensors | 22 | dispached_input, probs, indices, global_input_tokens_local_experts_indices = self.recompute_needed_tensors |
| 23 | ep_group = parallel_state.get_expert_model_parallel_group() | 23 | ep_group = parallel_state.get_expert_model_parallel_group() |
| 24 | tp_size = parallel_state.get_tensor_model_parallel_world_size() | 24 | tp_size = parallel_state.get_tensor_model_parallel_world_size() |
| 25 | + | ||
| 26 | + # Launch swap-in at the beginning of the backward pass. | ||
| 27 | + if self.unperm2_swap_manager: | ||
| 28 | + self.unperm2_swap_manager.async_swap_in(wait_stream=torch.npu.current_stream()) | ||
| 29 | + | ||
| 25 | if args.moe_tp_extend_ep: | 30 | if args.moe_tp_extend_ep: |
| 26 | ep_group = parallel_state.get_tensor_and_expert_parallel_group() | 31 | ep_group = parallel_state.get_tensor_and_expert_parallel_group() |
| 27 | if tp_size > 1: | 32 | if tp_size > 1: |
| @@ -37,7 +42,8 @@ def transformer_layer_backward_moe( | |||
| 37 | if backward_ag_shared_handle is not None: | 42 | if backward_ag_shared_handle is not None: |
| 38 | backward_ag_shared_handle.wait() | 43 | backward_ag_shared_handle.wait() |
| 39 | backward_ag_shared_handle = None | 44 | backward_ag_shared_handle = None |
| 40 | - if layer_output_grad is not None: | 45 | + # In case of unperm2 swap, layer_output_grad is required for probs_grad before router-backward |
| 46 | + if layer_output_grad is not None and not args.moe_unperm2_mem_optim_swap: | ||
| 41 | layer_output_grad.untyped_storage().resize_(0) | 47 | layer_output_grad.untyped_storage().resize_(0) |
| 42 | _, unperm1_out_grad, handle = async_all_to_all( | 48 | _, unperm1_out_grad, handle = async_all_to_all( |
| 43 | self.unperm_a2a_graph[1].grad, | 49 | self.unperm_a2a_graph[1].grad, |
| @@ -117,7 +123,22 @@ def transformer_layer_backward_moe( | |||
| 117 | if prob_handle: | 123 | if prob_handle: |
| 118 | prob_handle.wait() | 124 | prob_handle.wait() |
| 119 | run_graph_backward(self.perm1_graph, [perm1_out_grad, perm1_prob_out_grad]) | 125 | run_graph_backward(self.perm1_graph, [perm1_out_grad, perm1_prob_out_grad]) |
| 120 | - run_graph_backward(self.router_graph) | 126 | + |
| 127 | + # Swap-in unperm2 input for probs_grad computation in backward pass of router. | ||
| 128 | + if self.unperm2_swap_manager: | ||
| 129 | + self.unperm2_swap_manager.wait_swap_in() | ||
| 130 | + probs_grad = None | ||
| 131 | + if args.moe_unperm2_mem_optim_swap: | ||
| 132 | + # dprobs computation | ||
| 133 | + H = self.unperm2_swap_manager.npu_tensor.shape[-1] | ||
| 134 | + K = args.moe_router_topk | ||
| 135 | + probs_dtype = probs.dtype | ||
| 136 | + probs_grad = layer_output_grad.to(probs_dtype) * self.unperm2_swap_manager.npu_tensor.reshape(-1, K, H).to(probs_dtype) | ||
| 137 | + probs_grad = probs_grad.sum(dim=-1) | ||
| 138 | + layer_output_grad.untyped_storage().resize_(0) | ||
| 139 | + self.unperm2_swap_manager.npu_tensor.untyped_storage().resize_(0) | ||
| 140 | + run_graph_backward(self.router_graph, probs_grad) | ||
| 141 | + | ||
| 121 | run_graph_backward(self.pre_mlp_layernorm_graph) | 142 | run_graph_backward(self.pre_mlp_layernorm_graph) |
| 122 | run_graph_backward(self.attn_graph) | 143 | run_graph_backward(self.attn_graph) |
| 123 | 144 | ||
| @@ -178,7 +178,7 @@ def transformer_layer_forward_moe( | |||
| 178 | # but backward func of unperm1_out is needed, so resize the storage but keep tensor. | 178 | # but backward func of unperm1_out is needed, so resize the storage but keep tensor. |
| 179 | unperm1_out.untyped_storage().resize_(0) | 179 | unperm1_out.untyped_storage().resize_(0) |
| 180 | detached_unperm_a2a_out = detach_tensor(unperm_a2a_out, checkpoint_forward=checkpoint) | 180 | detached_unperm_a2a_out = detach_tensor(unperm_a2a_out, checkpoint_forward=checkpoint) |
| 181 | - route_expert_output, _ = alltoall_token_unperm2(self.mlp.token_dispatcher, detached_unperm_a2a_out) | 181 | + route_expert_output, unperm2_swap_manager = alltoall_token_unperm2(self.mlp.token_dispatcher, detached_unperm_a2a_out) |
| 182 | if get_args().moe_unperm2_mem_optim: | 182 | if get_args().moe_unperm2_mem_optim: |
| 183 | unperm_a2a_out.untyped_storage().resize_(0) | 183 | unperm_a2a_out.untyped_storage().resize_(0) |
| 184 | 184 | ||
| @@ -191,7 +191,7 @@ def transformer_layer_forward_moe( | |||
| 191 | share_experts_graph = None | 191 | share_experts_graph = None |
| 192 | mlp_output = route_expert_output | 192 | mlp_output = route_expert_output |
| 193 | 193 | ||
| 194 | - if recomp_norm: | 194 | + if recomp_norm and mlp_output.requires_grad: |
| 195 | mlp_output.register_hook(self.norm_ckpt2.recompute) | 195 | mlp_output.register_hook(self.norm_ckpt2.recompute) |
| 196 | 196 | ||
| 197 | 197 | ||
| @@ -232,6 +232,7 @@ def transformer_layer_forward_moe( | |||
| 232 | checkpointed=checkpoint | 232 | checkpointed=checkpoint |
| 233 | ) | 233 | ) |
| 234 | graph.act_ckpt_manager = act_ckpt_manager | 234 | graph.act_ckpt_manager = act_ckpt_manager |
| 235 | + graph.unperm2_swap_manager = unperm2_swap_manager | ||
| 235 | 236 | ||
| 236 | return output, context, graph | 237 | return output, context, graph |
| 237 | 238 | ||
| @@ -283,7 +284,8 @@ def transformer_layer_forward_dense( | |||
| 283 | 284 | ||
| 284 | if recomp_norm: | 285 | if recomp_norm: |
| 285 | self.norm_ckpt2.discard_output() | 286 | self.norm_ckpt2.discard_output() |
| 286 | - mlp_output_with_bias[0].register_hook(self.norm_ckpt2.recompute) | 287 | + if mlp_output_with_bias[0].requires_grad: |
| 288 | + mlp_output_with_bias[0].register_hook(self.norm_ckpt2.recompute) | ||
| 287 | 289 | ||
| 288 | 290 | ||
| 289 | with self.bias_dropout_add_exec_handler(): | 291 | with self.bias_dropout_add_exec_handler(): |
| @@ -63,6 +63,10 @@ def transformer_layer_forward_dense_backward_moe_overlaping( | |||
| 63 | recomp_norm = getattr(args, 'recompute_norm', False) | 63 | recomp_norm = getattr(args, 'recompute_norm', False) |
| 64 | bwd_dispached_input, bwd_probs, bwd_indices, global_input_tokens_local_experts_indices = bwd_layer_graph.recompute_needed_tensors | 64 | bwd_dispached_input, bwd_probs, bwd_indices, global_input_tokens_local_experts_indices = bwd_layer_graph.recompute_needed_tensors |
| 65 | 65 | ||
| 66 | + # Launch swap-in at the beginning of the backward pass. | ||
| 67 | + if bwd_layer_graph.unperm2_swap_manager: | ||
| 68 | + bwd_layer_graph.unperm2_swap_manager.async_swap_in(wait_stream=torch.npu.current_stream()) | ||
| 69 | + | ||
| 66 | # Unperm2 Bwd | 70 | # Unperm2 Bwd |
| 67 | # check if backward unpermutation alltoall is launched at bwd layer before | 71 | # check if backward unpermutation alltoall is launched at bwd layer before |
| 68 | if bwd_unperm_a2a_handle is None: | 72 | if bwd_unperm_a2a_handle is None: |
| @@ -200,7 +204,25 @@ def transformer_layer_forward_dense_backward_moe_overlaping( | |||
| 200 | shared_expert_grad.untyped_storage().resize_(0) | 204 | shared_expert_grad.untyped_storage().resize_(0) |
| 201 | run_graph_backward(bwd_layer_graph.shared_experts_graph, backward_ag_shared, keep_grad=True) # dw computation | 205 | run_graph_backward(bwd_layer_graph.shared_experts_graph, backward_ag_shared, keep_grad=True) # dw computation |
| 202 | WeightGradStore.end_decouple() | 206 | WeightGradStore.end_decouple() |
| 203 | - run_graph_backward(bwd_layer_graph.router_graph) | 207 | + |
| 208 | + # swap-in unperm2 input for probs_grad computation | ||
| 209 | + if bwd_layer_graph.unperm2_swap_manager: | ||
| 210 | + bwd_layer_graph.unperm2_swap_manager.wait_swap_in() | ||
| 211 | + probs_grad = None | ||
| 212 | + if args.moe_unperm2_mem_optim_swap: | ||
| 213 | + # dprobs computation | ||
| 214 | + output_grad = bwd_layer_output_grad | ||
| 215 | + if hasattr(bwd_layer_graph, 'last_layer_input_grad'): | ||
| 216 | + output_grad = bwd_layer_graph.last_layer_input_grad | ||
| 217 | + H = bwd_layer_graph.unperm2_swap_manager.npu_tensor.shape[-1] | ||
| 218 | + K = args.moe_router_topk | ||
| 219 | + probs_dtype = bwd_probs.dtype | ||
| 220 | + probs_grad = bwd_layer_graph.unperm2_swap_manager.npu_tensor.reshape(-1, K, H).to(probs_dtype) * output_grad.to(probs_dtype) | ||
| 221 | + output_grad.untyped_storage().resize_(0) | ||
| 222 | + bwd_layer_graph.unperm2_swap_manager.npu_tensor.untyped_storage().resize_(0) | ||
| 223 | + probs_grad = probs_grad.sum(dim=-1) | ||
| 224 | + run_graph_backward(bwd_layer_graph.router_graph, probs_grad) | ||
| 225 | + | ||
| 204 | run_graph_backward(bwd_layer_graph.pre_mlp_layernorm_graph, keep_graph=True) | 226 | run_graph_backward(bwd_layer_graph.pre_mlp_layernorm_graph, keep_graph=True) |
| 205 | WeightGradStore.start_decouple() | 227 | WeightGradStore.start_decouple() |
| 206 | run_graph_backward(bwd_layer_graph.attn_graph, keep_grad=True) | 228 | run_graph_backward(bwd_layer_graph.attn_graph, keep_grad=True) |
| @@ -272,6 +294,13 @@ def transformer_layer_forward_dense_backward_moe_overlaping( | |||
| 272 | checkpointed=checkpoint | 294 | checkpointed=checkpoint |
| 273 | ) | 295 | ) |
| 274 | 296 | ||
| 297 | + # save original layer output for probs_grad computation | ||
| 298 | + if args.moe_unperm2_mem_optim_swap \ | ||
| 299 | + and next_bwd_layer_graph is not None \ | ||
| 300 | + and getattr(next_bwd_layer_graph, 'is_moe_layer', False): | ||
| 301 | + | ||
| 302 | + next_bwd_layer_graph.last_layer_input_grad = bwd_layer_graph.layer_input.grad | ||
| 303 | + | ||
| 275 | for tensor in bwd_layer_graph.recompute_needed_tensors: | 304 | for tensor in bwd_layer_graph.recompute_needed_tensors: |
| 276 | if tensor is not None: | 305 | if tensor is not None: |
| 277 | tensor.untyped_storage().resize_(0) | 306 | tensor.untyped_storage().resize_(0) |
| @@ -457,7 +486,7 @@ def transformer_layer_forward_moe_backward_dense_overlaping( | |||
| 457 | 486 | ||
| 458 | with checkpoint_context: | 487 | with checkpoint_context: |
| 459 | detached_unperm_a2a_out = detach_tensor(unperm_a2a_out, checkpoint_forward=checkpoint) | 488 | detached_unperm_a2a_out = detach_tensor(unperm_a2a_out, checkpoint_forward=checkpoint) |
| 460 | - route_expert_output, _ = alltoall_token_unperm2(fwd_layer.mlp.token_dispatcher, detached_unperm_a2a_out) | 489 | + route_expert_output, unperm2_swap_manager = alltoall_token_unperm2(fwd_layer.mlp.token_dispatcher, detached_unperm_a2a_out) |
| 461 | if args.moe_unperm2_mem_optim: | 490 | if args.moe_unperm2_mem_optim: |
| 462 | unperm_a2a_out.untyped_storage().resize_(0) | 491 | unperm_a2a_out.untyped_storage().resize_(0) |
| 463 | 492 | ||
| @@ -522,6 +551,7 @@ def transformer_layer_forward_moe_backward_dense_overlaping( | |||
| 522 | checkpointed=checkpoint | 551 | checkpointed=checkpoint |
| 523 | ) | 552 | ) |
| 524 | graph.act_ckpt_manager = act_ckpt_manager | 553 | graph.act_ckpt_manager = act_ckpt_manager |
| 554 | + graph.unperm2_swap_manager = unperm2_swap_manager | ||
| 525 | 555 | ||
| 526 | for tensor in bwd_layer_graph.recompute_needed_tensors: | 556 | for tensor in bwd_layer_graph.recompute_needed_tensors: |
| 527 | if tensor is not None: | 557 | if tensor is not None: |
| @@ -702,13 +732,18 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 702 | tp_group = parallel_state.get_tensor_model_parallel_group() | 732 | tp_group = parallel_state.get_tensor_model_parallel_group() |
| 703 | use_shared_experts = hasattr(fwd_layer.mlp, 'shared_experts') and fwd_layer.mlp.shared_experts is not None | 733 | use_shared_experts = hasattr(fwd_layer.mlp, 'shared_experts') and fwd_layer.mlp.shared_experts is not None |
| 704 | recomp_norm = getattr(args, 'recompute_norm', False) | 734 | recomp_norm = getattr(args, 'recompute_norm', False) |
| 735 | + swap_unperm2 = getattr(args, 'moe_unperm2_mem_optim_swap', False) | ||
| 705 | bwd_dispached_input, bwd_probs, bwd_indices, global_input_tokens_local_experts_indices = bwd_layer_graph.recompute_needed_tensors | 736 | bwd_dispached_input, bwd_probs, bwd_indices, global_input_tokens_local_experts_indices = bwd_layer_graph.recompute_needed_tensors |
| 706 | a2a_hooked_on_attention = getattr(fwd_layer.self_attention, 'a2a_hooked_on_attention', False) | 737 | a2a_hooked_on_attention = getattr(fwd_layer.self_attention, 'a2a_hooked_on_attention', False) |
| 707 | 738 | ||
| 739 | + # Launch swap-in | ||
| 740 | + if bwd_layer_graph.unperm2_swap_manager: | ||
| 741 | + bwd_layer_graph.unperm2_swap_manager.async_swap_in(wait_stream=torch.npu.current_stream()) | ||
| 742 | + | ||
| 708 | # Unperm2 Bwd | 743 | # Unperm2 Bwd |
| 709 | # check if backward unpermutation alltoall is launched at bwd layer before | 744 | # check if backward unpermutation alltoall is launched at bwd layer before |
| 710 | if bwd_unperm_a2a_handle is None: | 745 | if bwd_unperm_a2a_handle is None: |
| 711 | - run_graph_backward(bwd_layer_graph.unperm2_graph, bwd_layer_output_grad) | 746 | + run_graph_backward(bwd_layer_graph.unperm2_graph, bwd_layer_output_grad, keep_grad=swap_unperm2) |
| 712 | # Async Unperm A2A | 747 | # Async Unperm A2A |
| 713 | if tp_size > 1 and a2a_hooked_on_attention: | 748 | if tp_size > 1 and a2a_hooked_on_attention: |
| 714 | set_async_alltoall_inputs( | 749 | set_async_alltoall_inputs( |
| @@ -959,7 +994,25 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 959 | bwd_prob_handle.wait() | 994 | bwd_prob_handle.wait() |
| 960 | run_graph_backward(bwd_layer_graph.perm1_graph, [perm1_out_grad, perm1_prob_out_grad]) | 995 | run_graph_backward(bwd_layer_graph.perm1_graph, [perm1_out_grad, perm1_prob_out_grad]) |
| 961 | perm1_out_grad.untyped_storage().resize_(0) | 996 | perm1_out_grad.untyped_storage().resize_(0) |
| 962 | - run_graph_backward(bwd_layer_graph.router_graph) | 997 | + |
| 998 | + # router backward | ||
| 999 | + if bwd_layer_graph.unperm2_swap_manager: | ||
| 1000 | + bwd_layer_graph.unperm2_swap_manager.wait_swap_in() | ||
| 1001 | + probs_grad = None | ||
| 1002 | + if swap_unperm2: | ||
| 1003 | + # dprobs computation | ||
| 1004 | + output_grad = bwd_layer_output_grad | ||
| 1005 | + if hasattr(bwd_layer_graph, 'last_layer_input_grad'): | ||
| 1006 | + output_grad = bwd_layer_graph.last_layer_input_grad | ||
| 1007 | + H = bwd_layer_graph.unperm2_swap_manager.npu_tensor.shape[-1] | ||
| 1008 | + K = args.moe_router_topk | ||
| 1009 | + probs_dtype = bwd_probs.dtype | ||
| 1010 | + probs_grad = bwd_layer_graph.unperm2_swap_manager.npu_tensor.reshape(-1, K, H).to(probs_dtype) * output_grad.to(probs_dtype) | ||
| 1011 | + output_grad.untyped_storage().resize_(0) | ||
| 1012 | + bwd_layer_graph.unperm2_swap_manager.npu_tensor.untyped_storage().resize_(0) | ||
| 1013 | + probs_grad = probs_grad.sum(dim=-1) | ||
| 1014 | + run_graph_backward(bwd_layer_graph.router_graph, probs_grad) | ||
| 1015 | + | ||
| 963 | run_graph_backward(bwd_layer_graph.pre_mlp_layernorm_graph, keep_graph=True) | 1016 | run_graph_backward(bwd_layer_graph.pre_mlp_layernorm_graph, keep_graph=True) |
| 964 | WeightGradStore.start_decouple() | 1017 | WeightGradStore.start_decouple() |
| 965 | run_graph_backward(bwd_layer_graph.attn_graph, keep_grad=True) | 1018 | run_graph_backward(bwd_layer_graph.attn_graph, keep_grad=True) |
| @@ -968,7 +1021,7 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 968 | unperm_a2a_out, unperm_a2a_handle = get_async_alltoall_outputs() | 1021 | unperm_a2a_out, unperm_a2a_handle = get_async_alltoall_outputs() |
| 969 | 1022 | ||
| 970 | if next_bwd_layer_graph is not None and getattr(next_bwd_layer_graph, 'is_moe_layer', False): | 1023 | if next_bwd_layer_graph is not None and getattr(next_bwd_layer_graph, 'is_moe_layer', False): |
| 971 | - run_graph_backward(next_bwd_layer_graph.unperm2_graph, bwd_layer_graph.layer_input.grad, keep_graph=True) | 1024 | + run_graph_backward(next_bwd_layer_graph.unperm2_graph, bwd_layer_graph.layer_input.grad, keep_graph=True, keep_grad=swap_unperm2) |
| 972 | 1025 | ||
| 973 | unperm_a2a_handle.wait() | 1026 | unperm_a2a_handle.wait() |
| 974 | unperm_a2a_handle = None | 1027 | unperm_a2a_handle = None |
| @@ -984,7 +1037,7 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 984 | ) | 1037 | ) |
| 985 | with checkpoint_context: | 1038 | with checkpoint_context: |
| 986 | detached_unperm_a2a_out = detach_tensor(unperm_a2a_out) | 1039 | detached_unperm_a2a_out = detach_tensor(unperm_a2a_out) |
| 987 | - route_expert_output, _ = alltoall_token_unperm2(fwd_layer.mlp.token_dispatcher, detached_unperm_a2a_out) | 1040 | + route_expert_output, unperm2_swap_manager = alltoall_token_unperm2(fwd_layer.mlp.token_dispatcher, detached_unperm_a2a_out) |
| 988 | if args.moe_unperm2_mem_optim: | 1041 | if args.moe_unperm2_mem_optim: |
| 989 | unperm_a2a_out.untyped_storage().resize_(0) | 1042 | unperm_a2a_out.untyped_storage().resize_(0) |
| 990 | 1043 | ||
| @@ -1050,6 +1103,14 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 1050 | checkpointed=checkpoint | 1103 | checkpointed=checkpoint |
| 1051 | ) | 1104 | ) |
| 1052 | graph.act_ckpt_manager = act_ckpt_manager | 1105 | graph.act_ckpt_manager = act_ckpt_manager |
| 1106 | + graph.unperm2_swap_manager = unperm2_swap_manager | ||
| 1107 | + | ||
| 1108 | + # save original layer output for probs_grad computation | ||
| 1109 | + if swap_unperm2 \ | ||
| 1110 | + and next_bwd_layer_graph is not None \ | ||
| 1111 | + and getattr(next_bwd_layer_graph, 'is_moe_layer', False): | ||
| 1112 | + | ||
| 1113 | + next_bwd_layer_graph.last_layer_input_grad = bwd_layer_graph.layer_input.grad | ||
| 1053 | 1114 | ||
| 1054 | for tensor in bwd_layer_graph.recompute_needed_tensors: | 1115 | for tensor in bwd_layer_graph.recompute_needed_tensors: |
| 1055 | if tensor is not None: | 1116 | if tensor is not None: |
| @@ -13,6 +13,7 @@ class FwdBwdOverlapFeature(MindSpeedFeature): | |||
| 13 | group = parser.add_argument_group(title=self.feature_name) | 13 | group = parser.add_argument_group(title=self.feature_name) |
| 14 | group.add_argument('--moe-fb-overlap', action='store_true') | 14 | group.add_argument('--moe-fb-overlap', action='store_true') |
| 15 | group.add_argument('--moe-unperm2-mem-optim', action='store_true') | 15 | group.add_argument('--moe-unperm2-mem-optim', action='store_true') |
| 16 | + group.add_argument('--moe-unperm2-mem-optim-swap', action='store_true') | ||
| 16 | 17 | ||
| 17 | 18 | ||
| 18 | 19 | ||
| @@ -28,9 +29,9 @@ class FwdBwdOverlapFeature(MindSpeedFeature): | |||
| 28 | self.incompatible_check(args, 'recompute_in_advance') | 29 | self.incompatible_check(args, 'recompute_in_advance') |
| 29 | self.incompatible_check(args, 'use_legacy_models') | 30 | self.incompatible_check(args, 'use_legacy_models') |
| 30 | if args.moe_fb_overlap and args.moe_token_dispatcher_type == 'allgather': | 31 | if args.moe_fb_overlap and args.moe_token_dispatcher_type == 'allgather': |
| 31 | - raise AssertionError('The fb overlap feature do not support allgather dispatcher') | 32 | + raise AssertionError('The fb overlap feature do not support allgather dispatcher.') |
| 32 | if args.moe_fb_overlap and args.moe_zero_memory == 'level1': | 33 | if args.moe_fb_overlap and args.moe_zero_memory == 'level1': |
| 33 | - raise AssertionError('fb overlap only support moe zero memory level 0') | 34 | + raise AssertionError('fb overlap only support moe zero memory level 0.') |
| 34 | 35 | ||
| 35 | self.dependency_check(args, 'n_shared_experts') | 36 | self.dependency_check(args, 'n_shared_experts') |
| 36 | self.dependency_check(args, 'moe_permutation_async_comm') | 37 | self.dependency_check(args, 'moe_permutation_async_comm') |
| @@ -38,6 +39,10 @@ class FwdBwdOverlapFeature(MindSpeedFeature): | |||
| 38 | self.dependency_check(args, 'moe_tp_extend_ep') | 39 | self.dependency_check(args, 'moe_tp_extend_ep') |
| 39 | if args.moe_unperm2_mem_optim and not args.moe_fb_overlap: | 40 | if args.moe_unperm2_mem_optim and not args.moe_fb_overlap: |
Z 这句没意义，上面三行应该回退掉？   | |||
| 40 | raise AssertionError('--moe-unperm2-mem-optim currently only can be used with --moe-fb-overlap') | 41 | raise AssertionError('--moe-unperm2-mem-optim currently only can be used with --moe-fb-overlap') |
| 42 | + if args.moe_unperm2_mem_optim_swap and not args.moe_fb_overlap: | ||
| 43 | + raise AssertionError('--moe-unperm2-mem-optim-swap currently only can be used with --moe-fb-overlap') | ||
| 44 | + if args.moe_unperm2_mem_optim and args.moe_unperm2_mem_optim_swap: | ||
| 45 | + raise AssertionError('--moe-unperm2-mem-optim and --moe-unperm2-mem-optim-swap are incompatible.') | ||
| 41 | 46 | ||
| 42 | 47 | ||
| 43 | def register_patches(self, patch_manager, args): | 48 | def register_patches(self, patch_manager, args): |
| @@ -0,0 +1,126 @@ | |||
| 1 | + | ||
| 2 | +#!/bin/bash | ||
| 3 | + | ||
| 4 | +export CUDA_DEVICE_MAX_CONNECTIONS=1 | ||
| 5 | +source "tests_extend/system_tests/env_npu.sh" | ||
| 6 | + | ||
| 7 | +NPUS_PER_NODE=16 | ||
| 8 | +MASTER_ADDR=localhost | ||
| 9 | +MASTER_PORT=6001 | ||
| 10 | +NNODES=1 | ||
| 11 | +NODE_RANK=0 | ||
| 12 | +WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES)) | ||
| 13 | + | ||
| 14 | +CKPT_DIR=./ckpt_llama | ||
| 15 | +DATA_PATH="/home/dataset/llama2/alpaca_text_document" | ||
| 16 | +TOKENIZER_MODEL="/home/dataset/model/llama-2-7b-hf/tokenizer.model" | ||
| 17 | + | ||
| 18 | +TP=2 | ||
| 19 | +PP=4 | ||
| 20 | +EP=2 | ||
| 21 | + | ||
| 22 | +DISTRIBUTED_ARGS=" | ||
| 23 | + --nproc_per_node $NPUS_PER_NODE \ | ||
| 24 | + --nnodes $NNODES \ | ||
| 25 | + --node_rank $NODE_RANK \ | ||
| 26 | + --master_addr $MASTER_ADDR \ | ||
| 27 | + --master_port $MASTER_PORT | ||
| 28 | +" | ||
| 29 | + | ||
| 30 | +DUALPIPE_ARGS=" | ||
| 31 | + --moe-unperm2-mem-optim-swap \ | ||
| 32 | + --moe-fb-overlap \ | ||
| 33 | + --schedules-method dualpipev \ | ||
| 34 | + --moe-zero-memory level0 \ | ||
| 35 | + --moe-tp-extend-ep \ | ||
| 36 | +" | ||
| 37 | + | ||
| 38 | +MOE_ARGS=" | ||
| 39 | + --expert-model-parallel-size ${EP} \ | ||
| 40 | + --moe-token-dispatcher-type alltoall \ | ||
| 41 | + --moe-permutation-async-comm \ | ||
| 42 | + -use-fused-moe-token-permute-and-unpermute \ | ||
| 43 | + --moe-grouped-gemm \ | ||
| 44 | + --n-shared-experts 1 \ | ||
| 45 | + --num-experts 32 \ | ||
| 46 | + --moe-router-topk 8 \ | ||
| 47 | + --moe-aux-loss-coeff 0.02 \ | ||
| 48 | + --moe-tp-extend-ep \ | ||
| 49 | + --recompute-norm \ | ||
| 50 | + --recompute-activation-function \ | ||
| 51 | +" | ||
| 52 | + | ||
| 53 | +GPT_ARGS=" | ||
| 54 | + --tensor-model-parallel-size ${TP} \ | ||
| 55 | + --pipeline-model-parallel-size ${PP} \ | ||
| 56 | + --use-flash-attn \ | ||
| 57 | + --use-fused-rotary-pos-emb \ | ||
| 58 | + --use-fused-swiglu \ | ||
| 59 | + --use-fused-rmsnorm \ | ||
| 60 | + --sequence-parallel \ | ||
| 61 | + --use-distributed-optimizer \ | ||
| 62 | + --num-layers 8 \ | ||
| 63 | + --noop-layers 7 \ | ||
| 64 | + --manual-gc \ | ||
| 65 | + --manual-gc-interval 50 \ | ||
| 66 | + --seq-length 4096 \ | ||
| 67 | + --max-position-embeddings 4096 \ | ||
| 68 | + --train-iters 50 \ | ||
| 69 | + --hidden-size 5120 \ | ||
| 70 | + --num-attention-heads 128 \ | ||
| 71 | + --ffn-hidden-size 1536 \ | ||
| 72 | + --make-vocab-size-divisible-by 128 \ | ||
| 73 | + --vocab-size 126464 \ | ||
| 74 | + --micro-batch-size 1 \ | ||
| 75 | + --global-batch-size 32 \ | ||
| 76 | + --tokenizer-type Llama2Tokenizer \ | ||
| 77 | + --tokenizer-model ${TOKENIZER_MODEL} \ | ||
| 78 | + --disable-bias-linear \ | ||
| 79 | + --lr-decay-style linear \ | ||
| 80 | + --lr-warmup-iters 0 \ | ||
| 81 | + --short-seq-prob 0.0 \ | ||
| 82 | + --attention-dropout 0.0 \ | ||
| 83 | + --hidden-dropout 0.0 \ | ||
| 84 | + --untie-embeddings-and-output-weights \ | ||
| 85 | + --init-method-std 0.006 \ | ||
| 86 | + --position-embedding-type rope \ | ||
| 87 | + --normalization RMSNorm \ | ||
| 88 | + --swiglu \ | ||
| 89 | + --no-masked-softmax-fusion \ | ||
| 90 | + --attention-softmax-in-fp32 \ | ||
| 91 | + --no-gradient-accumulation-fusion \ | ||
| 92 | + --bf16 \ | ||
| 93 | + --group-query-attention \ | ||
| 94 | + --num-query-groups 8 \ | ||
| 95 | + --lr 2.0e-4 \ | ||
| 96 | + --min-lr 2.0e-4 \ | ||
| 97 | + --weight-decay 0.1 \ | ||
| 98 | + --clip-grad 1.0 \ | ||
| 99 | + --adam-beta1 0.9 \ | ||
| 100 | + --adam-beta2 0.95 \ | ||
| 101 | + --rotary-base 100000 \ | ||
| 102 | + --norm-epsilon 1.0e-5 \ | ||
| 103 | +" | ||
| 104 | + | ||
| 105 | +DATA_ARGS=" | ||
| 106 | + --data-path $DATA_PATH \ | ||
| 107 | + --split 995,5,0 | ||
| 108 | +" | ||
| 109 | + | ||
| 110 | +OUTPUT_ARGS=" | ||
| 111 | + --log-throughput \ | ||
| 112 | + --log-interval 1 \ | ||
| 113 | + --save-interval 10000 \ | ||
| 114 | + --eval-interval 10000 \ | ||
| 115 | + --eval-iters 0 \ | ||
| 116 | +" | ||
| 117 | + | ||
| 118 | +torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \ | ||
| 119 | + --npu-deterministic \ | ||
| 120 | + $GPT_ARGS \ | ||
| 121 | + $MOE_ARGS \ | ||
| 122 | + $DATA_ARGS \ | ||
| 123 | + $OUTPUT_ARGS \ | ||
| 124 | + $DUALPIPE_ARGS \ | ||
| 125 | + 2>&1 | tee -i base-deter-nofuse.log | ||
| 126 | + | ||
| @@ -0,0 +1,166 @@ | |||
| 1 | +import pytest | ||
| 2 | +import unittest | ||
| 3 | +from unittest.mock import patch, MagicMock | ||
| 4 | + | ||
| 5 | +import torch | ||
| 6 | +import torch_npu | ||
| 7 | + | ||
| 8 | +from mindspeed import megatron_adaptor | ||
| 9 | +from mindspeed.core.pipeline_parallel.fb_overlap.modules.utils import ( | ||
| 10 | + TensorSwapManager, | ||
| 11 | + make_wait_swap_in_hook, | ||
| 12 | + make_async_swap_in_hook | ||
| 13 | + ) | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +class TestTensorSwapManager: | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + def setup_class(cls): | ||
| 20 | + # Ensure we have a device to test with | ||
| 21 | + if not torch_npu.npu.is_available(): | ||
| 22 | + pytest.skip("NPU device not available") | ||
| 23 | + | ||
| 24 | + def test_initialization(self): | ||
| 25 | + """Test basic initialization of TensorSwapManager.""" | ||
| 26 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 27 | + manager = TensorSwapManager(test_tensor) | ||
| 28 | + | ||
| 29 | + assert manager.npu_tensor is test_tensor | ||
| 30 | + assert manager.cpu_tensor is None | ||
| 31 | + assert manager.swap_out_event is None | ||
| 32 | + assert manager.swap_in_event is None | ||
| 33 | + assert not manager.under_swap_in | ||
| 34 | + | ||
| 35 | + def test_swap_group_initialization(self): | ||
| 36 | + """Test initialization with swap groups.""" | ||
| 37 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 38 | + group_name = "test_group" | ||
| 39 | + | ||
| 40 | + # First manager should create the group | ||
| 41 | + manager1 = TensorSwapManager(test_tensor, group_name) | ||
| 42 | + assert group_name in TensorSwapManager._ALL_SWAP_OUT_QUEUES | ||
| 43 | + assert len(TensorSwapManager._ALL_SWAP_OUT_QUEUES[group_name]) == 0 | ||
| 44 | + | ||
| 45 | + # Second manager should use existing group | ||
| 46 | + manager2 = TensorSwapManager(test_tensor, group_name) | ||
| 47 | + assert len(TensorSwapManager._ALL_SWAP_OUT_QUEUES[group_name]) == 0 | ||
| 48 | + | ||
| 49 | + def test_async_swap_out(self): | ||
| 50 | + """Test asynchronous swap-out functionality.""" | ||
| 51 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 52 | + manager = TensorSwapManager(test_tensor) | ||
| 53 | + | ||
| 54 | + # Perform async swap-out | ||
| 55 | + manager.async_swap_out() | ||
| 56 | + | ||
| 57 | + # Verify swap-out was initiated | ||
| 58 | + assert manager.cpu_tensor is not None | ||
| 59 | + assert manager.swap_out_event is not None | ||
| 60 | + assert manager.cpu_tensor.device.type == 'cpu' | ||
| 61 | + assert manager.cpu_tensor.is_pinned() | ||
| 62 | + | ||
| 63 | + def test_async_swap_out_wait_conditions(self): | ||
| 64 | + """Test swap-out with wait conditions.""" | ||
| 65 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 66 | + manager = TensorSwapManager(test_tensor) | ||
| 67 | + | ||
| 68 | + # Create dummy event and stream | ||
| 69 | + dummy_event = torch.npu.Event() | ||
| 70 | + dummy_event.record() | ||
| 71 | + dummy_stream = torch.npu.Stream() | ||
| 72 | + | ||
| 73 | + # Should complete without errors | ||
| 74 | + manager.async_swap_out(wait_event=dummy_event, wait_stream=dummy_stream) | ||
| 75 | + | ||
| 76 | + def test_wait_swap_out(self): | ||
| 77 | + """Test waiting for swap-out completion.""" | ||
| 78 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 79 | + manager = TensorSwapManager(test_tensor) | ||
| 80 | + | ||
| 81 | + manager.async_swap_out() | ||
| 82 | + manager.wait_swap_out() | ||
| 83 | + | ||
| 84 | + # Verify NPU storage was released | ||
| 85 | + assert manager.npu_tensor.untyped_storage().size() == 0 | ||
| 86 | + | ||
| 87 | + def test_async_swap_in(self): | ||
| 88 | + """Test asynchronous swap-in functionality.""" | ||
| 89 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 90 | + manager = TensorSwapManager(test_tensor) | ||
| 91 | + | ||
| 92 | + # First swap out | ||
| 93 | + manager.async_swap_out() | ||
| 94 | + manager.wait_swap_out() | ||
| 95 | + | ||
| 96 | + # Then swap back in | ||
| 97 | + manager.async_swap_in() | ||
| 98 | + | ||
| 99 | + # Verify swap-in was initiated | ||
| 100 | + assert manager.swap_in_event is not None | ||
| 101 | + assert manager.under_swap_in | ||
| 102 | + | ||
| 103 | + def test_wait_swap_in(self): | ||
| 104 | + """Test waiting for swap-in completion.""" | ||
| 105 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 106 | + manager = TensorSwapManager(test_tensor) | ||
| 107 | + | ||
| 108 | + # Full swap cycle | ||
| 109 | + manager.async_swap_out() | ||
| 110 | + manager.wait_swap_out() | ||
| 111 | + manager.async_swap_in() | ||
| 112 | + manager.wait_swap_in() | ||
| 113 | + | ||
| 114 | + # Verify state after swap-in | ||
| 115 | + assert manager.cpu_tensor is None | ||
| 116 | + assert not manager.under_swap_in | ||
| 117 | + assert manager.npu_tensor.untyped_storage().size() > 0 | ||
| 118 | + | ||
| 119 | + def test_wait_all_swap_out(self): | ||
| 120 | + """Test waiting for all swaps in a group.""" | ||
| 121 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 122 | + group_name = "group_test" | ||
| 123 | + | ||
| 124 | + # Create two managers in the same group | ||
| 125 | + manager1 = TensorSwapManager(test_tensor, group_name) | ||
| 126 | + manager2 = TensorSwapManager(test_tensor.clone(), group_name) | ||
| 127 | + | ||
| 128 | + # Initiate swaps | ||
| 129 | + manager1.async_swap_out() | ||
| 130 | + manager2.async_swap_out() | ||
| 131 | + | ||
| 132 | + # Wait for all in group | ||
| 133 | + TensorSwapManager.wait_all_swap_out(group_name) | ||
| 134 | + | ||
| 135 | + # Verify both were processed | ||
| 136 | + assert manager1.npu_tensor.untyped_storage().size() == 0 | ||
| 137 | + assert manager2.npu_tensor.untyped_storage().size() == 0 | ||
| 138 | + assert len(TensorSwapManager._ALL_SWAP_OUT_QUEUES[group_name]) == 0 | ||
| 139 | + | ||
| 140 | + def test_hook_functions(self): | ||
| 141 | + """Test the hook creation functions.""" | ||
| 142 | + test_tensor = torch.randn(5, 5, device='npu') | ||
| 143 | + manager = TensorSwapManager(test_tensor) | ||
| 144 | + | ||
| 145 | + # Test wait swap-in hook | ||
| 146 | + wait_hook = make_wait_swap_in_hook(manager) | ||
| 147 | + manager.async_swap_out() | ||
| 148 | + manager.wait_swap_out() | ||
| 149 | + manager.async_swap_in() | ||
| 150 | + wait_hook() # Should wait for swap-in | ||
| 151 | + | ||
| 152 | + # Test async swap-in hook | ||
| 153 | + managers = [TensorSwapManager(test_tensor.clone()) for _ in range(3)] | ||
| 154 | + async_hook = make_async_swap_in_hook(managers) | ||
| 155 | + | ||
| 156 | + # First swap out all | ||
| 157 | + for m in managers: | ||
| 158 | + m.async_swap_out() | ||
| 159 | + m.wait_swap_out() | ||
| 160 | + | ||
| 161 | + # Trigger swap-in | ||
| 162 | + async_hook() | ||
| 163 | + | ||
| 164 | + # Verify all initiated swap-in | ||
| 165 | + for m in managers: | ||
| 166 | + assert m.swap_in_event is not None | ||
copyright