已合并
Add distributed methods patch #16598
AtomGit-Bot创建于 2024年12月3日
Add distributed methods patch #16598
已合并
从refs/pull/16598/head合入到master
共 11 个文件变更+35-97
| @@ -1665,6 +1665,11 @@ | |||
| 1665 | "cast", | 1665 | "cast", |
| 1666 | "dataclass" | 1666 | "dataclass" |
| 1667 | ], | 1667 | ], |
| 1668 | + "torch.distributed.distributed_c10d": [ | ||
| 1669 | + "gather", | ||
| 1670 | + "gather_object", | ||
| 1671 | + "batch_isend_irecv" | ||
| 1672 | + ], | ||
| 1668 | "torch.distributed.elastic.rendezvous.c10d_rendezvous_backend": [ | 1673 | "torch.distributed.elastic.rendezvous.c10d_rendezvous_backend": [ |
| 1669 | "Any", | 1674 | "Any", |
| 1670 | "FileStore", | 1675 | "FileStore", |
| @@ -18,7 +18,7 @@ from typing import ( | |||
| 18 | import torch_npu | 18 | import torch_npu |
| 19 | from torch_npu.testing.testcase import TestCase, run_tests | 19 | from torch_npu.testing.testcase import TestCase, run_tests |
| 20 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | 20 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU |
| 21 | -from torch_npu.npu.amp.sharded_grad_scaler import ShardedGradScaler as NpuShardedGradScaler | 21 | +from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler as NpuShardedGradScaler |
| 22 | 22 | ||
| 23 | import torch | 23 | import torch |
| 24 | import torch.distributed as dist | 24 | import torch.distributed as dist |
| @@ -6,7 +6,7 @@ import torch | |||
| 6 | import torch.distributed as dist | 6 | import torch.distributed as dist |
| 7 | 7 | ||
| 8 | import torch_npu | 8 | import torch_npu |
| 9 | -from torch_npu.npu.amp.sharded_grad_scaler import ShardedGradScaler as NpuShardedGradScaler | 9 | +from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler as NpuShardedGradScaler |
| 10 | from torch_npu.testing.testcase import TestCase, run_tests | 10 | from torch_npu.testing.testcase import TestCase, run_tests |
| 11 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU | 11 | from torch_npu.testing.common_distributed import skipIfUnsupportMultiNPU |
| 12 | 12 | ||
| @@ -809,30 +809,12 @@ | |||
| 809 | "torch_npu.contrib.module.roi_align.ROIAlign.forward": { | 809 | "torch_npu.contrib.module.roi_align.ROIAlign.forward": { |
| 810 | "signature": "(self, input_tensor, rois)" | 810 | "signature": "(self, input_tensor, rois)" |
| 811 | }, | 811 | }, |
| 812 | - "torch_npu.distributed.batch_isend_irecv": { | ||
| 813 | - "signature": "(p2p_op_list)" | ||
| 814 | - }, | ||
| 815 | - "torch_npu.distributed.gather": { | ||
| 816 | - "signature": "(tensor, gather_list=None, dst=0, group=None, async_op=False)" | ||
| 817 | - }, | ||
| 818 | - "torch_npu.distributed.gather_object": { | ||
| 819 | - "signature": "(obj, object_gather_list=None, dst=0, group=None)" | ||
| 820 | - }, | ||
| 821 | "torch_npu.distributed.is_available": { | 812 | "torch_npu.distributed.is_available": { |
| 822 | "signature": "()" | 813 | "signature": "()" |
| 823 | }, | 814 | }, |
| 824 | "torch_npu.distributed.is_hccl_available": { | 815 | "torch_npu.distributed.is_hccl_available": { |
| 825 | "signature": "()" | 816 | "signature": "()" |
| 826 | }, | 817 | }, |
| 827 | - "torch_npu.distributed.distributed_c10d.batch_isend_irecv": { | ||
| 828 | - "signature": "(p2p_op_list)" | ||
| 829 | - }, | ||
| 830 | - "torch_npu.distributed.distributed_c10d.gather": { | ||
| 831 | - "signature": "(tensor, gather_list=None, dst=0, group=None, async_op=False)" | ||
| 832 | - }, | ||
| 833 | - "torch_npu.distributed.distributed_c10d.gather_object": { | ||
| 834 | - "signature": "(obj, object_gather_list=None, dst=0, group=None)" | ||
| 835 | - }, | ||
| 836 | "torch_npu.distributed.distributed_c10d.is_hccl_available": { | 818 | "torch_npu.distributed.distributed_c10d.is_hccl_available": { |
| 837 | "signature": "()" | 819 | "signature": "()" |
| 838 | }, | 820 | }, |
| @@ -842,12 +824,6 @@ | |||
| 842 | "torch_npu.distributed.reinit_process_group": { | 824 | "torch_npu.distributed.reinit_process_group": { |
| 843 | "signature": "(group=None, rebuild_link=True)" | 825 | "signature": "(group=None, rebuild_link=True)" |
| 844 | }, | 826 | }, |
| 845 | - "torch_npu.distributed.rpc.backend_registry.construct_rpc_backend_options": { | ||
| 846 | - "signature": "(backend, rpc_timeout=60.0, init_method='env://', **kwargs)" | ||
| 847 | - }, | ||
| 848 | - "torch_npu.distributed.rpc.backend_registry.init_backend": { | ||
| 849 | - "signature": "(backend, *args, **kwargs)" | ||
| 850 | - }, | ||
| 851 | "torch_npu.distributed.rpc.options.NPUTensorPipeRpcBackendOptions": { | 827 | "torch_npu.distributed.rpc.options.NPUTensorPipeRpcBackendOptions": { |
| 852 | "signature": "(*, num_worker_threads: int = 16, rpc_timeout: float = 60.0, init_method: str = 'env://', device_maps: Optional[Dict[str, Dict[Union[int, str, torch.device], Union[int, str, torch.device]]]] = None, devices: Optional[List[Union[int, str, torch.device]]] = None, _transports: Optional[List] = None, _channels: Optional[List] = None)" | 828 | "signature": "(*, num_worker_threads: int = 16, rpc_timeout: float = 60.0, init_method: str = 'env://', device_maps: Optional[Dict[str, Dict[Union[int, str, torch.device], Union[int, str, torch.device]]]] = None, devices: Optional[List[Union[int, str, torch.device]]] = None, _transports: Optional[List] = None, _channels: Optional[List] = None)" |
| 853 | }, | 829 | }, |
| @@ -1343,33 +1319,6 @@ | |||
| 1343 | "torch_npu.npu.amp.GradScaler._npu_update_scale": { | 1319 | "torch_npu.npu.amp.GradScaler._npu_update_scale": { |
| 1344 | "signature": "(self)" | 1320 | "signature": "(self)" |
| 1345 | }, | 1321 | }, |
| 1346 | - "torch_npu.npu.amp.ShardedGradScaler": { | ||
| 1347 | - "signature": "(init_scale: float = 65536.0, backoff_factor: float = 0.5, growth_factor: float = 2.0, growth_interval: int = 2000, enabled: bool = True, process_group: Optional[torch.distributed.distributed_c10d.ProcessGroup] = None)" | ||
| 1348 | - }, | ||
| 1349 | - "torch_npu.npu.amp.ShardedGradScaler.scale": { | ||
| 1350 | - "signature": "(self, outputs: Union[torch.Tensor, List[torch.Tensor]]) -> Union[torch.Tensor, List[torch.Tensor]]" | ||
| 1351 | - }, | ||
| 1352 | - "torch_npu.npu.amp.ShardedGradScaler._foreach_non_finite_check_and_unscale_cpu_": { | ||
| 1353 | - "signature": "(self, grads: List, found_inf: torch.Tensor, inv_scale: torch.Tensor) -> None" | ||
| 1354 | - }, | ||
| 1355 | - "torch_npu.npu.amp.ShardedGradScaler._unscale_grads_": { | ||
| 1356 | - "signature": "(self, optimizer: torch.optim.sgd.SGD, inv_scale: torch.Tensor, found_inf: torch.Tensor, allow_fp16: bool = True) -> Dict[torch.device, torch.Tensor]" | ||
| 1357 | - }, | ||
| 1358 | - "torch_npu.npu.amp.ShardedGradScaler.unscale_": { | ||
| 1359 | - "signature": "(self, optimizer: torch.optim.sgd.SGD) -> None" | ||
| 1360 | - }, | ||
| 1361 | - "torch_npu.npu.amp.ShardedGradScaler.step": { | ||
| 1362 | - "signature": "(self, optimizer: torch.optim.sgd.SGD, *args, **kwargs) -> Optional[float]" | ||
| 1363 | - }, | ||
| 1364 | - "torch_npu.npu.amp.ShardedGradScaler._amp_update_scale_cpu_": { | ||
| 1365 | - "signature": "(self, found_inf) -> None" | ||
| 1366 | - }, | ||
| 1367 | - "torch_npu.npu.amp.ShardedGradScaler._amp_update_scale_npu_": { | ||
| 1368 | - "signature": "(self, found_inf) -> None" | ||
| 1369 | - }, | ||
| 1370 | - "torch_npu.npu.amp.ShardedGradScaler.update": { | ||
| 1371 | - "signature": "(self, new_scale=None) -> None" | ||
| 1372 | - }, | ||
| 1373 | "torch_npu.npu.amp.autocast": { | 1322 | "torch_npu.npu.amp.autocast": { |
| 1374 | "signature": "(enabled: bool = True, dtype: torch.dtype = torch.float16, cache_enabled: bool = True)" | 1323 | "signature": "(enabled: bool = True, dtype: torch.dtype = torch.float16, cache_enabled: bool = True)" |
| 1375 | }, | 1324 | }, |
| @@ -1436,33 +1385,6 @@ | |||
| 1436 | "torch_npu.npu.amp.grad_scaler.GradScaler._npu_update_scale": { | 1385 | "torch_npu.npu.amp.grad_scaler.GradScaler._npu_update_scale": { |
| 1437 | "signature": "(self)" | 1386 | "signature": "(self)" |
| 1438 | }, | 1387 | }, |
| 1439 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler": { | ||
| 1440 | - "signature": "(init_scale: float = 65536.0, backoff_factor: float = 0.5, growth_factor: float = 2.0, growth_interval: int = 2000, enabled: bool = True, process_group: Optional[torch.distributed.distributed_c10d.ProcessGroup] = None)" | ||
| 1441 | - }, | ||
| 1442 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler.scale": { | ||
| 1443 | - "signature": "(self, outputs: Union[torch.Tensor, List[torch.Tensor]]) -> Union[torch.Tensor, List[torch.Tensor]]" | ||
| 1444 | - }, | ||
| 1445 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler._foreach_non_finite_check_and_unscale_cpu_": { | ||
| 1446 | - "signature": "(self, grads: List, found_inf: torch.Tensor, inv_scale: torch.Tensor) -> None" | ||
| 1447 | - }, | ||
| 1448 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler._unscale_grads_": { | ||
| 1449 | - "signature": "(self, optimizer: torch.optim.sgd.SGD, inv_scale: torch.Tensor, found_inf: torch.Tensor, allow_fp16: bool = True) -> Dict[torch.device, torch.Tensor]" | ||
| 1450 | - }, | ||
| 1451 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler.unscale_": { | ||
| 1452 | - "signature": "(self, optimizer: torch.optim.sgd.SGD) -> None" | ||
| 1453 | - }, | ||
| 1454 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler.step": { | ||
| 1455 | - "signature": "(self, optimizer: torch.optim.sgd.SGD, *args, **kwargs) -> Optional[float]" | ||
| 1456 | - }, | ||
| 1457 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler._amp_update_scale_cpu_": { | ||
| 1458 | - "signature": "(self, found_inf) -> None" | ||
| 1459 | - }, | ||
| 1460 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler._amp_update_scale_npu_": { | ||
| 1461 | - "signature": "(self, found_inf) -> None" | ||
| 1462 | - }, | ||
| 1463 | - "torch_npu.npu.amp.sharded_grad_scaler.ShardedGradScaler.update": { | ||
| 1464 | - "signature": "(self, new_scale=None) -> None" | ||
| 1465 | - }, | ||
| 1466 | "torch_npu.npu.autocast_utils.get_amp_supported_dtype": { | 1388 | "torch_npu.npu.autocast_utils.get_amp_supported_dtype": { |
| 1467 | "signature": "()" | 1389 | "signature": "()" |
| 1468 | }, | 1390 | }, |
| @@ -41,6 +41,7 @@ import torch_npu.optim | |||
| 41 | import torch_npu.dynamo | 41 | import torch_npu.dynamo |
| 42 | import torch_npu._C | 42 | import torch_npu._C |
| 43 | from torch_npu import profiler | 43 | from torch_npu import profiler |
| 44 | +from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler | ||
| 44 | from torch_npu.contrib.function import npu_functional | 45 | from torch_npu.contrib.function import npu_functional |
| 45 | from torch_npu.contrib.module import npu_modules | 46 | from torch_npu.contrib.module import npu_modules |
| 46 | from torch_npu.utils import _apply_module_patch, _add_tensor_methods, _add_collect_env_methods, \ | 47 | from torch_npu.utils import _apply_module_patch, _add_tensor_methods, _add_collect_env_methods, \ |
| @@ -128,7 +129,7 @@ def _apply_distributed_patches(): | |||
| 128 | 129 | ||
| 129 | 130 | ||
| 130 | def _apply_sharded_grad_scaler_patch(): | 131 | def _apply_sharded_grad_scaler_patch(): |
| 131 | - torch.distributed.fsdp.sharded_grad_scaler.ShardedGradScaler = torch_npu.npu.amp.ShardedGradScaler | 132 | + torch.distributed.fsdp.sharded_grad_scaler.ShardedGradScaler = _ShardedGradScaler |
| 132 | 133 | ||
| 133 | 134 | ||
| 134 | def _apply_class_patches(): | 135 | def _apply_class_patches(): |
| @@ -144,6 +145,17 @@ def _apply_class_patches(): | |||
| 144 | apply_sanitizer_patch() | 145 | apply_sanitizer_patch() |
| 145 | _apply_sharded_grad_scaler_patch() | 146 | _apply_sharded_grad_scaler_patch() |
| 146 | add_perf_dump_patch() | 147 | add_perf_dump_patch() |
| 148 | + _apply_distributed_methods_patch() | ||
| 149 | + | ||
| 150 | + | ||
| 151 | +def _apply_distributed_methods_patch(): | ||
| 152 | + torch._C._distributed_c10d._verify_params_across_processes = torch_npu.distributed._verify_params_across_processes | ||
| 153 | + torch.distributed.batch_isend_irecv = torch_npu.distributed.distributed_c10d._batch_isend_irecv | ||
| 154 | + torch.distributed.distributed_c10d.batch_isend_irecv = torch_npu.distributed.distributed_c10d._batch_isend_irecv | ||
| 155 | + torch.distributed.gather = torch_npu.distributed.distributed_c10d._gather | ||
| 156 | + torch.distributed.distributed_c10d.gather = torch_npu.distributed.distributed_c10d._gather | ||
| 157 | + torch.distributed.gather_object = torch_npu.distributed.distributed_c10d._gather_object | ||
| 158 | + torch.distributed.distributed_c10d.gather_object = torch_npu.distributed.distributed_c10d._gather_object | ||
D | |||
| 147 | 159 | ||
| 148 | 160 | ||
| 149 | torch.utils.rename_privateuse1_backend("npu") | 161 | torch.utils.rename_privateuse1_backend("npu") |
| @@ -329,7 +329,7 @@ public: | |||
| 329 | at::Tensor& inputTensors, | 329 | at::Tensor& inputTensors, |
| 330 | const c10d::ReduceOptions& opts = c10d::ReduceOptions()); | 330 | const c10d::ReduceOptions& opts = c10d::ReduceOptions()); |
| 331 | 331 | ||
| 332 | - c10::intrusive_ptr<c10d::Work>batch_isend_irecv( | 332 | + c10::intrusive_ptr<c10d::Work> batch_isend_irecv( |
| 333 | std::vector<std::string>& op_type, | 333 | std::vector<std::string>& op_type, |
| 334 | std::vector<at::Tensor>& tensors, | 334 | std::vector<at::Tensor>& tensors, |
| 335 | std::vector<uint32_t> remote_rank_list); | 335 | std::vector<uint32_t> remote_rank_list); |
| @@ -2,7 +2,7 @@ import torch_npu | |||
| 2 | from torch_npu.utils._error_code import ErrCode, dist_error | 2 | from torch_npu.utils._error_code import ErrCode, dist_error |
| 3 | 3 | ||
| 4 | __all__ = [ | 4 | __all__ = [ |
| 5 | - "batch_isend_irecv", "gather", "gather_object", "is_hccl_available", "_verify_params_across_processes", "reinit_process_group" | 5 | + "is_hccl_available", "_verify_params_across_processes", "reinit_process_group" |
| 6 | ] | 6 | ] |
| 7 | 7 | ||
| 8 | 8 | ||
| @@ -30,6 +30,6 @@ from torch_npu._C._distributed_c10d import ( | |||
| 30 | 30 | ||
| 31 | 31 | ||
| 32 | from torch_npu.distributed import rendezvous | 32 | from torch_npu.distributed import rendezvous |
| 33 | -from .distributed_c10d import batch_isend_irecv, gather, gather_object, is_hccl_available, reinit_process_group | 33 | +from .distributed_c10d import is_hccl_available, reinit_process_group |
| 34 | 34 | ||
| 35 | rendezvous._rendezvous_init() | 35 | rendezvous._rendezvous_init() |
| @@ -11,10 +11,10 @@ from torch.distributed.distributed_c10d import _get_default_group, get_group_ran | |||
| 11 | get_backend, GatherOptions, _update_default_pg, _world, _unregister_all_process_groups, _pg_map,\ | 11 | get_backend, GatherOptions, _update_default_pg, _world, _unregister_all_process_groups, _pg_map,\ |
| 12 | ProcessGroup, default_pg_timeout, _unregister_process_group | 12 | ProcessGroup, default_pg_timeout, _unregister_process_group |
| 13 | 13 | ||
| 14 | -__all__ = ["batch_isend_irecv", "gather", "gather_object", "is_hccl_available", "reinit_process_group"] | 14 | +__all__ = ["is_hccl_available", "reinit_process_group"] |
| 15 | 15 | ||
| 16 | 16 | ||
| 17 | -def batch_isend_irecv(p2p_op_list): | 17 | +def _batch_isend_irecv(p2p_op_list): |
| 18 | group = p2p_op_list[0].group | 18 | group = p2p_op_list[0].group |
| 19 | device = p2p_op_list[0].tensor.device | 19 | device = p2p_op_list[0].tensor.device |
| 20 | is_multi_pg = True | 20 | is_multi_pg = True |
| @@ -47,7 +47,7 @@ def batch_isend_irecv(p2p_op_list): | |||
| 47 | return reqs | 47 | return reqs |
| 48 | 48 | ||
| 49 | 49 | ||
| 50 | -def gather(tensor, gather_list=None, dst=0, group=None, async_op=False): | 50 | +def _gather(tensor, gather_list=None, dst=0, group=None, async_op=False): |
| 51 | """ | 51 | """ |
| 52 | Gathers a list of tensors in a single process. | 52 | Gathers a list of tensors in a single process. |
| 53 | 53 | ||
| @@ -130,7 +130,7 @@ def gather(tensor, gather_list=None, dst=0, group=None, async_op=False): | |||
| 130 | return None | 130 | return None |
| 131 | 131 | ||
| 132 | 132 | ||
| 133 | -def gather_object(obj, object_gather_list=None, dst=0, group=None): | 133 | +def _gather_object(obj, object_gather_list=None, dst=0, group=None): |
| 134 | """ | 134 | """ |
| 135 | Note: | 135 | Note: |
| 136 | Avoid gather_object to use gather func defined in origin distributed_c10d. | 136 | Avoid gather_object to use gather func defined in origin distributed_c10d. |
| @@ -172,7 +172,7 @@ def gather_object(obj, object_gather_list=None, dst=0, group=None): | |||
| 172 | for i in range(group_size) | 172 | for i in range(group_size) |
| 173 | ] | 173 | ] |
| 174 | # All ranks call gather with equal-sized tensors. | 174 | # All ranks call gather with equal-sized tensors. |
| 175 | - gather( | 175 | + _gather( |
| 176 | input_tensor, | 176 | input_tensor, |
| 177 | gather_list=output_tensors if my_rank == dst else None, # type: ignore[possibly-undefined] | 177 | gather_list=output_tensors if my_rank == dst else None, # type: ignore[possibly-undefined] |
| 178 | dst=dst, | 178 | dst=dst, |
| @@ -9,7 +9,7 @@ from torch.distributed.rpc import constants as rpc_constants | |||
| 9 | import torch_npu._C | 9 | import torch_npu._C |
| 10 | from torch_npu.utils._error_code import ErrCode, dist_error | 10 | from torch_npu.utils._error_code import ErrCode, dist_error |
| 11 | 11 | ||
| 12 | -__all__ = ["construct_rpc_backend_options", "init_backend"] | 12 | +__all__ = [] |
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | def _get_device_count_info(): | 15 | def _get_device_count_info(): |
| @@ -176,7 +176,7 @@ def _backend_type_repr(self): | |||
| 176 | return "BackendType." + self.name | 176 | return "BackendType." + self.name |
| 177 | 177 | ||
| 178 | 178 | ||
| 179 | -def construct_rpc_backend_options( | 179 | +def _construct_rpc_backend_options( |
| 180 | backend, | 180 | backend, |
| 181 | rpc_timeout=rpc_constants.DEFAULT_RPC_TIMEOUT_SEC, | 181 | rpc_timeout=rpc_constants.DEFAULT_RPC_TIMEOUT_SEC, |
| 182 | init_method=rpc_constants.DEFAULT_INIT_METHOD, | 182 | init_method=rpc_constants.DEFAULT_INIT_METHOD, |
| @@ -188,7 +188,7 @@ def construct_rpc_backend_options( | |||
| 188 | ) | 188 | ) |
| 189 | 189 | ||
| 190 | 190 | ||
| 191 | -def init_backend(backend, *args, **kwargs): | 191 | +def _init_backend(backend, *args, **kwargs): |
| 192 | return backend.value.init_backend_handler(*args, **kwargs) | 192 | return backend.value.init_backend_handler(*args, **kwargs) |
| 193 | 193 | ||
| 194 | 194 | ||
| @@ -1,7 +1,6 @@ | |||
| 1 | __all__ = [ | 1 | __all__ = [ |
| 2 | - "autocast", "GradScaler", "custom_fwd", "custom_bwd", "ShardedGradScaler" | 2 | + "autocast", "GradScaler", "custom_fwd", "custom_bwd" |
| 3 | ] | 3 | ] |
| 4 | 4 | ||
| 5 | from .autocast_mode import autocast, custom_fwd, custom_bwd # noqa: F401 | 5 | from .autocast_mode import autocast, custom_fwd, custom_bwd # noqa: F401 |
| 6 | -from .grad_scaler import GradScaler # noqa: F401 | 6 | +from .grad_scaler import GradScaler # noqa: F401 |
| 7 | -from .sharded_grad_scaler import ShardedGradScaler | ||
| @@ -10,7 +10,7 @@ from torch_npu.npu.utils import npu_check_overflow | |||
| 10 | from torch_npu.npu.amp.grad_scaler import _MultiDeviceReplicator, GradScaler, OptState | 10 | from torch_npu.npu.amp.grad_scaler import _MultiDeviceReplicator, GradScaler, OptState |
| 11 | from torch_npu.utils._error_code import ErrCode, pta_error | 11 | from torch_npu.utils._error_code import ErrCode, pta_error |
| 12 | 12 | ||
| 13 | -__all__ = ["ShardedGradScaler"] | 13 | +__all__ = [] |
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | log = logging.getLogger(__name__) | 16 | log = logging.getLogger(__name__) |
| @@ -37,7 +37,7 @@ class _GeneralMultiDeviceReplicator(_MultiDeviceReplicator): | |||
| 37 | self._per_device_tensors: Dict[torch.device, torch.Tensor] = {} | 37 | self._per_device_tensors: Dict[torch.device, torch.Tensor] = {} |
| 38 | 38 | ||
| 39 | 39 | ||
| 40 | -class ShardedGradScaler(GradScaler): | 40 | +class _ShardedGradScaler(GradScaler): |
| 41 | """ | 41 | """ |
| 42 | ShardedGradScaler helps perform gradient scaling in a shard aware manner. It extends | 42 | ShardedGradScaler helps perform gradient scaling in a shard aware manner. It extends |
| 43 | functionality from GradScaler: | 43 | functionality from GradScaler: |
改为非公开接口之后，_apply_distributed_patches会忽略这些方法？