已合并
Optimize the patch for FSDP #35275
zhenyu创建于 5月11日
Optimize the patch for FSDP #35275
已合并
共 6 个文件变更+181-237
| @@ -182,6 +182,10 @@ select = [ | |||
| 182 | ] | 182 | ] |
| 183 | 183 | ||
| 184 | [tool.ruff.lint.per-file-ignores] | 184 | [tool.ruff.lint.per-file-ignores] |
| 185 | +"torch_npu/distributed/fsdp/_add_fsdp_patch.py" = [ | ||
| 186 | + "UP007", | ||
| 187 | + "UP045", | ||
| 188 | +] | ||
| 185 | "__init__.py" = [ | 189 | "__init__.py" = [ |
| 186 | "F401", | 190 | "F401", |
| 187 | ] | 191 | ] |
| @@ -1,88 +0,0 @@ | |||
| 1 | -import logging | ||
| 2 | -from unittest.mock import patch | ||
| 3 | - | ||
| 4 | -import torch | ||
| 5 | -import torch.nn as nn | ||
| 6 | -from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState | ||
| 7 | -from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState | ||
| 8 | -from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup, AllGatherState | ||
| 9 | -from torch.distributed.utils import _to_kwargs | ||
| 10 | -from torch_npu.distributed.fsdp._add_fsdp_patch import _patched_finalize_backward | ||
| 11 | -import torch_npu.distributed.fsdp._add_fsdp_patch as add_fsdp_patch | ||
| 12 | -from torch_npu.testing.testcase import TestCase, run_tests | ||
| 13 | - | ||
| 14 | - | ||
| 15 | -class TestAddFsdpPatch(TestCase): | ||
| 16 | - def test_get_param_all_gather_inputs_compiled_autograd(self): | ||
| 17 | - class MockFSDPParam: | ||
| 18 | - def __init__(self): | ||
| 19 | - self.all_gather_inputs = [torch.tensor([3.0, 4.0])] | ||
| 20 | - self.param_dtype = torch.float32 | ||
| 21 | - self.offload_to_cpu = False | ||
| 22 | - self._sharded_local_tensor = torch.tensor([1.0, 2.0]) | ||
| 23 | - self.sharded_state = ShardedState.SHARDED | ||
| 24 | - self._sharded_param_data = torch.tensor([3.0, 4.0]) | ||
| 25 | - self._sharded_post_forward_param_data = torch.tensor([5.0, 6.0]) | ||
| 26 | - self.device = torch.device("cpu") | ||
| 27 | - | ||
| 28 | - fsdp_param = MockFSDPParam() | ||
| 29 | - fsdp_params = [fsdp_param] | ||
| 30 | - pass | ||
| 31 | - | ||
| 32 | - def test_patched_finalize_backward_with_events(self): | ||
| 33 | - class MockFSDPParamGroup: | ||
| 34 | - def __init__(self): | ||
| 35 | - self.fsdp_params = [] | ||
| 36 | - self._all_gather_result = MockAllGatherResult() | ||
| 37 | - self._post_forward_indices = [1, 2, 3] | ||
| 38 | - | ||
| 39 | - def _wait_for_post_backward(self): | ||
| 40 | - pass | ||
| 41 | - | ||
| 42 | - class MockAllGatherResult: | ||
| 43 | - def __init__(self): | ||
| 44 | - self.all_gather_event = MockEvent() | ||
| 45 | - self.all_gather_work = MockWork() | ||
| 46 | - | ||
| 47 | - class MockEvent: | ||
| 48 | - def synchronize(self): | ||
| 49 | - pass | ||
| 50 | - | ||
| 51 | - def wait(self, *args): | ||
| 52 | - pass | ||
| 53 | - | ||
| 54 | - class MockWork: | ||
| 55 | - def wait(self): | ||
| 56 | - pass | ||
| 57 | - | ||
| 58 | - class MockFSDPParam: | ||
| 59 | - def __init__(self): | ||
| 60 | - self.grad_offload_event = MockEvent() | ||
| 61 | - | ||
| 62 | - mock_group = MockFSDPParamGroup() | ||
| 63 | - mock_group.fsdp_params = [MockFSDPParam()] | ||
| 64 | - | ||
| 65 | - _patched_finalize_backward(mock_group) | ||
| 66 | - | ||
| 67 | - self.assertIsNone(mock_group._all_gather_result) | ||
| 68 | - self.assertEqual(len(mock_group._post_forward_indices), 0) | ||
| 69 | - | ||
| 70 | - def test_get_param_all_gather_inputs_no_foreach_copy(self): | ||
| 71 | - class MockFSDPParam: | ||
| 72 | - def __init__(self): | ||
| 73 | - self.param_dtype = torch.float32 | ||
| 74 | - self.offload_to_cpu = True | ||
| 75 | - self._sharded_local_tensor = torch.tensor([1.0, 2.0]) | ||
| 76 | - self.sharded_state = ShardedState.SHARDED | ||
| 77 | - self._sharded_param_data = torch.tensor([3.0, 4.0]) | ||
| 78 | - self._sharded_post_forward_param_data = torch.tensor([5.0, 6.0]) | ||
| 79 | - self.device = torch.device("cpu") | ||
| 80 | - self.all_gather_inputs = [torch.tensor([7.0, 8.0])] | ||
| 81 | - | ||
| 82 | - fsdp_param = MockFSDPParam() | ||
| 83 | - fsdp_params = [fsdp_param] | ||
| 84 | - pass | ||
| 85 | - | ||
| 86 | - | ||
| 87 | -if __name__ == "__main__": | ||
| 88 | - run_tests() | ||
| @@ -0,0 +1,127 @@ | |||
| 1 | +import importlib | ||
| 2 | +import os | ||
| 3 | +import unittest | ||
| 4 | + | ||
| 5 | + | ||
| 6 | +RUN_MANUAL_NATIVE_TESTS = ( | ||
| 7 | + os.getenv("TORCH_NPU_RUN_FSDP_NATIVE_COLLECTIVES_TESTS") == "1" | ||
| 8 | +) | ||
| 9 | +MANUAL_NATIVE_SKIP_REASON = ( | ||
| 10 | + "manual-only FSDP native collectives checks; " | ||
| 11 | + "set TORCH_NPU_RUN_FSDP_NATIVE_COLLECTIVES_TESTS=1 to run" | ||
| 12 | +) | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | +class TestFSDPNativeCollectives(unittest.TestCase): | ||
| 17 | + def _require_npu(self): | ||
| 18 | + import torch | ||
| 19 | + | ||
| 20 | + importlib.import_module("torch_npu") | ||
| 21 | + | ||
| 22 | + if not torch.npu.is_available(): | ||
| 23 | + self.skipTest("NPU is required for FSDP native collectives checks") | ||
| 24 | + torch.npu.set_device(0) | ||
| 25 | + return torch, torch.device("npu:0") | ||
| 26 | + | ||
| 27 | + def test_foreach_copy_same_device_npu(self): | ||
| 28 | + torch, device = self._require_npu() | ||
| 29 | + | ||
| 30 | + src_tensors = [ | ||
| 31 | + torch.tensor([1.0, 2.0], device=device), | ||
| 32 | + torch.tensor([3.0, 4.0, 5.0], device=device), | ||
| 33 | + ] | ||
| 34 | + dst_tensors = [torch.empty_like(tensor) for tensor in src_tensors] | ||
| 35 | + | ||
| 36 | + torch._foreach_copy_(dst_tensors, src_tensors) | ||
| 37 | + torch.npu.synchronize() | ||
| 38 | + | ||
| 39 | + for actual, expected in zip(dst_tensors, src_tensors): | ||
| 40 | + torch.testing.assert_close(actual.cpu(), expected.cpu()) | ||
| 41 | + | ||
| 42 | + def test_native_all_gather_copy_in_op_copies_rank_input(self): | ||
| 43 | + torch, device = self._require_npu() | ||
| 44 | + importlib.import_module("torch.distributed.fsdp._fully_shard._fsdp_collectives") | ||
| 45 | + | ||
| 46 | + all_gather_inputs = [ | ||
| 47 | + torch.tensor([1.0, 2.0], device=device), | ||
| 48 | + torch.tensor([3.0, 4.0, 5.0], device=device), | ||
| 49 | + ] | ||
| 50 | + inp_split_sizes = [tensor.numel() for tensor in all_gather_inputs] | ||
| 51 | + all_gather_input_numel = sum(inp_split_sizes) | ||
| 52 | + rank = 1 | ||
| 53 | + world_size = 2 | ||
| 54 | + all_gather_output = torch.full( | ||
| 55 | + (all_gather_input_numel * world_size,), -1.0, device=device | ||
| 56 | + ) | ||
| 57 | + | ||
| 58 | + rank_input, gathered_output = torch.ops.fsdp.all_gather_copy_in( | ||
| 59 | + all_gather_inputs, | ||
| 60 | + all_gather_output, | ||
| 61 | + inp_split_sizes, | ||
| 62 | + all_gather_input_numel, | ||
| 63 | + rank, | ||
| 64 | + ) | ||
| 65 | + torch.npu.synchronize() | ||
| 66 | + | ||
| 67 | + expected_rank_input = torch.cat(all_gather_inputs).cpu() | ||
| 68 | + expected_output = torch.cat( | ||
| 69 | + [ | ||
| 70 | + torch.full((all_gather_input_numel,), -1.0, device=device), | ||
| 71 | + torch.cat(all_gather_inputs), | ||
| 72 | + ] | ||
| 73 | + ).cpu() | ||
| 74 | + | ||
| 75 | + torch.testing.assert_close(rank_input.cpu(), expected_rank_input) | ||
| 76 | + torch.testing.assert_close(gathered_output.cpu(), expected_output) | ||
| 77 | + self.assertEqual(gathered_output.data_ptr(), all_gather_output.data_ptr()) | ||
| 78 | + | ||
| 79 | + def test_native_get_param_all_gather_inputs_uses_foreach_copy_path(self): | ||
| 80 | + torch, device = self._require_npu() | ||
| 81 | + fsdp_collectives = importlib.import_module( | ||
| 82 | + "torch.distributed.fsdp._fully_shard._fsdp_collectives" | ||
| 83 | + ) | ||
| 84 | + from torch.distributed.fsdp._fully_shard._fsdp_param import ShardedState | ||
| 85 | + | ||
| 86 | + class MockFSDPParam: | ||
| 87 | + def __init__(self, sharded_state, sharded_data, post_forward_data): | ||
| 88 | + self.param_dtype = sharded_data.dtype | ||
| 89 | + self.offload_to_cpu = False | ||
| 90 | + self._sharded_local_tensor = sharded_data | ||
| 91 | + self.sharded_state = sharded_state | ||
| 92 | + self._sharded_param_data = sharded_data | ||
| 93 | + self._sharded_post_forward_param_data = post_forward_data | ||
| 94 | + self.device = sharded_data.device | ||
| 95 | + | ||
| 96 | + sharded_data = torch.tensor([1.0, 2.0], device=device) | ||
| 97 | + post_forward_data = torch.tensor([3.0, 4.0, 5.0], device=device) | ||
| 98 | + fsdp_params = [ | ||
| 99 | + MockFSDPParam( | ||
| 100 | + ShardedState.SHARDED, | ||
| 101 | + sharded_data, | ||
| 102 | + torch.full_like(sharded_data, -1.0), | ||
| 103 | + ), | ||
| 104 | + MockFSDPParam( | ||
| 105 | + ShardedState.SHARDED_POST_FORWARD, | ||
| 106 | + torch.full_like(post_forward_data, -1.0), | ||
| 107 | + post_forward_data, | ||
| 108 | + ), | ||
| 109 | + ] | ||
| 110 | + | ||
| 111 | + all_gather_inputs = fsdp_collectives._get_param_all_gather_inputs(fsdp_params) | ||
| 112 | + torch.npu.synchronize() | ||
| 113 | + | ||
| 114 | + self.assertEqual(len(all_gather_inputs), 2) | ||
| 115 | + self.assertEqual([len(inputs) for inputs in all_gather_inputs], [1, 1]) | ||
| 116 | + torch.testing.assert_close(all_gather_inputs[0][0].cpu(), sharded_data.cpu()) | ||
| 117 | + torch.testing.assert_close( | ||
| 118 | + all_gather_inputs[1][0].cpu(), post_forward_data.cpu() | ||
| 119 | + ) | ||
| 120 | + self.assertNotEqual(all_gather_inputs[0][0].data_ptr(), sharded_data.data_ptr()) | ||
| 121 | + self.assertNotEqual( | ||
| 122 | + all_gather_inputs[1][0].data_ptr(), post_forward_data.data_ptr() | ||
| 123 | + ) | ||
| 124 | + | ||
| 125 | + | ||
| 126 | +if __name__ == "__main__": | ||
| 127 | + unittest.main() | ||
| @@ -400,15 +400,6 @@ class TestTorchNpuBootstrap(TestCase): | |||
| 400 | import torch.distributed.distributed_c10d as c10d | 400 | import torch.distributed.distributed_c10d as c10d |
| 401 | import torch.distributed.launcher.api as launcher_api | 401 | import torch.distributed.launcher.api as launcher_api |
| 402 | from torch.distributed.fsdp import sharded_grad_scaler | 402 | from torch.distributed.fsdp import sharded_grad_scaler |
| 403 | - from torch.distributed.fsdp._fully_shard import _fsdp_collectives | ||
| 404 | - from torch.distributed.fsdp._fully_shard._fsdp_param_group import ( | ||
| 405 | - FSDPParamGroup, | ||
| 406 | - ) | ||
| 407 | - from torch_npu.distributed.fsdp._add_fsdp_patch import ( | ||
| 408 | - _patched_finalize_backward, | ||
| 409 | - _patched_get_param_all_gather_inputs, | ||
| 410 | - _patched_all_gather_copy_in, | ||
| 411 | - ) | ||
| 412 | from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler | 403 | from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler |
| 413 | 404 | ||
| 414 | assert torch._C._distributed_c10d._verify_params_across_processes is ( | 405 | assert torch._C._distributed_c10d._verify_params_across_processes is ( |
| @@ -443,14 +434,6 @@ class TestTorchNpuBootstrap(TestCase): | |||
| 443 | assert callable(launcher_api._get_addr_and_port) | 434 | assert callable(launcher_api._get_addr_and_port) |
| 444 | 435 | ||
| 445 | assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler | 436 | assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler |
| 446 | - assert FSDPParamGroup.finalize_backward is _patched_finalize_backward | ||
| 447 | - assert _fsdp_collectives._get_param_all_gather_inputs is ( | ||
| 448 | - _patched_get_param_all_gather_inputs | ||
| 449 | - ) | ||
| 450 | - assert torch.ops.fsdp.all_gather_copy_in is _patched_all_gather_copy_in | ||
| 451 | - assert torch.ops.fsdp.all_gather_copy_in.default is ( | ||
| 452 | - _patched_all_gather_copy_in | ||
| 453 | - ) | ||
| 454 | """ | 437 | """ |
| 455 | ) | 438 | ) |
| 456 | 439 | ||
| @@ -3,7 +3,6 @@ from torch_npu._init.patches.patch_manager import PatchManager | |||
| 3 | 3 | ||
| 4 | 4 | ||
| 5 | def apply_torch_api_patches(): | 5 | def apply_torch_api_patches(): |
| 6 | - from torch_npu.distributed.fsdp._add_fsdp_patch import _apply_fsdp_patch | ||
| 7 | from torch_npu.multiprocessing.reductions import _add_reductions_methods | 6 | from torch_npu.multiprocessing.reductions import _add_reductions_methods |
| 8 | from torch_npu.utils._module import _apply_module_patch | 7 | from torch_npu.utils._module import _apply_module_patch |
| 9 | from torch_npu.utils._optim import add_optim_method | 8 | from torch_npu.utils._optim import add_optim_method |
| @@ -21,4 +20,3 @@ def apply_torch_api_patches(): | |||
| 21 | _add_collect_env_methods() | 20 | _add_collect_env_methods() |
| 22 | add_optim_method() | 21 | add_optim_method() |
| 23 | _add_reductions_methods() | 22 | _add_reductions_methods() |
| 24 | - _apply_fsdp_patch() | ||
| @@ -1,36 +1,33 @@ | |||
| 1 | -from collections import defaultdict | ||
| 2 | -from functools import reduce | ||
| 3 | -from typing import cast, Optional, Sequence, Union | ||
| 4 | import operator | 1 | import operator |
| 2 | +from collections import defaultdict | ||
| 3 | +from collections.abc import Sequence | ||
| 4 | +from functools import reduce | ||
| 5 | +from typing import Optional, Union | ||
| 5 | 6 | ||
| 6 | import torch | 7 | import torch |
| 7 | -from torch import distributed as dist | ||
| 8 | from torch.distributed.fsdp import fully_shard as torch_fully_shard | 8 | from torch.distributed.fsdp import fully_shard as torch_fully_shard |
| 9 | from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState | 9 | from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState |
| 10 | -from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState | ||
| 11 | from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup | 10 | from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup |
| 12 | from torch.distributed.fsdp._fully_shard._fsdp_state import FSDPState | 11 | from torch.distributed.fsdp._fully_shard._fsdp_state import FSDPState |
| 13 | 12 | ||
| 14 | -import torch_npu | ||
| 15 | - | ||
| 16 | 13 | ||
| 17 | _FSDP_ENHANCE_PATCH_APPLIED = False | 14 | _FSDP_ENHANCE_PATCH_APPLIED = False |
| 18 | 15 | ||
| 19 | 16 | ||
| 20 | class FSDPMemCache: | 17 | class FSDPMemCache: |
| 21 | def __init__(self): | 18 | def __init__(self): |
| 22 | - self.buffers = defaultdict(list) # dtype -> buffer list | 19 | + self.buffers = defaultdict(list) # dtype -> buffer list |
| 23 | - self.used = defaultdict(list) # dtype -> bool list | 20 | + self.used = defaultdict(list) # dtype -> bool list |
| 24 | 21 | ||
| 25 | def _get_storage_ptr(self, tensor: torch.Tensor) -> int: | 22 | def _get_storage_ptr(self, tensor: torch.Tensor) -> int: |
| 26 | return tensor.storage().data_ptr() | 23 | return tensor.storage().data_ptr() |
| 27 | 24 | ||
| 28 | def allocate( | 25 | def allocate( |
| 29 | self, | 26 | self, |
| 30 | - size: Sequence[int | torch.SymInt], | 27 | + size: Sequence[Union[int, torch.SymInt]], |
| 31 | *, | 28 | *, |
| 32 | dtype: torch.dtype, | 29 | dtype: torch.dtype, |
| 33 | - device: torch.device | 30 | + device: torch.device, |
| 34 | ) -> torch.Tensor: | 31 | ) -> torch.Tensor: |
| 35 | buffer_list = self.buffers[dtype] | 32 | buffer_list = self.buffers[dtype] |
| 36 | used = self.used[dtype] | 33 | used = self.used[dtype] |
| @@ -65,97 +62,6 @@ class FSDPMemCache: | |||
| 65 | _fsdp_mem_cache = FSDPMemCache() | 62 | _fsdp_mem_cache = FSDPMemCache() |
| 66 | 63 | ||
| 67 | 64 | ||
| 68 | -def _patched_finalize_backward(self): | ||
| 69 | - self._wait_for_post_backward() | ||
| 70 | - for fsdp_param in self.fsdp_params: | ||
| 71 | - if fsdp_param.grad_offload_event is not None: | ||
| 72 | - fsdp_param.grad_offload_event.synchronize() | ||
| 73 | - fsdp_param.grad_offload_event = None | ||
| 74 | - if self._all_gather_result is not None: | ||
| 75 | - # If there was a mistargeted unshard without a corresponding wait, | ||
| 76 | - # then we wait here and clear the unshard | ||
| 77 | - event = self._all_gather_result.all_gather_event | ||
| 78 | - if event is not None: | ||
| 79 | - torch.npu.current_stream().wait_event(event) | ||
| 80 | - work = self._all_gather_result.all_gather_work | ||
| 81 | - if isinstance(work, dist.distributed_c10d.Work): | ||
| 82 | - work.wait() | ||
| 83 | - self._all_gather_result = None | ||
| 84 | - self._post_forward_indices.clear() | ||
| 85 | - | ||
| 86 | - | ||
| 87 | -def _patched_get_param_all_gather_inputs( | ||
| 88 | - fsdp_params: list[FSDPParam], | ||
| 89 | -) -> list[list[torch.Tensor]]: | ||
| 90 | - # Intentionally try to run a fast-path that bypasses abstractions for the | ||
| 91 | - # common FSDP case of bf16/fp32 mixed precision in order to use foreach | ||
| 92 | - # copy for lower CPU overhead and more efficient copying in eager | ||
| 93 | - def use_foreach_copy(fsdp_param: FSDPParam) -> bool: | ||
| 94 | - return ( | ||
| 95 | - fsdp_param.param_dtype is not None | ||
| 96 | - and not fsdp_param.offload_to_cpu | ||
| 97 | - and not hasattr(fsdp_param._sharded_local_tensor, "fsdp_pre_all_gather") | ||
| 98 | - ) | ||
| 99 | - | ||
| 100 | - param_all_gather_inputs: list[list[torch.Tensor]] = [[] for _ in fsdp_params] | ||
| 101 | - foreach_copy_indices: list[int] = [] | ||
| 102 | - foreach_copy_inputs: list[torch.Tensor] = [] | ||
| 103 | - foreach_copy_input_numels: list[int] = [] | ||
| 104 | - | ||
| 105 | - # 1st pass: for foreach-copy parameters, get inputs and metadata for the | ||
| 106 | - # foreach copy, and for the others, actually get their all-gather inputs | ||
| 107 | - for i, fsdp_param in enumerate(fsdp_params): | ||
| 108 | - if use_foreach_copy(fsdp_param): | ||
| 109 | - foreach_copy_indices.append(i) | ||
| 110 | - all_gather_input = ( | ||
| 111 | - fsdp_param._sharded_param_data | ||
| 112 | - if fsdp_param.sharded_state == ShardedState.SHARDED | ||
| 113 | - else cast(torch.Tensor, fsdp_param._sharded_post_forward_param_data) | ||
| 114 | - ) | ||
| 115 | - foreach_copy_inputs.append(all_gather_input) | ||
| 116 | - foreach_copy_input_numels.append(all_gather_input.numel()) | ||
| 117 | - else: | ||
| 118 | - param_all_gather_inputs[i] = fsdp_param.all_gather_inputs | ||
| 119 | - | ||
| 120 | - # 2nd pass: use foreach copy to compute the remaining all-gather inputs | ||
| 121 | - if foreach_copy_inputs: | ||
| 122 | - fsdp_param_0 = fsdp_params[foreach_copy_indices[0]] | ||
| 123 | - param_dtype, device = fsdp_param_0.param_dtype, fsdp_param_0.device | ||
| 124 | - flat_foreach_copy_input = torch.empty( | ||
| 125 | - (sum(foreach_copy_input_numels),), device=device, dtype=param_dtype | ||
| 126 | - ) | ||
| 127 | - splits = torch.split(flat_foreach_copy_input, foreach_copy_input_numels) | ||
| 128 | - # patch in npu: set non_blocking=True | ||
| 129 | - if splits[0].device == foreach_copy_inputs[0].device: | ||
| 130 | - torch._foreach_copy_(splits, foreach_copy_inputs, non_blocking=True) | ||
| 131 | - else: | ||
| 132 | - torch._foreach_copy_(splits, foreach_copy_inputs) | ||
| 133 | - for i, split in zip(foreach_copy_indices, splits): | ||
| 134 | - param_all_gather_inputs[i] = [split] | ||
| 135 | - | ||
| 136 | - return param_all_gather_inputs | ||
| 137 | - | ||
| 138 | - | ||
| 139 | -def _patched_all_gather_copy_in( | ||
| 140 | - all_gather_inputs: list[torch.Tensor], | ||
| 141 | - all_gather_output: torch.Tensor, | ||
| 142 | - inp_split_sizes: list[int], | ||
| 143 | - all_gather_input_numel: int, | ||
| 144 | - rank: int, | ||
| 145 | -) -> tuple[torch.Tensor, torch.Tensor]: | ||
| 146 | - all_gather_input = all_gather_output.narrow( | ||
| 147 | - 0, all_gather_input_numel * rank, all_gather_input_numel | ||
| 148 | - ) | ||
| 149 | - foreach_copy_dsts = torch.split(all_gather_input, inp_split_sizes) | ||
| 150 | - with torch.no_grad(): | ||
| 151 | - # patch in npu: set non_blocking=True | ||
| 152 | - if foreach_copy_dsts[0].device == all_gather_inputs[0].device: | ||
| 153 | - torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs, non_blocking=True) | ||
| 154 | - else: | ||
| 155 | - torch._foreach_copy_(foreach_copy_dsts, all_gather_inputs) | ||
| 156 | - return all_gather_input, all_gather_output | ||
| 157 | - | ||
| 158 | - | ||
| 159 | def _patched_fsdp_param_group_init(original_func): | 65 | def _patched_fsdp_param_group_init(original_func): |
| 160 | def wrapper(self, *args, **kwargs): | 66 | def wrapper(self, *args, **kwargs): |
| 161 | original_func(self, *args, **kwargs) | 67 | original_func(self, *args, **kwargs) |
| @@ -164,6 +70,7 @@ def _patched_fsdp_param_group_init(original_func): | |||
| 164 | use_mem_cache = getattr(self.modules[0], "_use_mem_cache", False) | 70 | use_mem_cache = getattr(self.modules[0], "_use_mem_cache", False) |
| 165 | self._all_gather_comm._use_mem_cache = use_mem_cache | 71 | self._all_gather_comm._use_mem_cache = use_mem_cache |
| 166 | self._reduce_scatter_comm._use_mem_cache = use_mem_cache | 72 | self._reduce_scatter_comm._use_mem_cache = use_mem_cache |
| 73 | + | ||
| 167 | return wrapper | 74 | return wrapper |
| 168 | 75 | ||
| 169 | 76 | ||
| @@ -183,13 +90,19 @@ def _patched_wait_all_gather_streams_on_event(original_func): | |||
| 183 | def wrapper(self, event: Optional[torch.Event]): | 90 | def wrapper(self, event: Optional[torch.Event]): |
| 184 | original_func(self, event) | 91 | original_func(self, event) |
| 185 | # if previous layer deferred free for overlap, free its output in current comm_ctx.all_gather_state | 92 | # if previous layer deferred free for overlap, free its output in current comm_ctx.all_gather_state |
| 186 | - if self._training_state == TrainingState.FORWARD and self.comm_ctx.all_gather_state: | 93 | + if ( |
| 187 | - prev_all_gather_output = self.comm_ctx.all_gather_state.all_gather_result.all_gather_output | 94 | + self._training_state == TrainingState.FORWARD |
| 95 | + and self.comm_ctx.all_gather_state | ||
| 96 | + ): | ||
| 97 | + prev_all_gather_output = ( | ||
| 98 | + self.comm_ctx.all_gather_state.all_gather_result.all_gather_output | ||
| 99 | + ) | ||
| 188 | _fsdp_mem_cache.free(prev_all_gather_output) | 100 | _fsdp_mem_cache.free(prev_all_gather_output) |
| 189 | # if current layer no need to defer free, free output after all_gather_copy_out event | 101 | # if current layer no need to defer free, free output after all_gather_copy_out event |
| 190 | elif self._all_gather_result: | 102 | elif self._all_gather_result: |
| 191 | all_gather_output = self._all_gather_result.all_gather_output | 103 | all_gather_output = self._all_gather_result.all_gather_output |
| 192 | _fsdp_mem_cache.free(all_gather_output) | 104 | _fsdp_mem_cache.free(all_gather_output) |
| 105 | + | ||
| 193 | return wrapper | 106 | return wrapper |
| 194 | 107 | ||
| 195 | 108 | ||
| @@ -201,7 +114,9 @@ def _patched_reduce_scatter_allocate( | |||
| 201 | device: torch.device, | 114 | device: torch.device, |
| 202 | ) -> torch.Tensor: | 115 | ) -> torch.Tensor: |
| 203 | # foreach_reduce allocates memory for both input and output, we only cache the input, i.e. on the first call | 116 | # foreach_reduce allocates memory for both input and output, we only cache the input, i.e. on the first call |
| 204 | - if getattr(self, "_use_mem_cache", False) and not getattr(self, "_mem_cache_flag", False): | 117 | + if getattr(self, "_use_mem_cache", False) and not getattr( |
| 118 | + self, "_mem_cache_flag", False | ||
| 119 | + ): | ||
| 205 | self._mem_cache_flag = True | 120 | self._mem_cache_flag = True |
| 206 | return _fsdp_mem_cache.allocate(size, dtype=dtype, device=device) | 121 | return _fsdp_mem_cache.allocate(size, dtype=dtype, device=device) |
| 207 | return torch.empty(*size, dtype=dtype, device=device) | 122 | return torch.empty(*size, dtype=dtype, device=device) |
| @@ -215,6 +130,7 @@ def _patched_foreach_reduce(original_foreach_reduce): | |||
| 215 | reduce_scatter_comm = kwargs.get("reduce_scatter_comm", args[4]) | 130 | reduce_scatter_comm = kwargs.get("reduce_scatter_comm", args[4]) |
| 216 | reduce_scatter_comm._mem_cache_flag = False | 131 | reduce_scatter_comm._mem_cache_flag = False |
| 217 | return out | 132 | return out |
| 133 | + | ||
| 218 | return wrapper | 134 | return wrapper |
| 219 | 135 | ||
| 220 | 136 | ||
| @@ -227,6 +143,7 @@ def _patched_post_forward(original_post_forward): | |||
| 227 | out = original_post_forward(self, module, args, out) | 143 | out = original_post_forward(self, module, args, out) |
| 228 | self._skip_post_forward = True | 144 | self._skip_post_forward = True |
| 229 | return out | 145 | return out |
| 146 | + | ||
| 230 | return wrapper | 147 | return wrapper |
| 231 | 148 | ||
| 232 | 149 | ||
| @@ -235,6 +152,7 @@ def _patched_post_backward(original_post_backward): | |||
| 235 | def wrapper(self): | 152 | def wrapper(self): |
| 236 | self._skip_post_forward = False | 153 | self._skip_post_forward = False |
| 237 | original_post_backward(self) | 154 | original_post_backward(self) |
| 155 | + | ||
| 238 | return wrapper | 156 | return wrapper |
| 239 | 157 | ||
| 240 | 158 | ||
| @@ -243,7 +161,7 @@ def move_attr(src_obj, src_attr, dst_obj, dst_attr): | |||
| 243 | src_value = getattr(src_obj, src_attr, None) | 161 | src_value = getattr(src_obj, src_attr, None) |
| 244 | setattr(dst_obj, dst_attr, src_value) | 162 | setattr(dst_obj, dst_attr, src_value) |
| 245 | if type(src_value) in (list, tuple, set, dict): | 163 | if type(src_value) in (list, tuple, set, dict): |
| 246 | - setattr(src_obj, src_attr, type(src_value)()) # [] / () / set() / {} | 164 | + setattr(src_obj, src_attr, type(src_value)()) # [] / () / set() / {} |
| 247 | else: | 165 | else: |
| 248 | setattr(src_obj, src_attr, None) | 166 | setattr(src_obj, src_attr, None) |
| 249 | 167 | ||
| @@ -262,24 +180,18 @@ def _patched_fsdp_state_post_forward(original_post_forward): | |||
| 262 | skip_post_forward = getattr(self._fsdp_param_group, "_skip_post_forward", False) | 180 | skip_post_forward = getattr(self._fsdp_param_group, "_skip_post_forward", False) |
| 263 | if skip_post_forward or self._training_state == TrainingState.PRE_BACKWARD: | 181 | if skip_post_forward or self._training_state == TrainingState.PRE_BACKWARD: |
| 264 | if self._backup_forward_fetch is not None: | 182 | if self._backup_forward_fetch is not None: |
| 265 | - move_attr(self, "_backup_forward_fetch", self, "_states_to_forward_prefetch") | 183 | + move_attr( |
| 184 | + self, "_backup_forward_fetch", self, "_states_to_forward_prefetch" | ||
| 185 | + ) | ||
| 266 | return original_post_forward(self, module, args, out) | 186 | return original_post_forward(self, module, args, out) |
| 267 | 187 | ||
| 268 | # backup forward prefetch state before original post_forward | 188 | # backup forward prefetch state before original post_forward |
| 269 | move_attr(self, "_states_to_forward_prefetch", self, "_backup_forward_fetch") | 189 | move_attr(self, "_states_to_forward_prefetch", self, "_backup_forward_fetch") |
| 270 | return original_post_forward(self, module, args, out) | 190 | return original_post_forward(self, module, args, out) |
| 191 | + | ||
| 271 | return wrapper | 192 | return wrapper |
| 272 | 193 | ||
| 273 | 194 | ||
| 274 | -def _apply_fsdp_patch(): | ||
| 275 | - # essential patch to run on NPU | ||
| 276 | - FSDPParamGroup.finalize_backward = _patched_finalize_backward | ||
| 277 | - torch.distributed.fsdp._fully_shard._fsdp_collectives._get_param_all_gather_inputs \ | ||
| 278 | - = _patched_get_param_all_gather_inputs | ||
| 279 | - torch.ops.fsdp.all_gather_copy_in = _patched_all_gather_copy_in | ||
| 280 | - torch.ops.fsdp.all_gather_copy_in.default = _patched_all_gather_copy_in | ||
| 281 | - | ||
| 282 | - | ||
| 283 | def _apply_fsdp_enhance_patch(): | 195 | def _apply_fsdp_enhance_patch(): |
| 284 | global _FSDP_ENHANCE_PATCH_APPLIED | 196 | global _FSDP_ENHANCE_PATCH_APPLIED |
| 285 | if _FSDP_ENHANCE_PATCH_APPLIED: | 197 | if _FSDP_ENHANCE_PATCH_APPLIED: |
| @@ -287,21 +199,29 @@ def _apply_fsdp_enhance_patch(): | |||
| 287 | 199 | ||
| 288 | # support using memory cache for FSDP comm ops | 200 | # support using memory cache for FSDP comm ops |
| 289 | FSDPParamGroup.__init__ = _patched_fsdp_param_group_init(FSDPParamGroup.__init__) | 201 | FSDPParamGroup.__init__ = _patched_fsdp_param_group_init(FSDPParamGroup.__init__) |
| 290 | - FSDPParamGroup._wait_all_gather_streams_on_event \ | 202 | + FSDPParamGroup._wait_all_gather_streams_on_event = ( |
| 291 | - = _patched_wait_all_gather_streams_on_event(FSDPParamGroup._wait_all_gather_streams_on_event) | 203 | + _patched_wait_all_gather_streams_on_event( |
| 292 | - torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultAllGather.allocate = _patched_all_gather_allocate | 204 | + FSDPParamGroup._wait_all_gather_streams_on_event |
| 293 | - torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultReduceScatter.allocate \ | 205 | + ) |
| 294 | - = _patched_reduce_scatter_allocate | 206 | + ) |
| 295 | - origin_foreach_reduce = torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce | 207 | + torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultAllGather.allocate = ( |
| 296 | - torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce \ | 208 | + _patched_all_gather_allocate |
| 297 | - = _patched_foreach_reduce(origin_foreach_reduce) | 209 | + ) |
| 210 | + torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultReduceScatter.allocate = _patched_reduce_scatter_allocate | ||
| 211 | + origin_foreach_reduce = ( | ||
| 212 | + torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce | ||
| 213 | + ) | ||
| 214 | + torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce = ( | ||
| 215 | + _patched_foreach_reduce(origin_foreach_reduce) | ||
| 216 | + ) | ||
| 298 | # _fsdp_param_group imported these functions before patching | 217 | # _fsdp_param_group imported these functions before patching |
| 299 | - torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultAllGather.allocate = _patched_all_gather_allocate | 218 | + torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultAllGather.allocate = ( |
| 300 | - torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultReduceScatter.allocate \ | 219 | + _patched_all_gather_allocate |
| 301 | - = _patched_reduce_scatter_allocate | 220 | + ) |
| 302 | - torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce \ | 221 | + torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultReduceScatter.allocate = _patched_reduce_scatter_allocate |
| 303 | - = _patched_foreach_reduce(origin_foreach_reduce) | 222 | + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce = ( |
| 304 | - | 223 | + _patched_foreach_reduce(origin_foreach_reduce) |
| 224 | + ) | ||
| 305 | 225 | ||
| 306 | # optimize communication, e.g. removing redundant all-gather when recomputing in backward | 226 | # optimize communication, e.g. removing redundant all-gather when recomputing in backward |
| 307 | FSDPState._post_forward = _patched_fsdp_state_post_forward(FSDPState._post_forward) | 227 | FSDPState._post_forward = _patched_fsdp_state_post_forward(FSDPState._post_forward) |