Merged
[fix][fully_shard] adjust prefetch, unshard position #707
MengXY107Add on May 28
[fix][fully_shard] adjust prefetch, unshard position #707
Merged
Total: 2 files changed+7-8
| @@ -178,12 +178,12 @@ class HSDPSchedulerV2: | |||
| 178 | cast_fn = functools.partial(self.platform.cast_fp_tensor, self.mp_policy.param_dtype) | 178 | cast_fn = functools.partial(self.platform.cast_fp_tensor, self.mp_policy.param_dtype) |
| 179 | args = self.platform.apply_to_tensors(cast_fn, args) | 179 | args = self.platform.apply_to_tensors(cast_fn, args) |
| 180 | kwargs = self.platform.apply_to_tensors(cast_fn, kwargs) | 180 | kwargs = self.platform.apply_to_tensors(cast_fn, kwargs) |
| 181 | + with self.platform.profiler_record(f"pre_forward unshard:{self.hsdp_state.module_name}"): | ||
| 182 | + self.hsdp_state.unshard() | ||
| 181 | for prefetch_cell in self.forward_prefetch_cells: | 183 | for prefetch_cell in self.forward_prefetch_cells: |
| 182 | with self.platform.profiler_record(f"pre_forward prefetch:" | 184 | with self.platform.profiler_record(f"pre_forward prefetch:" |
| 183 | f"{prefetch_cell.hsdp_scheduler.hsdp_state.module_name}"): | 185 | f"{prefetch_cell.hsdp_scheduler.hsdp_state.module_name}"): |
| 184 | prefetch_cell.hsdp_scheduler.hsdp_state.prefetch() | 186 | prefetch_cell.hsdp_scheduler.hsdp_state.prefetch() |
| 185 | - with self.platform.profiler_record(f"pre_forward unshard:{self.hsdp_state.module_name}"): | ||
| 186 | - self.hsdp_state.unshard() | ||
| 187 | return args, kwargs | 187 | return args, kwargs |
| 188 | 188 | ||
| 189 | def _lazy_init_all_states(self): | 189 | def _lazy_init_all_states(self): |
| @@ -241,13 +241,13 @@ class HSDPSchedulerV2: | |||
| 241 | def _hsdp_backward_pre_hook(self, cell, grad_outputs): | 241 | def _hsdp_backward_pre_hook(self, cell, grad_outputs): |
| 242 | """Backward pre hook to unsharded parameter for backward process.""" | 242 | """Backward pre hook to unsharded parameter for backward process.""" |
| 243 | self.scheduler_state = FSDPSchedulerState.PRE_BACKWARD | 243 | self.scheduler_state = FSDPSchedulerState.PRE_BACKWARD |
| 244 | + if self.reshard_after_forward: | ||
| 245 | + with self.platform.profiler_record(f"pre_backward unshard:{self.hsdp_state.module_name}"): | ||
| 246 | + self.hsdp_state.unshard(unshard_replicate=False) | ||
| 244 | for prefetch_cell in self.backward_prefetch_cells: | 247 | for prefetch_cell in self.backward_prefetch_cells: |
| 245 | with self.platform.profiler_record(f"pre_backward prefetch:" | 248 | with self.platform.profiler_record(f"pre_backward prefetch:" |
| 246 | f"{prefetch_cell.hsdp_scheduler.hsdp_state.module_name}"): | 249 | f"{prefetch_cell.hsdp_scheduler.hsdp_state.module_name}"): |
| 247 | prefetch_cell.hsdp_scheduler.hsdp_state.prefetch(unshard_replicate=False) | 250 | prefetch_cell.hsdp_scheduler.hsdp_state.prefetch(unshard_replicate=False) |
| 248 | - if self.reshard_after_forward: | ||
| 249 | - with self.platform.profiler_record(f"pre_backward unshard:{self.hsdp_state.module_name}"): | ||
| 250 | - self.hsdp_state.unshard(unshard_replicate=False) | ||
| 251 | 251 | ||
| 252 | # pylint: disable=W0613 | 252 | # pylint: disable=W0613 |
| 253 | def _hsdp_backward_hook(self, cell, grad_inputs, grad_outputs): | 253 | def _hsdp_backward_hook(self, cell, grad_inputs, grad_outputs): |
| @@ -18,7 +18,6 @@ from typing import Any, Optional, Union | |||
| 18 | import dataclasses | 18 | import dataclasses |
| 19 | from collections import OrderedDict | 19 | from collections import OrderedDict |
| 20 | 20 | ||
| 21 | -import contextlib | ||
| 22 | import numpy as np | 21 | import numpy as np |
| 23 | import mindspore as ms | 22 | import mindspore as ms |
| 24 | import mindspore.common.dtype as mstype | 23 | import mindspore.common.dtype as mstype |
| @@ -29,10 +28,10 @@ from mindspore import mint | |||
| 29 | from mindspore.common.api import _no_grad | 28 | from mindspore.common.api import _no_grad |
| 30 | from mindspore.common._grad_function import _Function | 29 | from mindspore.common._grad_function import _Function |
| 31 | from mindspore.common.dtype import type_size_in_bytes | 30 | from mindspore.common.dtype import type_size_in_bytes |
| 31 | +from mindspore.common.recompute import null_context_fn | ||
| 32 | from mindspore.common.parameter import Parameter | 32 | from mindspore.common.parameter import Parameter |
| 33 | from mindspore.common.tensor import Tensor | 33 | from mindspore.common.tensor import Tensor |
| 34 | from mindspore.common.initializer import initializer | 34 | from mindspore.common.initializer import initializer |
| 35 | -from mindspore.common.recompute import null_context_fn | ||
| 36 | from mindspore.communication import GlobalComm | 35 | from mindspore.communication import GlobalComm |
| 37 | from mindspore.communication import get_group_size | 36 | from mindspore.communication import get_group_size |
| 38 | from mindspore.communication import create_group as new_group | 37 | from mindspore.communication import create_group as new_group |
| @@ -1780,7 +1779,7 @@ class MindSporePlatform(Platform): | |||
| 1780 | @staticmethod | 1779 | @staticmethod |
| 1781 | def profiler_record(name): | 1780 | def profiler_record(name): |
| 1782 | """Profiler context manager for recording operations using mindspore.profiler.""" | 1781 | """Profiler context manager for recording operations using mindspore.profiler.""" |
| 1783 | - return contextlib.nullcontext() | 1782 | + return ms.profiler.common.record_function.RecordFunction(name) |
| 1784 | 1783 | ||
| 1785 | def str_to_dtype(self, dtype_str: str) -> Any: | 1784 | def str_to_dtype(self, dtype_str: str) -> Any: |
| 1786 | """Resolve checkpoint dtype strings (``mindspore.*`` or short ``str(Tensor.dtype)`` e.g. ``Float32``).""" | 1785 | """Resolve checkpoint dtype strings (``mindspore.*`` or short ``str(Tensor.dtype)`` e.g. ``Float32``).""" |