已合并
refactor fsdp patch and add torch_npu.distributed.fsdp.fully_shard #32148
jizewei创建于 3月21日
refactor fsdp patch and add torch_npu.distributed.fsdp.fully_shard #32148
已合并
共 3 个文件变更+46-9
| @@ -563,6 +563,9 @@ | |||
| 563 | "torch_npu.distributed.distributed_c10d.reinit_process_group": { | 563 | "torch_npu.distributed.distributed_c10d.reinit_process_group": { |
| 564 | "signature": "(group=None, rebuild_link=True)" | 564 | "signature": "(group=None, rebuild_link=True)" |
| 565 | }, | 565 | }, |
| 566 | + "torch_npu.distributed.fsdp.fully_shard": { | ||
| 567 | + "signature": "(*args, **kwargs)" | ||
| 568 | + }, | ||
| 566 | "torch_npu.distributed.reinit_process_group": { | 569 | "torch_npu.distributed.reinit_process_group": { |
| 567 | "signature": "(group=None, rebuild_link=True)" | 570 | "signature": "(group=None, rebuild_link=True)" |
| 568 | }, | 571 | }, |
| @@ -0,0 +1,5 @@ | |||
| 1 | +from ._add_fsdp_patch import fully_shard | ||
| 2 | + | ||
| 3 | +fully_shard.__module__ = __name__ | ||
| 4 | + | ||
| 5 | +__all__ = ["fully_shard"] | ||
| @@ -5,6 +5,7 @@ import operator | |||
| 5 | 5 | ||
| 6 | import torch | 6 | import torch |
| 7 | from torch import distributed as dist | 7 | from torch import distributed as dist |
| 8 | +from torch.distributed.fsdp import fully_shard as torch_fully_shard | ||
| 8 | from torch.distributed.fsdp._fully_shard._fsdp_common import compiled_autograd_enabled, TrainingState | 9 | from torch.distributed.fsdp._fully_shard._fsdp_common import compiled_autograd_enabled, TrainingState |
| 9 | from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState | 10 | from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam, ShardedState |
| 10 | from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup | 11 | from torch.distributed.fsdp._fully_shard._fsdp_param_group import FSDPParamGroup |
| @@ -13,6 +14,9 @@ from torch.distributed.fsdp._fully_shard._fsdp_state import FSDPState | |||
| 13 | import torch_npu | 14 | import torch_npu |
| 14 | 15 | ||
| 15 | 16 | ||
| 17 | +_FSDP_ENHANCE_PATCH_APPLIED = False | ||
| 18 | + | ||
| 19 | + | ||
| 16 | class FSDPMemCache: | 20 | class FSDPMemCache: |
| 17 | def __init__(self): | 21 | def __init__(self): |
| 18 | self.buffers = defaultdict(list) # dtype -> buffer list | 22 | self.buffers = defaultdict(list) # dtype -> buffer list |
| @@ -271,24 +275,49 @@ def _patched_fsdp_state_post_forward(original_post_forward): | |||
| 271 | 275 | ||
| 272 | 276 | ||
| 273 | def _apply_fsdp_patch(): | 277 | def _apply_fsdp_patch(): |
| 274 | - FSDPState._post_forward = _patched_fsdp_state_post_forward(FSDPState._post_forward) | 278 | + # essential patch to run on NPU |
| 275 | - FSDPParamGroup.__init__ = _patched_fsdp_param_group_init(FSDPParamGroup.__init__) | ||
| 276 | - FSDPParamGroup._wait_all_gather_streams_on_event \ | ||
| 277 | - = _patched_wait_all_gather_streams_on_event(FSDPParamGroup._wait_all_gather_streams_on_event) | ||
| 278 | - FSDPParamGroup.post_forward = _patched_post_forward(FSDPParamGroup.post_forward) | ||
| 279 | - FSDPParamGroup.post_backward = _patched_post_backward(FSDPParamGroup.post_backward) | ||
| 280 | FSDPParamGroup.finalize_backward = _patched_finalize_backward | 279 | FSDPParamGroup.finalize_backward = _patched_finalize_backward |
| 281 | - torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultAllGather.allocate = _patched_all_gather_allocate | ||
| 282 | - torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultReduceScatter.allocate = _patched_reduce_scatter_allocate | ||
| 283 | torch.distributed.fsdp._fully_shard._fsdp_collectives._get_param_all_gather_inputs \ | 280 | torch.distributed.fsdp._fully_shard._fsdp_collectives._get_param_all_gather_inputs \ |
| 284 | = _patched_get_param_all_gather_inputs | 281 | = _patched_get_param_all_gather_inputs |
| 285 | torch.ops.fsdp.all_gather_copy_in = _patched_all_gather_copy_in | 282 | torch.ops.fsdp.all_gather_copy_in = _patched_all_gather_copy_in |
| 286 | torch.ops.fsdp.all_gather_copy_in.default = _patched_all_gather_copy_in | 283 | torch.ops.fsdp.all_gather_copy_in.default = _patched_all_gather_copy_in |
| 284 | + | ||
| 285 | + | ||
| 286 | +def _apply_fsdp_enhance_patch(): | ||
| 287 | + global _FSDP_ENHANCE_PATCH_APPLIED | ||
| 288 | + if _FSDP_ENHANCE_PATCH_APPLIED: | ||
| 289 | + return | ||
| 290 | + | ||
| 291 | + # support using memory cache for FSDP comm ops | ||
| 292 | + FSDPParamGroup.__init__ = _patched_fsdp_param_group_init(FSDPParamGroup.__init__) | ||
| 293 | + FSDPParamGroup._wait_all_gather_streams_on_event \ | ||
| 294 | + = _patched_wait_all_gather_streams_on_event(FSDPParamGroup._wait_all_gather_streams_on_event) | ||
| 295 | + torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultAllGather.allocate = _patched_all_gather_allocate | ||
| 296 | + torch.distributed.fsdp._fully_shard._fsdp_collectives.DefaultReduceScatter.allocate \ | ||
| 297 | + = _patched_reduce_scatter_allocate | ||
| 287 | origin_foreach_reduce = torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce | 298 | origin_foreach_reduce = torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce |
| 288 | torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce \ | 299 | torch.distributed.fsdp._fully_shard._fsdp_collectives.foreach_reduce \ |
| 289 | = _patched_foreach_reduce(origin_foreach_reduce) | 300 | = _patched_foreach_reduce(origin_foreach_reduce) |
| 290 | # _fsdp_param_group imported these functions before patching | 301 | # _fsdp_param_group imported these functions before patching |
| 291 | torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultAllGather.allocate = _patched_all_gather_allocate | 302 | torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultAllGather.allocate = _patched_all_gather_allocate |
| 292 | - torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultReduceScatter.allocate = _patched_reduce_scatter_allocate | 303 | + torch.distributed.fsdp._fully_shard._fsdp_param_group.DefaultReduceScatter.allocate \ |
| 304 | + = _patched_reduce_scatter_allocate | ||
| 293 | torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce \ | 305 | torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce \ |
| 294 | = _patched_foreach_reduce(origin_foreach_reduce) | 306 | = _patched_foreach_reduce(origin_foreach_reduce) |
| 307 | + | ||
| 308 | + | ||
| 309 | + # optimize communication, e.g. removing redundant all-gather when recomputing in backward | ||
| 310 | + FSDPState._post_forward = _patched_fsdp_state_post_forward(FSDPState._post_forward) | ||
| 311 | + FSDPParamGroup.post_forward = _patched_post_forward(FSDPParamGroup.post_forward) | ||
| 312 | + FSDPParamGroup.post_backward = _patched_post_backward(FSDPParamGroup.post_backward) | ||
| 313 | + | ||
| 314 | + _FSDP_ENHANCE_PATCH_APPLIED = True | ||
| 315 | + | ||
| 316 | + | ||
| 317 | +def fully_shard(*args, **kwargs): | ||
| 318 | + _apply_fsdp_enhance_patch() | ||
| 319 | + return torch_fully_shard(*args, **kwargs) | ||
| 320 | + | ||
| 321 | + | ||
| 322 | +fully_shard.state = torch_fully_shard.state | ||
| 323 | +fully_shard.__doc__ = torch_fully_shard.__doc__ | ||