Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / fully_shard / param.py: 46%
567 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""MindSpore fully_shard parameter lifecycle and gradient communication."""
16from dataclasses import dataclass
17from typing import Any, Callable, List, Optional, Tuple, cast
18import itertools
19import mindspore as ms
20from mindspore import nn
21from mindspore.common.api import _no_grad
22from mindspore import Parameter
23import mindspore.mint.distributed as dist
24from hyper_parallel.core.fully_shard.utils import (
25 MixedPrecisionPolicy,
26 CPUOffloadPolicy,
27 OffloadPolicy,
28 DataParallelMeshInfo,
29 DDPMeshInfo,
30 FSDPMeshInfo,
31 HSDPMeshInfo,
32 SourceShardMetaInfo,
33)
34from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
35from hyper_parallel.core.dtensor.dtensor import DTensor
36from hyper_parallel.core.dtensor.layout import Layout
37from hyper_parallel.core.fully_shard.hsdp_param import HSDPParamV2
38from hyper_parallel.core.fully_shard.hsdp_utils import (
39 ShardedState,
40 apply_gradient_scaling_factor,
41 unwrap_dtensor_param,
42)
43from hyper_parallel.core.dtensor.placement_types import Replicate, Shard, StridedShard
44from hyper_parallel.core.fully_shard.hsdp_utils import ParamModuleInfo
45from hyper_parallel.platform.mindspore.fully_shard._version_utils import copy_without_bumping_version
46from hyper_parallel.platform.mindspore.utils import normalize_runtime_device
47from hyper_parallel.platform.mindspore.fully_shard.pack_utils import (
48 build_rs_plan,
49 pack_for_reduce_scatter,
50 unpack_from_all_gather,
51)
54def _pack_for_reduce_scatter(local_tensor: ms.Tensor, shard_dim: int, world_size: int) -> ms.Tensor:
55 """Pack one local gradient into the row-major reduce-scatter layout.
57 MindSpore currently aligns with the torch non-comm-fusion V1 path:
59 - shard on dim 0: identity flatten
60 - shard on non-dim0: chunk on shard dim, then concatenate on dim 0
61 """
62 if world_size <= 1 or shard_dim == 0:
63 return local_tensor
64 chunks = ms.mint.chunk(local_tensor, world_size, dim=shard_dim)
65 return ms.mint.cat(chunks, dim=0).contiguous()
68def _to_dtype_if_needed(
69 tensor: ms.Tensor, dtype: Optional[ms.Type]
70) -> ms.Tensor:
71 """Cast tensor to the given dtype if it differs from current dtype."""
72 if isinstance(dtype, ms.Type) and tensor.dtype != dtype:
73 return tensor.to(dtype)
74 return tensor
77def make_contiguous_strides_for(shape, row_major=True):
78 """
79 Compute strides for a contiguous tensor of the given shape.
81 Args:
82 shape (tuple of int): The shape of the tensor. Each dimension must be a non-negative integer.
83 row_major (bool):
84 - If True (default), returns C-style (row-major) strides: last dimension changes fastest.
85 - If False, returns strides where the last two dimensions are Fortran-style
86 (i.e., for batched matrix operations in BLAS/LAPACK): second-to-last dim changes fastest.
88 Returns:
89 tuple of int: The computed strides.
91 Examples:
92 >>> make_contiguous_strides_for((2, 3, 4))
93 (12, 4, 1)
94 >>> make_contiguous_strides_for((2, 3, 4), row_major=False)
95 (12, 1, 3)
96 >>> make_contiguous_strides_for((5,))
97 (1,)
98 >>> make_contiguous_strides_for((5,), row_major=False)
99 (1,)
100 >>> make_contiguous_strides_for(())
101 ()
102 """
103 if not isinstance(shape, (tuple, list)):
104 raise TypeError("shape must be a tuple or list of non-negative integers")
106 # Validate shape elements
107 for dim in shape:
108 if not isinstance(dim, int) or dim < 0:
109 raise ValueError("All dimensions in shape must be non-negative integers")
111 if not shape:
112 return ()
114 # Compute C-style (row-major) strides: stride[i] = product(shape[i+1:])
115 strides = []
116 multiplier = 1
117 # Traverse shape in reverse order
118 for size in reversed(shape):
119 strides.append(multiplier)
120 multiplier *= max(size, 1) # handle size=0 gracefully (treat as 1 for stride calc)
122 # Reverse to get correct order
123 c_strides = tuple(reversed(strides))
125 if row_major:
126 return c_strides
127 # For column-major: only affect last two dimensions
128 if len(shape) < 2:
129 return c_strides
130 # In Fortran-style for matrices:
131 # stride of last dim = 1
132 # stride of second-to-last dim = shape[-1]
133 # But note: in batched case (..., M, N), we want strides (..., N, 1) → wait!
134 # However, the original PyTorch logic returns: result[:-2] + (1, max(shape[-2], 1))
135 # Let's follow that exactly:
136 # Example: shape=(B, M, N) → c_strides=(M*N, N, 1)
137 # col-major → (M*N, 1, M)
138 # So: keep all but last two, then (1, shape[-2])
139 return c_strides[:-2] + (1, max(shape[-2], 1))
142@dataclass
143class ReduceScatterCommCtx:
144 """Per-parameter reduce-scatter output and asynchronous work."""
146 reduce_scatter_output: Optional[ms.Tensor] = None
147 reduce_scatter_handle: Optional[Any] = None
150@dataclass
151class AllReduceCommCtx:
152 """Per-parameter all-reduce output and asynchronous work."""
154 all_reduce_output: Optional[ms.Tensor] = None
155 all_reduce_handle: Optional[Any] = None
158@dataclass
159class AllGatherCommCtx:
160 """Per-parameter all-gather buffers and asynchronous work."""
162 allgather_input: Optional[ms.Tensor] = None
163 allgather_output: Optional[ms.Tensor] = None
164 allgather_handle: Optional[Any] = None
167class MindSporeHSDPParamV2(HSDPParamV2):
168 """
169 MindSpore HSDP parameter.
170 """
172 def __init__(
173 self,
174 param: Parameter,
175 module_info: ParamModuleInfo,
176 mesh_info: DataParallelMeshInfo,
177 shard_placement_fn: Optional[Callable[[Parameter], Optional[Shard]]] = None,
178 mp_policy: Optional[MixedPrecisionPolicy] = None,
179 offload_policy: Optional[OffloadPolicy] = None,
180 device: Optional[str] = None,
181 source_shard_info: Optional[SourceShardMetaInfo] = None,
182 ):
183 self._module_info: ParamModuleInfo = module_info
184 self.mesh_info = mesh_info
185 self.mp_policy = mp_policy
186 self.device = device
187 self.orig_dtype = None
188 self.param_dtype = None
189 self.reduce_dtype = None
190 self.offload_to_cpu: bool = isinstance(offload_policy, CPUOffloadPolicy)
191 self.pin_memory = (
192 self.offload_to_cpu and cast(CPUOffloadPolicy, offload_policy).pin_memory
193 )
194 self._orig_param_hooks: List[Callable] = []
195 self.grad_offload_event: Optional[ms.runtime.Event] = None
196 dtensor_payload = unwrap_dtensor_param(param)
197 if (dtensor_payload is not None) != (
198 source_shard_info is not None and source_shard_info.origin_is_dtensor
199 ):
200 raise ValueError(
201 "source_shard_info.origin_is_dtensor must be True exactly for native DTensor parameters, "
202 f"got parameter type {type(param).__name__} and source_shard_info={source_shard_info}"
203 )
204 self.source_shard_info = source_shard_info
205 self._orig_param_is_dtensor = (
206 source_shard_info is not None and source_shard_info.origin_is_dtensor
207 )
208 self._orig_dtensor_mesh = source_shard_info.mesh if self._orig_param_is_dtensor else None
209 self._orig_dtensor_placements = (
210 tuple(source_shard_info.placements) if self._orig_param_is_dtensor else None
211 )
212 self._spmd_shard_mesh_dim = self.mesh_info.shard_mesh_dim
213 self._spmd_replicate_mesh_dim = self.mesh_info.replicate_mesh_dim
214 self._init_sharded_param(param, shard_placement_fn)
215 self._save_backward_hooks(param)
216 self.unsharded_param_buffers: List[ms.Tensor] = []
217 self.unsharded_accumulated_grad = None
218 self._unsharded_param: Optional[Parameter] = None
219 self._param_fqn: Optional[str] = None
220 # Communication attributes for prefetch pattern
221 self.allgather_comm_ctx = AllGatherCommCtx()
222 self.reduce_scatter_comm_ctx = ReduceScatterCommCtx()
223 self.all_reduce_comm_ctx = AllReduceCommCtx()
224 self._accumulated_allreduced_grad = True
225 self._reduce_partial_output = None
226 self._post_load_hook_handle = (
227 module_info.module.register_load_state_dict_post_hook(
228 lambda *args, **kwargs: self.reset_sharded_param()
229 )
230 )
231 self.gradient_scaling_factor = None
233 @property
234 def accumulated_allreduced_grad(self) -> bool:
235 return self._accumulated_allreduced_grad
237 @accumulated_allreduced_grad.setter
238 def accumulated_allreduced_grad(self, value: bool) -> None:
239 self._accumulated_allreduced_grad = value
241 @property
242 def reduce_partial_output(self) -> Optional[ms.Tensor]:
243 """Return reduce-scatter results accumulated before the final micro-step."""
244 return self._reduce_partial_output
246 @reduce_partial_output.setter
247 def reduce_partial_output(self, value: Optional[ms.Tensor]) -> None:
248 self._reduce_partial_output = value
250 def reduce_comm_dtype(self, grad: Optional[ms.Tensor] = None):
251 """Resolve the communication dtype owned by this parameter."""
252 if self.reduce_dtype is not None:
253 return self.reduce_dtype
254 if grad is not None:
255 return grad.dtype
256 if self.unsharded_accumulated_grad is not None:
257 return self.unsharded_accumulated_grad_data.dtype
258 if self.unsharded_param is not None and self.unsharded_param.grad is not None:
259 return self.unsharded_grad_data.dtype
260 return self.orig_dtype
262 def _get_base_spmd_placements(self) -> tuple:
263 """Return source-layout placements prefixed by explicit data-parallel axes."""
264 if self.source_shard_info is not None:
265 self._spmd_mesh = DeviceMesh.concatenate(
266 [self.mesh_info.mesh, self.source_shard_info.mesh]
267 )
268 dp_prefix = tuple(Replicate() for _ in range(self.mesh_info.mesh.ndim))
269 return dp_prefix + tuple(self.source_shard_info.placements)
270 self._spmd_mesh = self.mesh_info.mesh
271 return tuple(Replicate() for _ in range(self._spmd_mesh.ndim))
273 def _apply_data_parallel_placements(
274 self, placements: list, shard_placement: Shard
275 ) -> tuple:
276 """Apply the parameter-specific DDP/FSDP layout to source placements."""
277 if len(placements) != self._spmd_mesh.ndim:
278 raise AssertionError(
279 f"Expected {self._spmd_mesh.ndim} unified placements, got "
280 f"{len(placements)}: {placements}"
281 )
282 if (
283 isinstance(self.mesh_info, DDPMeshInfo)
284 and self._spmd_replicate_mesh_dim is not None
285 and not self._orig_param_is_dtensor
286 ):
287 placements[self._spmd_replicate_mesh_dim] = Replicate()
288 if isinstance(self.mesh_info, FSDPMeshInfo) and self._spmd_shard_mesh_dim is not None:
289 placements[self._spmd_shard_mesh_dim] = self._get_data_parallel_shard_placement(
290 placements, shard_placement
291 )
292 return tuple(placements)
294 def _get_data_parallel_shard_placement(self, placements: list, shard_placement: Shard):
295 """Return the explicit fully_shard placement on the unified SPMD mesh."""
296 split_factor = 1
297 shard_mesh_dim = getattr(self, "_spmd_shard_mesh_dim", None)
298 for mesh_idx, placement in enumerate(placements):
299 if mesh_idx == shard_mesh_dim:
300 continue
301 if placement.is_shard(shard_placement.dim):
302 split_factor *= self._spmd_mesh.mesh_shape[mesh_idx]
303 if split_factor > 1:
304 return StridedShard(shard_placement.dim, split_factor=split_factor)
305 return shard_placement
307 def _release_full_param_storage_if_safe(self, param_data: ms.Tensor) -> None:
308 """Release the temporary full-parameter storage once the sharded param is installed.
310 Skip storage reclamation only for meta tensors. Both plain Tensor inputs and DTensor local
311 tensors should drop their original storage after the sharded Parameter has been installed
312 onto the owning modules.
313 """
314 if param_data.is_meta:
315 return
316 storage = param_data.untyped_storage()
317 if storage.size() != 0:
318 storage.resize_(0)
320 def _iter_backward_hooks(self, param: Parameter) -> List[Callable]:
321 """Return backward hooks registered on a MindSpore Tensor/Parameter."""
322 hooks_getter = getattr(param, "hooks", None)
323 if callable(hooks_getter):
324 try:
325 return list(hooks_getter())
326 except (AttributeError, RuntimeError, TypeError, ValueError):
327 pass
329 backward_hooks = getattr(param, "_backward_hooks", None)
330 if backward_hooks is None:
331 return []
332 if hasattr(backward_hooks, "values"):
333 return list(backward_hooks.values())
334 return list(backward_hooks)
336 def _save_backward_hooks(self, param: Parameter) -> None:
337 """Save user-registered parameter backward hooks for later parameter swaps."""
338 if not hasattr(self, "_orig_param_hooks"):
339 self._orig_param_hooks = []
340 if not hasattr(self, "_saved_hook_ids"):
341 self._saved_hook_ids = set()
343 for hook_func in self._iter_backward_hooks(param):
344 hook_func_id = id(hook_func)
345 if hook_func_id not in self._saved_hook_ids:
346 self._orig_param_hooks.append(hook_func)
347 self._saved_hook_ids.add(hook_func_id)
349 def _migrate_backward_hooks(self, new_param: Parameter) -> None:
350 """Migrate saved user backward hooks to the active sharded/unsharded parameter."""
351 if not getattr(self, "_orig_param_hooks", None):
352 return
353 if hasattr(new_param, "migrate_backward_hooks_run_once"):
354 return
355 register_hook = getattr(new_param, "register_hook", None)
356 if not callable(register_hook):
357 return
359 for hook_func in self._orig_param_hooks:
360 try:
361 if getattr(new_param, "requires_grad", False):
362 register_hook(hook_func)
363 except (RuntimeError, TypeError, ValueError):
364 pass
365 new_param.migrate_backward_hooks_run_once = True
367 @_no_grad()
368 def _init_sharded_param(
369 self,
370 param: Parameter,
371 shard_placement_fn: Optional[Callable],
372 ) -> None:
373 param_device = normalize_runtime_device(param.device)
374 if param_device not in ("meta", self.device):
375 raise AssertionError(
376 f"Expects the parameter to already be moved to device {self.device} but got {param.device}"
377 )
378 hsdp_placement = shard_placement_fn(param) if shard_placement_fn else None
379 if hsdp_placement is None:
380 hsdp_placement = Shard(0)
381 elif hsdp_placement.dim < 0:
382 # if dim is negative, add the number of dimensions of the parameter
383 hsdp_placement = Shard(hsdp_placement.dim + param.ndim)
385 if not isinstance(hsdp_placement, Shard):
386 raise AssertionError(
387 f"Expected Shard, got {type(hsdp_placement)}: {hsdp_placement}"
388 )
390 self.hsdp_placement = hsdp_placement
391 base_placements = list(self._get_base_spmd_placements())
392 self._spmd_placements = self._apply_data_parallel_placements(base_placements, hsdp_placement)
393 param_data = unwrap_dtensor_param(param).to_local() if self._orig_param_is_dtensor else param
395 shard_dim = hsdp_placement.dim
396 self._orig_size = param_data.shape
397 self._contiguous_orig_stride = make_contiguous_strides_for(self._orig_size)
399 if isinstance(self.mesh_info, FSDPMeshInfo):
400 self.shard_rank = self.mesh_info.shard_mesh_rank
401 self.shard_world_size = self.mesh_info.shard_mesh_size
402 else:
403 self.shard_rank = 0
404 self.shard_world_size = 1
405 if isinstance(self.mesh_info, DDPMeshInfo):
406 self.replicate_world_size = self.mesh_info.replicate_mesh_size
407 else:
408 self.replicate_world_size = 1
409 self.is_replicate_param = (
410 isinstance(self.mesh_info, DDPMeshInfo)
411 and not isinstance(self.mesh_info, HSDPMeshInfo)
412 )
413 self.is_sharded = self.shard_world_size > 1
415 if param_data.shape[shard_dim] % self.shard_world_size != 0:
416 raise NotImplementedError(
417 f"Uneven sharding on dim {shard_dim} not supported: "
418 f"shape={param_data.shape}, world_size={self.shard_world_size}"
419 )
420 chunks = ms.mint.chunk(param_data, self.shard_world_size, dim=shard_dim)
421 sharded_param = chunks[self.shard_rank].clone().contiguous()
422 self.sharded_size = sharded_param.shape
423 self.contiguous_sharded_stride = make_contiguous_strides_for(self.sharded_size)
424 self._sharded_param_data = sharded_param.view(-1)
426 self._sharding_spec = Layout.from_device_mesh(self._spmd_mesh)
427 self._sharding_spec.set_placements(self._spmd_placements)
428 self._sharding_spec.placement_to_tensor_map(param.ndim)
430 shard_dtensor = DTensor.from_local(sharded_param, self._spmd_mesh, self._spmd_placements)
431 self.sharded_param = Parameter(shard_dtensor, name=param.name)
432 set_requires_grad_if_needed(param, self.sharded_param)
433 self.sharded_param.grad = None
435 self._setattr_on_modules(self.sharded_param)
436 self._release_full_param_storage_if_safe(param_data)
437 self.sharded_param._hsdp_param_initialized = True
438 self.sharded_state = ShardedState.SHARDED
439 self.param_dtype = None
441 def init_dtype_attrs(self, mp_policy: MixedPrecisionPolicy):
442 param_dtype, reduce_dtype = (mp_policy.param_dtype, mp_policy.reduce_dtype)
443 self.orig_dtype = self.sharded_param.dtype
444 if reduce_dtype == param_dtype:
445 reduce_dtype = None
446 if param_dtype == self.orig_dtype:
447 param_dtype = None
448 self.param_dtype = param_dtype
449 self.reduce_dtype = reduce_dtype
451 def init_unsharded_param_buffers(
452 self,
453 all_gather_input_numels: list[int],
454 all_gather_input_dtypes: list[ms.Type],
455 world_size: int,
456 device: str,
457 force_recreate: bool = False,
458 ):
459 if not force_recreate and len(self.unsharded_param_buffers) > 0:
460 return # already initialized
461 self.unsharded_param_buffers = [
462 ms.mint.empty([numel * world_size], dtype=dtype, device=device.split(':')[0])
463 for numel, dtype in zip(all_gather_input_numels, all_gather_input_dtypes)
464 ]
466 def init_unsharded_param(self):
467 """
468 Initialize unsharded parameter from all-gather outputs.
470 This reconstructs the full parameter after all-gather by unpacking the
471 gathered flat buffer back to the original tensor layout.
472 """
473 unsharded_param = self._get_unsharded_param_from_all_gather_output()
474 if self._unsharded_param is not None:
475 # Keep the Parameter identity stable across forward-reshard-backward
476 # cycles so backward hooks continue to read gradients from the same
477 # object that participated in the forward graph.
478 if self._orig_param_is_dtensor:
479 self._unsharded_param.set_data(unsharded_param)
480 else:
481 self._unsharded_param.data = unsharded_param
482 set_requires_grad_if_needed(self.sharded_param, self._unsharded_param)
483 self._unsharded_param.grad = None
484 return
485 if self._orig_param_is_dtensor:
486 self._unsharded_param = Parameter(
487 unsharded_param,
488 name=self.sharded_param.name,
489 requires_grad=self.sharded_param.requires_grad,
490 )
491 return
492 # For MindSpore, if use `Parameter(tensor)`, Parameter will create a new Tensor instead of a view.
493 # Here we need to share storage, so we use the `.data = tensor` approach to create shared storage.
494 self._unsharded_param = Parameter(
495 [],
496 name=self.sharded_param.name,
497 requires_grad=False,
498 )
499 self._unsharded_param.data = unsharded_param
500 if self.sharded_param.requires_grad:
501 self._unsharded_param.requires_grad = True
503 def _get_unsharded_param_from_all_gather_output(self):
504 """Reconstruct the full local parameter view from the packed all-gather output."""
505 if len(self.unsharded_param_buffers) != 1:
506 raise AssertionError(
507 f"Expected 1 unsharded_param_buffer, got {len(self.unsharded_param_buffers)}"
508 )
509 unsharded_tensor = self.unsharded_param_buffers[0]
510 plan = build_rs_plan(
511 self,
512 self._sharded_local_tensor,
513 self.shard_world_size if self.is_sharded else 1,
514 )
515 unsharded_param = unpack_from_all_gather(unsharded_tensor, plan)
516 if getattr(self, "_orig_param_is_dtensor", False):
517 unsharded_param = DTensor.from_local(
518 unsharded_param,
519 self._orig_dtensor_mesh,
520 self._orig_dtensor_placements,
521 )
522 return unsharded_param
524 def to_sharded(self) -> None:
525 self._setattr_on_modules(self.sharded_param)
526 self.free_unsharded_param()
527 self.allgather_comm_ctx.allgather_input = None
528 self.allgather_comm_ctx.allgather_output = None
529 self.allgather_comm_ctx.allgather_handle = None
530 self.sharded_state = ShardedState.SHARDED
532 def to_unsharded(self) -> None:
533 set_requires_grad_if_needed(self.sharded_param, self._unsharded_param)
534 self._setattr_on_modules(self._unsharded_param)
535 self.sharded_state = ShardedState.UNSHARDED
537 def _setattr_on_modules(self, param: Parameter) -> None:
538 if getattr(self._module_info.module.__setattr__, "__func__", None) is nn.Cell.__setattr__:
539 # fast path
540 self._module_info.module._params[self._module_info.param_name] = param
541 else:
542 # slow path
543 setattr(self._module_info.module, self._module_info.param_name, param)
544 if hasattr(self, "sharded_param"):
545 self._save_backward_hooks(self.sharded_param)
546 self._migrate_backward_hooks(param)
548 # Iterate through all modules that share this parameter to prevent pointer desync.
549 for shared_module, shared_param_name in zip(
550 self._module_info.shared_modules, self._module_info.shared_param_names
551 ):
552 if getattr(shared_module.__setattr__, "__func__", None) is nn.Cell.__setattr__:
553 shared_module._params[shared_param_name] = param
554 else:
555 setattr(shared_module, shared_param_name, param)
557 def to_sharded_dtensor(self, tensor: ms.Tensor) -> DTensor:
558 """
559 Converts a local tensor representing either the sharded parameter or
560 sharded gradient to DTensor.
561 """
562 return DTensor.from_local(
563 tensor,
564 self._sharding_spec.mesh,
565 self._sharding_spec.placements
566 )
568 def _to_local_unsharded_grad(self, grad):
569 """Normalize a pending gradient to the local tensor expected by fully_shard collectives."""
570 return self._normalize_unsharded_grad_to_local(grad, reduce_partial_dtensor=False)
572 def to_accumulated_grad_if_needed(self) -> None:
573 if self._unsharded_param.grad is None:
574 return
575 unsharded_grad = self._unsharded_param.grad
576 self._unsharded_param.grad = None
577 if self.reduce_dtype is not None and unsharded_grad.dtype != self.reduce_dtype:
578 unsharded_grad = unsharded_grad.to(self.reduce_dtype)
579 if self.unsharded_accumulated_grad is None:
580 self.unsharded_accumulated_grad = unsharded_grad
581 else:
582 self.unsharded_accumulated_grad = ms.mint.add(
583 self.unsharded_accumulated_grad,
584 unsharded_grad,
585 )
587 def accumulate_unsharded_grad_if_needed(self) -> None:
588 if (
589 self.unsharded_accumulated_grad is not None
590 and self.unsharded_param.grad is not None
591 ):
592 # need to handle the gradient
593 self.unsharded_accumulated_grad = ms.mint.add(
594 self.unsharded_accumulated_grad,
595 self._to_local_unsharded_grad(self.unsharded_param.grad),
596 )
597 self.unsharded_param.grad = None
599 def alloc_unsharded_param_buffers(self) -> None:
600 for tensor in self.unsharded_param_buffers:
601 expected_size = tensor.numel() * tensor.itemsize
603 storage = tensor.untyped_storage()
604 if storage.size() != expected_size:
605 storage.resize_(expected_size)
607 def free_unsharded_param(self) -> None:
608 for tensor in itertools.chain(
609 self.unsharded_param_buffers
610 ):
611 storage = tensor.untyped_storage()
612 if storage.size() != 0:
613 storage.resize_(0)
615 @property
616 def all_gather_inputs(self) -> list[ms.Tensor]:
617 self._assert_in_states(ShardedState.SHARDED)
618 sharded_param_data = self._sharded_param_data
619 if self.offload_to_cpu:
620 sharded_param_data = sharded_param_data.to(
621 self.device, non_blocking=True
622 )
623 if self.param_dtype is not None and self.param_dtype != sharded_param_data.dtype:
624 return [sharded_param_data.to(self.param_dtype)]
625 return [sharded_param_data]
627 @property
628 def unsharded_param(self) -> Parameter:
629 """Return the full unsharded parameter after all-gather."""
630 return self._unsharded_param
632 @property
633 def unsharded_grad_data(self) -> ms.Tensor:
634 """
635 Get the unsharded gradient data as a local tensor.
636 """
637 grad = self.unsharded_param.grad
638 if grad is None:
639 raise AssertionError("Expects unsharded_param.grad to not be None")
640 return self._to_local_unsharded_grad(grad)
642 @property
643 def unsharded_accumulated_grad_data(self) -> ms.Tensor:
644 """
645 Get the unsharded accumulated gradient data as a local tensor.
646 """
647 grad = self.unsharded_accumulated_grad
648 return grad
650 @property
651 def _sharded_local_tensor(self) -> ms.Tensor:
652 """Return the underlying local tensor of the sharded DTensor parameter."""
653 return cast(DTensor, self.sharded_param)._local_tensor
655 def _sharded_param_storage_dtype(self) -> Optional[ms.Type]:
656 """Return the dtype of the sharded parameter's on-device storage."""
657 if not hasattr(self.sharded_param, "dtype"):
658 return None
659 dtype = self.sharded_param.dtype
660 if isinstance(dtype, ms.Type):
661 return dtype
662 return None
664 def _assert_in_states(self, *states: ShardedState) -> None:
665 """Assert current state is one of expected states."""
666 if self.sharded_state not in states:
667 raise AssertionError(
668 f"Expected sharded_state in {states}, got {self.sharded_state}"
669 )
671 def _is_same_sharded_local_tensor(self, local_tensor: ms.Tensor) -> bool:
672 """Whether the cached flat shard view already points to the ``local_tensor`` storage."""
673 if not isinstance(self._sharded_param_data, ms.Tensor):
674 return False
675 cached_storage = self._sharded_param_data.untyped_storage()
676 local_storage = local_tensor.untyped_storage()
677 # when sharding param with shape (1, ...) over 2 ranks
678 # local_tensor on rank 1 can be size 0, data_ptr() can be 0
679 return (
680 cached_storage.data_ptr() > 0
681 and cached_storage.data_ptr() == local_storage.data_ptr()
682 )
684 def _validate_sharded_local_tensor_shape(self, local_tensor: ms.Tensor) -> None:
685 """Validate that a replaced local tensor still matches the expected shard shape."""
686 if local_tensor.shape != self.sharded_size:
687 raise AssertionError(
688 f"Expected sharded_size to be {self.sharded_size}, got {local_tensor.shape}"
689 )
691 def _pin_sharded_local_tensor_if_needed(self, local_tensor: ms.Tensor) -> Tuple[ms.Tensor, bool]:
692 """Pin the local tensor memory when CPU offload requires it."""
693 if self.pin_memory and not local_tensor.is_pinned():
694 return local_tensor.to("cpu").pin_memory(), True
695 return local_tensor, False
697 def _assert_sharded_param_is_dtensor(self) -> None:
698 """Assert that ``self.sharded_param`` is backed by a DTensor."""
699 if not isinstance(self.sharded_param, DTensor):
700 raise AssertionError(f"Expected DTensor, got {type(self.sharded_param)}")
702 def _refresh_sharded_local_tensor_view(
703 self,
704 local_tensor: ms.Tensor,
705 shard_dim: int,
706 length: int,
707 ) -> None:
708 """Refresh ``self.sharded_param`` to point to a local tensor view."""
709 # Only change the local tensor object if needed
710 with _no_grad():
711 local_view = local_tensor.narrow(dim=shard_dim, start=0, length=length)
712 set_requires_grad_if_needed(self.sharded_param, local_view)
713 self.sharded_param._local_tensor = local_view
714 if not self.sharded_param._local_tensor.is_contiguous():
715 raise AssertionError(
716 "Expected sharded_param._local_tensor to be contiguous"
717 )
719 def reset_sharded_param(self) -> None:
720 """Reset the sharded param after ``load_state_dict``."""
721 module_info = self._module_info
722 new_param = getattr(module_info.module, module_info.param_name)
723 if new_param is not self.sharded_param:
724 if isinstance(new_param, DTensor):
725 self.sharded_param = new_param
726 if not getattr(self.sharded_param, "_hsdp_param_initialized", None):
727 # reset _hsdp_param_initialized flag.
728 self.sharded_param._hsdp_param_initialized = True
729 elif isinstance(new_param, ms.Tensor):
730 # if new_param is Tensor, don't re-ref 'self.sharded_param'
731 # just update self.sharded_param._local_tensor and self.sharded_param_data.
732 pass
734 local_tensor = new_param._local_tensor if isinstance(new_param, DTensor) else new_param
735 if local_tensor.is_meta:
736 return
737 # local_tensor can be padded twice
738 # 1st time in fully_shard(model)
739 # 2nd time in model(input) lazy_init
740 # 2nd time should be no-op if parameters remain unchanged
741 # 2nd time shouldn't be no-op if people call model.load_state_dict(...) before lazy_init
742 # this makes it possible for trainer to call `sd = model.state_dict()` before the training loop
743 # and use `sd` without calling .state_dict() per iteration
744 same_local_tensor = self._is_same_sharded_local_tensor(local_tensor)
745 shard_dim = self.hsdp_placement.dim
746 length = local_tensor.shape[shard_dim] if local_tensor.numel() > 0 else 0
747 if not same_local_tensor:
748 self._validate_sharded_local_tensor_shape(local_tensor)
749 local_tensor, pinned_local_tensor = self._pin_sharded_local_tensor_if_needed(local_tensor)
750 updated_local_tensor = not same_local_tensor or pinned_local_tensor
751 if not same_local_tensor:
752 self._sharded_param_data = local_tensor.view(-1)
753 self._assert_sharded_param_is_dtensor()
754 if updated_local_tensor:
755 self._refresh_sharded_local_tensor_view(local_tensor, shard_dim, length)
756 self._sharding_spec = cast(DTensor, self.sharded_param).layout
758 @_no_grad()
759 def _get_unsharded_param_data(
760 self,
761 async_op: bool = False,
762 ) -> Tuple[ms.Tensor, ms.Tensor, Optional[Any]]:
763 """
764 Perform all-gather to get unsharded parameter data.
766 Args:
767 async_op: Whether to execute asynchronously.
769 Returns:
770 (all_gather_input, unsharded_param, handle): Communication input,
771 unsharded parameter data, and communication handle.
772 """
773 # Optimizer steps may refresh the underlying local tensor storage. Re-sync
774 # the cached flat shard view before reading all_gather_inputs for the next
775 # unshard cycle.
776 self.reset_sharded_param()
777 all_gather_input = self.all_gather_inputs[0]
779 # If parameter is not sharded (below threshold), no communication needed
780 if not self.is_sharded:
781 self.init_unsharded_param_buffers(
782 all_gather_input_numels=[all_gather_input.numel()],
783 all_gather_input_dtypes=[all_gather_input.dtype],
784 world_size=1,
785 device=all_gather_input.device.split(':')[0],
786 )
787 self.alloc_unsharded_param_buffers()
788 copy_without_bumping_version(self.unsharded_param_buffers[0], all_gather_input)
789 return all_gather_input, self.unsharded_param_buffers[0], None
791 # Initialize output buffer
792 self.init_unsharded_param_buffers(
793 all_gather_input_numels=[all_gather_input.numel()],
794 all_gather_input_dtypes=[all_gather_input.dtype],
795 world_size=self.shard_world_size,
796 device=self._sharded_param_data.device.split(':')[0],
797 )
798 self.alloc_unsharded_param_buffers()
800 # Get communication group
801 shard_group = self.mesh_info.shard_process_group if isinstance(self.mesh_info, FSDPMeshInfo) else None
803 if shard_group is None or self.shard_world_size <= 1:
804 # No communication needed, just copy
805 copy_without_bumping_version(self.unsharded_param_buffers[0], all_gather_input)
806 return all_gather_input, self.unsharded_param_buffers[0], None
808 # Execute all_gather_into_tensor
809 handle = dist.all_gather_into_tensor(
810 self.unsharded_param_buffers[0],
811 all_gather_input,
812 group=shard_group,
813 async_op=async_op,
814 )
816 return all_gather_input, self.unsharded_param_buffers[0], handle
818 def unshard(self, async_op: bool = False) -> None:
819 if self.allgather_comm_ctx.allgather_output is not None:
820 # Already triggered by HSDPState.prefetch(), so return directly.
821 return # no-op
823 all_gather_input, output, handle = self._get_unsharded_param_data(async_op=async_op)
824 self.allgather_comm_ctx.allgather_input = all_gather_input
825 self.allgather_comm_ctx.allgather_output = output
826 self.allgather_comm_ctx.allgather_handle = handle
828 def wait_for_unshard(self) -> None:
829 self._assert_in_states(ShardedState.SHARDED)
831 if self.allgather_comm_ctx.allgather_handle is not None:
832 self.allgather_comm_ctx.allgather_handle.wait()
833 self.allgather_comm_ctx.allgather_handle = None
834 self.allgather_comm_ctx.allgather_input = None
836 self.init_unsharded_param()
837 self.to_unsharded()
839 def shard(self) -> None:
840 """
841 Transition parameter from unsharded back to sharded state.
842 """
843 self._assert_in_states(ShardedState.UNSHARDED)
844 self.to_sharded()
846 def reduce_scatter_output(self):
847 """Return cached reduce-scatter output after waiting pending async work."""
848 if self.reduce_scatter_comm_ctx.reduce_scatter_handle is not None:
849 self.reduce_scatter_comm_ctx.reduce_scatter_handle.wait()
850 self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
851 return self.reduce_scatter_comm_ctx.reduce_scatter_output
853 def clear_reduce_scatter_output(self):
854 """Clear cached reduce-scatter output."""
855 self.reduce_scatter_comm_ctx.reduce_scatter_output = None
857 def reduce_scatter_grad(
858 self,
859 async_op: bool = True,
860 dtype: Optional[ms.Type] = None,
861 reduce_op: str = "avg",
862 output_buffer: Optional[ms.Tensor] = None,
863 ) -> None:
864 """
865 Perform reduce-scatter on gradient to reduce and shard the full gradient.
867 Args:
868 async_op: Whether to execute asynchronously.
869 dtype: reduce dtype.
870 reduce_op: do reduce-scatter avg or sum.
871 output_buffer: Optional pre-allocated output for fused all-reduce groups.
873 The output and optional asynchronous handle are stored in
874 ``reduce_scatter_comm_ctx``.
875 """
876 # Choose gradient source based on use_accumulated_grad flag
877 if self.unsharded_accumulated_grad is not None:
878 grad = self.unsharded_accumulated_grad_data
879 else:
880 grad = self.unsharded_grad_data
881 reduce_dtype = dtype or self.reduce_comm_dtype(grad)
882 grad = grad.to(reduce_dtype)
883 grad = grad.contiguous()
884 shard_group_info = getattr(self, "sharded_group_info", None)
885 shard_group = shard_group_info.group if shard_group_info is not None else None
886 shard_group_size = shard_group_info.rank_size if shard_group_info is not None else 1
887 if shard_group is None and isinstance(self.mesh_info, FSDPMeshInfo):
888 shard_group = self.mesh_info.shard_process_group
889 shard_group_size = self.shard_world_size
890 plan_world_size = (
891 shard_group_size
892 if self.is_sharded and shard_group is not None and shard_group_size > 1
893 else 1
894 )
895 plan = build_rs_plan(self, grad, plan_world_size)
896 grad_flat = pack_for_reduce_scatter(grad, plan).reshape(-1)
897 # apply gradient_scaling_factor (reduce-scatter leg)
898 apply_gradient_scaling_factor(grad_flat, self.gradient_scaling_factor)
899 # If parameter is not sharded (below threshold), no reduce-scatter needed
900 if not self.is_sharded:
901 if output_buffer is not None:
902 copy_without_bumping_version(output_buffer, grad_flat)
903 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
904 else:
905 self.reduce_scatter_comm_ctx.reduce_scatter_output = grad_flat
906 self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
907 return
909 if shard_group is None or shard_group_size <= 1:
910 if output_buffer is not None:
911 copy_without_bumping_version(output_buffer, grad_flat)
912 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
913 else:
914 self.reduce_scatter_comm_ctx.reduce_scatter_output = grad_flat
915 self.reduce_scatter_comm_ctx.reduce_scatter_handle = None
916 return
918 # Calculate output size
919 output_numel = grad_flat.numel() // shard_group_size
920 if output_buffer is not None:
921 if output_buffer.numel() != output_numel:
922 raise ValueError(
923 f"output_buffer size mismatch: expected {output_numel}, got {output_buffer.numel()}"
924 )
925 if output_buffer.dtype != reduce_dtype:
926 raise ValueError(
927 f"output_buffer dtype mismatch: expected {reduce_dtype}, got {output_buffer.dtype}"
928 )
929 self.reduce_scatter_comm_ctx.reduce_scatter_output = output_buffer
930 else:
931 self.reduce_scatter_comm_ctx.reduce_scatter_output = ms.mint.empty(
932 output_numel, dtype=reduce_dtype, device=grad.device.split(":")[0]
933 )
935 # Ascend HCCL DistCommReduceScatter rejects non-contiguous tensors.
936 # ``pack_for_reduce_scatter`` on a shard-dim-0 path returns the input
937 # tensor as-is (potentially a view from to_local() / redistribute()),
938 # and the trailing ``.reshape(-1)`` may yield a view. Force contiguous
939 # storage here (no-op when already contig).
940 grad_flat = grad_flat.contiguous()
942 # Execute reduce_scatter_tensor
943 self.reduce_scatter_comm_ctx.reduce_scatter_handle = dist.reduce_scatter_tensor(
944 self.reduce_scatter_comm_ctx.reduce_scatter_output,
945 grad_flat,
946 op=reduce_op,
947 group=shard_group,
948 async_op=async_op,
949 )
951 def zero_grad(self):
952 """Reset the sharded parameter's gradient buffers to None."""
953 self.sharded_param.grad = None
954 if hasattr(self.sharded_param, "main_grad"):
955 self.sharded_param.main_grad = None
957 def all_reduce_grad(
958 self,
959 async_op: bool = True,
960 reduce_op: str = "avg",
961 ) -> None:
962 """
963 Perform all-reduce on gradient (across replicate dimension in HSDP mode).
965 Args:
966 async_op: Whether to execute asynchronously.
967 reduce_op: Reduction operation accepted by ``mint.distributed``.
969 The output and optional asynchronous handle are stored in
970 ``all_reduce_comm_ctx``.
971 """
972 grad = self.reduce_scatter_comm_ctx.reduce_scatter_output
973 if grad is None:
974 raise RuntimeError("all_reduce_grad requires a completed reduce-scatter output.")
975 if self.reduce_dtype is not None and self.reduce_dtype != grad.dtype:
976 grad = grad.to(self.reduce_dtype)
977 reduce_group = (
978 self.mesh_info.replicate_process_group
979 if isinstance(self.mesh_info, DDPMeshInfo)
980 else None
981 )
982 if reduce_group is None or self.replicate_world_size <= 1:
983 self.all_reduce_comm_ctx.all_reduce_output = grad
984 self.all_reduce_comm_ctx.all_reduce_handle = None
985 return
987 # Ascend HCCL DistCommAllReduce rejects non-contiguous tensors.
988 # ``grad`` here may be a view returned by ``_to_local_unsharded_grad``
989 # (DTensor.to_local() / redistribute().to_local()) or by autograd.
990 # ``Tensor.contiguous()`` is itself a no-op when storage is already
991 # contiguous, so the unconditional call is safe and avoids the
992 # ``is_contiguous()`` query (which has been observed to under-detect
993 # non-contig views from DTensor on this MS version).
994 grad = grad.contiguous()
996 self.all_reduce_comm_ctx.all_reduce_output = grad
997 self.all_reduce_comm_ctx.all_reduce_handle = dist.all_reduce(
998 grad,
999 op=reduce_op,
1000 group=reduce_group,
1001 async_op=async_op,
1002 )
1004 def all_reduce_output(self):
1005 """Return cached all-reduce output after waiting pending async work."""
1006 if self.all_reduce_comm_ctx.all_reduce_handle is not None:
1007 self.all_reduce_comm_ctx.all_reduce_handle.wait()
1008 self.all_reduce_comm_ctx.all_reduce_handle = None
1009 return self.all_reduce_comm_ctx.all_reduce_output
1011 def clear_all_reduce_output(self):
1012 """Clear cached all-reduce output."""
1013 self.all_reduce_comm_ctx.all_reduce_output = None
1015 def clear_unsharded_source_grad(self) -> None:
1016 """Release the unsharded gradient after its communication input is safe."""
1017 if self.unsharded_accumulated_grad is not None:
1018 self.unsharded_accumulated_grad = None
1019 if self.unsharded_param is not None and self.unsharded_param.grad is not None:
1020 self.unsharded_param.grad = None
1022 def apply_reduced_grad(self, reduced_grad):
1023 """
1024 Apply reduced gradient to the sharded parameter.
1026 Reshapes ``reduced_grad`` to match the local shard, optionally
1027 offloads to CPU, then accumulates or assigns onto ``grad`` or
1028 ``main_grad`` depending on the mixed-precision policy.
1029 Args:
1030 reduced_grad (ms.Tensor): Gradient after reduce-scatter
1031 and/or all-reduce.
1032 """
1033 if self.mp_policy.apply_grad_on_fp32_main_grad:
1034 if not hasattr(self.sharded_param, "main_grad"):
1035 self.sharded_param.main_grad = None
1036 sharded_grad = self.sharded_param.main_grad
1037 else:
1038 sharded_grad = self.sharded_param.grad
1040 reduced_grad = reduced_grad.reshape(-1).narrow(
1041 0, 0, self._sharded_local_tensor.numel()
1042 ).view(self.sharded_size)
1043 if not self.mp_policy.apply_grad_on_fp32_main_grad:
1044 reduced_grad = _to_dtype_if_needed(reduced_grad, self.orig_dtype)
1045 reduced_grad = _to_dtype_if_needed(
1046 reduced_grad, self._sharded_param_storage_dtype()
1047 )
1048 to_accumulate_grad = sharded_grad is not None
1049 need_synchronize = False
1050 if self.offload_to_cpu:
1051 non_blocking = self.pin_memory and not to_accumulate_grad
1052 reduced_grad = reduced_grad.to(
1053 "cpu", non_blocking=non_blocking
1054 )
1055 need_synchronize = True
1056 if sharded_grad is None:
1057 if self.mp_policy.apply_grad_on_fp32_main_grad:
1058 self.sharded_param.main_grad = self.to_sharded_dtensor(reduced_grad)
1059 self.sharded_param.grad = None
1060 else:
1061 self.sharded_param.grad = self.to_sharded_dtensor(reduced_grad)
1062 else:
1063 if self.mp_policy.apply_grad_on_fp32_main_grad:
1064 accumulated_grad = ms.mint.add(
1065 self.sharded_param.main_grad._local_tensor,
1066 reduced_grad,
1067 )
1068 self.sharded_param.main_grad = self.to_sharded_dtensor(accumulated_grad)
1069 self.sharded_param.grad = None
1070 else:
1071 accumulated_grad = ms.mint.add(
1072 self.sharded_param.grad._local_tensor,
1073 reduced_grad,
1074 )
1075 self.sharded_param.grad = self.to_sharded_dtensor(accumulated_grad)
1077 self.clear_unsharded_source_grad()
1078 return need_synchronize
1080 def all_reduce_tp_replicate_grad_inplace(
1081 self,
1082 reduced_grad: ms.Tensor,
1083 reduce_op: str,
1084 ) -> None:
1085 """All-reduce a final gradient over replicated source-layout axes."""
1086 if self.source_shard_info is None or not self.source_shard_info.placements:
1087 return
1088 source_mesh = self.source_shard_info.mesh
1089 replicate_mesh_dims = tuple(
1090 mesh_dim
1091 for mesh_dim, placement in enumerate(self.source_shard_info.placements)
1092 if placement.is_replicate()
1093 )
1094 if not replicate_mesh_dims:
1095 return
1096 if source_mesh.mesh_dim_names is None:
1097 raise ValueError(
1098 "TP shard mesh must define mesh_dim_names to all-reduce replicated gradients."
1099 )
1100 replicate_dim_names = tuple(
1101 source_mesh.mesh_dim_names[mesh_dim]
1102 for mesh_dim in replicate_mesh_dims
1103 )
1104 replicate_mesh = source_mesh[replicate_dim_names].flatten()
1105 if replicate_mesh.size() <= 1:
1106 return
1107 dist.all_reduce(
1108 reduced_grad,
1109 op=reduce_op,
1110 group=replicate_mesh.get_group(),
1111 async_op=False,
1112 )
1115def set_requires_grad_if_needed(
1116 src_tensor: ms.Tensor, dst_tensor: ms.Tensor
1117) -> None:
1118 """Synchronize the requires_grad flag from src_tensor to dst_tensor if they differ."""
1119 if src_tensor.requires_grad != dst_tensor.requires_grad:
1120 dst_tensor.requires_grad_(src_tensor.requires_grad)