Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / api.py: 79%
327 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-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"""hybrid shard data parallel interface"""
16from collections import namedtuple
17from typing import Any, List, Mapping, cast, Optional, Union
19from hyper_parallel.platform.platform import PlatformType
20from hyper_parallel.core.fully_shard.utils import MixedPrecisionPolicy, OffloadPolicy
21from hyper_parallel import DeviceMesh, init_device_mesh
22from hyper_parallel.platform import get_platform
23from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor
24from hyper_parallel.core.fully_shard.hsdp_utils import (
25 get_managed_modules_parameters,
26 is_dtensor_managed_param,
27 get_dtensor_managed_mesh,
28)
30platform = get_platform()
32origin_class_to_extend_class = {}
35def _resolve_comm_fusion_zero_copy_default(
36 platform_type: PlatformType,
37 comm_fusion: bool,
38 comm_fusion_zero_copy: Optional[bool],
39) -> bool:
40 """Resolve backend-specific default for the comm_fusion zero-copy path."""
41 if comm_fusion_zero_copy is not None:
42 return comm_fusion_zero_copy
43 if not comm_fusion:
44 return False
45 if platform_type == PlatformType.PYTORCH:
46 return True
47 if platform_type == PlatformType.MINDSPORE:
48 return False
49 return False
52def _check_strict_keys(
53 module: platform.Module, state_dict: Mapping[str, Any],
54) -> None:
55 """Raise ``RuntimeError`` if *state_dict* keys do not match *module*."""
56 expected_keys = set(module.state_dict().keys())
57 missing = expected_keys - set(state_dict.keys())
58 unexpected = set(state_dict.keys()) - expected_keys
59 error_msgs: list[str] = []
60 if missing:
61 error_msgs.append(
62 "Missing key(s): " + ", ".join(repr(k) for k in sorted(missing))
63 )
64 if unexpected:
65 error_msgs.append(
66 "Unexpected key(s): " + ", ".join(repr(k) for k in sorted(unexpected))
67 )
68 if error_msgs:
69 raise RuntimeError(
70 f"Error(s) in loading state_dict for "
71 f"{module.__class__.__name__}:\n\t"
72 + "\n\t".join(error_msgs)
73 )
76def _resolve_local_tensor(
77 key: str, val: platform.Tensor, target: DTensor,
78) -> platform.Tensor:
79 """Return the local shard tensor to be loaded into *target*."""
80 if isinstance(val, DTensor):
81 return val.to_local()
82 local_shape = tuple(target.local_shape)
83 global_shape = tuple(target.shape)
84 val_shape = tuple(val.shape)
85 if val_shape == local_shape:
86 return val
87 if val_shape == global_shape:
88 wrapped = distribute_tensor(
89 val, target.device_mesh,
90 target.layout.alias_placements if target.layout else target.placements,
91 )
92 return wrapped.to_local()
94 raise ValueError(
95 f"load '{key}': plain tensor shape {val_shape} "
96 f"matches neither local shard {local_shape} "
97 f"nor global {global_shape}."
98 )
101class _UnshardHandle:
102 """Unshard handle for user call HSDPModule.unshard(async_op=True)"""
103 def __init__(self, hsdp_state=None):
104 """
105 Initialize an async unshard handle.
107 Args:
108 hsdp_state (HSDPState, optional): The state to wait on. None means a no-op handle.
109 """
110 self._hsdp_state = hsdp_state
112 def wait(self):
113 """Block until the async unshard operation completes."""
114 if self._hsdp_state is not None:
115 self._hsdp_state.wait_for_unshard()
116 self._hsdp_state = None
119class HSDPModule:
120 """
121 The hsdp block of neural networks with hsdp interface.
123 Supported Platforms:
124 ``MindSpore`` ``torch``
125 """
127 def __init__(self):
128 """Initialize HSDPModule."""
129 self.hsdp_scheduler = None # Initialized in hsdp_init()
131 # pylint: disable=C0415
132 def hsdp_init(self, platform_type, module, mesh, reshard_after_forward,
133 shard_placement_fn, mp_policy, offload_policy, ignored_params, replicate_params, device,
134 comm_fusion, comm_fusion_zero_copy: Optional[bool] = None):
135 """init hsdp2 scheduler."""
136 scheduler_class = None
137 if platform_type == PlatformType.MINDSPORE:
138 from hyper_parallel.platform.mindspore.fully_shard.scheduler import MindSporeHSDPSchedulerV2
139 scheduler_class = MindSporeHSDPSchedulerV2
140 else:
141 from hyper_parallel.platform.torch.fully_shard.scheduler import TorchHSDPSchedulerV2
142 scheduler_class = TorchHSDPSchedulerV2
144 resolved_comm_fusion_zero_copy = _resolve_comm_fusion_zero_copy_default(
145 platform_type,
146 comm_fusion,
147 comm_fusion_zero_copy,
148 )
150 self.hsdp_scheduler = scheduler_class(module,
151 mesh,
152 reshard_after_forward,
153 shard_placement_fn,
154 mp_policy,
155 offload_policy,
156 ignored_params,
157 replicate_params,
158 device,
159 comm_fusion,
160 resolved_comm_fusion_zero_copy,
161 )
163 def set_requires_gradient_sync(self, requires_grad_sync):
164 r"""
165 set requires grad sync flag.
166 Args:
167 requires_grad_sync(bool): requires_grad_sync is used to control gradient sync process.
168 Raises:
169 ValueError: If `requires_grad_sync` is not bool.
170 """
171 if not isinstance(requires_grad_sync, bool):
172 raise ValueError(f"requires_grad_sync must be bool but got {requires_grad_sync}.")
173 if not hasattr(self, "hsdp_scheduler"):
174 raise ValueError("call hsdp interface first.")
176 for _, module in platform.get_cells_and_names(self):
177 if isinstance(module, HSDPModule):
178 module.hsdp_scheduler.set_requires_grad_sync(requires_grad_sync)
180 def zero_grad(self):
181 """zero accumunication grads"""
182 if not hasattr(self, "hsdp_scheduler"):
183 raise ValueError("call hsdp interface first.")
184 if platform.platform_type == PlatformType.PYTORCH:
185 return super().zero_grad()
186 for _, module in platform.get_cells_and_names(self):
187 if isinstance(module, HSDPModule):
188 module.hsdp_scheduler.zero_grad()
190 def set_modules_to_forward_prefetch(self, modules):
191 """set forward prefetch module list to prefetch all gather for unsharded parameters"""
192 if not isinstance(modules, (tuple, list)):
193 raise ValueError("modules must be HSDPModule list")
194 for module in modules:
195 if not isinstance(module, HSDPModule):
196 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
197 if not hasattr(self, "hsdp_scheduler"):
198 raise ValueError("call hsdp interface first.")
199 self.hsdp_scheduler.set_forward_prefetch_cells(modules)
201 def set_modules_to_backward_prefetch(self, modules):
202 """set backward prefetch module list to prefetch all gather for unsharded parameters"""
203 if not isinstance(modules, (tuple, list)):
204 raise ValueError("modules must be HSDPModule list")
205 for module in modules:
206 if not isinstance(module, HSDPModule):
207 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
208 if not hasattr(self, "hsdp_scheduler"):
209 raise ValueError("call fully_shard interface first.")
210 self.hsdp_scheduler.set_backward_prefetch_cells(modules)
212 def reshard(self) -> None:
213 """reshard all sharded parameters"""
214 if not self.hsdp_scheduler:
215 raise ValueError("hsdp_scheduler is None")
216 hsdp_state = self.hsdp_scheduler.hsdp_state
217 if hsdp_state:
218 hsdp_state.shard()
220 def unshard(self, async_op: bool = False):
221 """unshard all sharded parameters"""
222 if not isinstance(async_op, bool):
223 raise ValueError(f"async_op should be a bool, got {type(async_op)}")
224 if not self.hsdp_scheduler:
225 raise ValueError("hsdp_scheduler is None")
226 hsdp_state = self.hsdp_scheduler.hsdp_state
227 if hsdp_state:
228 hsdp_state.unshard(async_op) # pylint: disable=too-many-function-args
229 if async_op:
230 return _UnshardHandle(hsdp_state=hsdp_state)
231 return None
233 def load_state_dict(
234 self,
235 state_dict: Mapping[str, Any],
236 strict: bool = True,
237 assign: bool = False,
238 ):
239 """
240 Load state dict by copying directly into local shards.
242 Bypasses ``super().load_state_dict()`` because the standard PyTorch
243 implementation triggers ``copy_`` through the DTensor dispatcher, which
244 is not registered in the hyper-parallel layout system.
246 Each value in ``state_dict`` is dispatched by type:
247 - hyper DTensor: extract local shard and copy directly.
248 - plain Tensor whose shape == local shard shape: copy as-is.
249 - plain Tensor whose shape == global shape: distribute via
250 ``distribute_tensor``, then copy the local shard.
252 Args:
253 state_dict (Mapping[str, Any]): Fully-qualified parameter/buffer
254 names mapped to tensors (DTensor or plain Tensor).
255 strict (bool): If ``True`` (default), missing or unexpected keys
256 raise ``RuntimeError``, matching ``nn.Module.load_state_dict``
257 semantics.
258 assign (bool): When ``True`` *and* every value in ``state_dict`` is
259 already a hyper DTensor, defer to the standard
260 ``nn.Module.load_state_dict(assign=True)``, which replaces the
261 module's parameters/buffers with the given DTensors instead of
262 copying into existing storage. This is required when loading
263 sharded DTensors onto a meta-device model (e.g.
264 ``cpu_ram_efficient_loading``). If ``state_dict`` contains any
265 plain tensor (local-shard or global shape), ``assign`` is
266 ignored and the copy/distribute path below is used so the
267 target stays a properly sharded DTensor.
269 Raises:
270 RuntimeError: When ``strict`` is ``True`` and keys do not match.
271 ValueError: When a plain tensor shape matches neither the local
272 shard shape nor the global shape of the target DTensor.
273 """
274 if assign and state_dict and all(
275 isinstance(val, DTensor) for val in state_dict.values()
276 ):
277 return super().load_state_dict(state_dict, strict=strict, assign=True)
278 self_module = cast(platform.Module, self)
280 target_map: dict[str, platform.Tensor] = {}
281 for name, p in platform.parameters_dict(self_module):
282 target_map[name] = p
283 for name, b in self_module.named_buffers():
284 target_map[name] = b
286 if strict:
287 _check_strict_keys(self_module, state_dict)
289 with platform.no_grad():
290 for key, val in state_dict.items():
291 target = target_map.get(key)
292 if target is None:
293 continue
295 if isinstance(target, DTensor):
296 val = _resolve_local_tensor(key, val, target)
297 platform.load_into_param(target, val)
299 # Trigger load_state_dict post-hooks so that HSDP internal
300 # bookkeeping (e.g. _sharded_param_data) stays in sync.
301 # Pass an IncompatibleKeys with the same attribute names as PyTorch
302 # so external hooks can safely read .missing_keys/.unexpected_keys.
303 _IK = namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"])
304 incompatible_keys = _IK([], [])
305 for _, module in platform.get_cells_and_names(self_module):
306 hooks = module._load_state_dict_post_hooks # pylint: disable=protected-access
307 for hook in hooks.values():
308 hook(module, incompatible_keys)
310 def set_is_last_backward(self, is_last_backward: bool):
311 """set is_last_backward flag"""
312 self.hsdp_scheduler.scheduler_ctx.is_last_backward = is_last_backward
314 def set_requires_all_reduce(self, requires_all_reduce: bool, *, recurse: bool = True) -> None:
315 """set requires_all_reduce flag"""
316 if not isinstance(requires_all_reduce, bool):
317 raise ValueError(
318 f"requires_all_reduce should be a bool, got {type(requires_all_reduce)}"
319 )
320 if not recurse:
321 raise NotImplementedError(
322 "Currently impl is equal to recurse=True, "
323 "need support module_param mapping."
324 )
325 self_module = cast(platform.Module, self)
326 for _, module in platform.get_cells_and_names(self_module):
327 if isinstance(module, HSDPModule):
328 module.hsdp_scheduler.set_requires_all_reduce(requires_all_reduce)
330 def set_reshard_after_forward(self, reshard_after_forward: bool, recurse: bool = True) -> None:
331 """set reshard_after_forward flag"""
332 if not isinstance(reshard_after_forward, bool):
333 raise ValueError(
334 f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}"
335 )
336 if not recurse:
337 raise NotImplementedError(
338 "Currently impl is equal to recurse=True, "
339 "need support module_param mapping."
340 )
341 self_module = cast(platform.Module, self)
342 for _, module in platform.get_cells_and_names(self_module):
343 if isinstance(module, HSDPModule):
344 module.hsdp_scheduler.set_reshard_after_forward(reshard_after_forward)
346 def set_reshard_after_backward(self, reshard_after_backward: bool, recurse: bool = True) -> None:
347 """set reshard_after_backward flag"""
348 if not isinstance(reshard_after_backward, bool):
349 raise ValueError(
350 f"reshard_after_backward should be a bool, got {type(reshard_after_backward)}"
351 )
352 if not recurse:
353 raise NotImplementedError(
354 "Currently impl is equal to recurse=True, "
355 "need support module_param mapping."
356 )
357 self_module = cast(platform.Module, self)
358 for _, module in platform.get_cells_and_names(self_module):
359 if isinstance(module, HSDPModule):
360 module.hsdp_scheduler.set_reshard_after_backward(reshard_after_backward)
362 def set_reduce_op_type(self, reduce_op_type, recurse: bool = True) -> None:
363 """
364 set reduce_op_type for all reduce operations in HSDP
365 support reduce_op_type "avg" and "sum", default is "avg"
366 """
367 self_module = cast(platform.Module, self)
368 if recurse:
369 sub_modules = [m for _, m in platform.get_cells_and_names(self_module)]
370 else:
371 sub_modules = [self_module]
372 for module in sub_modules:
373 if isinstance(module, HSDPModule):
374 hsdp_state = module.hsdp_scheduler.hsdp_state
375 if hsdp_state:
376 hsdp_state.set_reduce_op_type(reduce_op_type)
378 def set_gradient_scaling_factor(self, factor=None):
379 """
380 Set a multiplicative scaling factor applied to gradients after
381 reduce-scatter / all-reduce and before they are written into
382 ``sharded_param.grad``.
384 ``factor`` may be ``None`` (disable scaling), a Python ``float``/``int``,
385 or a 0-dim/1-element tensor. Setting ``factor`` to ``None`` (the default
386 on construction) skips the scaling op entirely so no extra device-side
387 ``mul_`` is launched on the hot path.
389 Args:
390 factor (None | float | int | platform.Tensor): Scaling coefficient.
391 Use ``None`` to disable scaling.
393 Raises:
394 ValueError: If ``factor`` is not one of the supported types or is a
395 tensor with more than one element.
396 """
397 if factor is not None:
398 if isinstance(factor, bool):
399 raise ValueError(
400 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
401 f"but got bool {factor}."
402 )
403 if isinstance(factor, platform.Tensor):
404 if factor.numel() != 1:
405 raise ValueError(
406 f"gradient_scaling_factor tensor must have exactly 1 element, "
407 f"but got shape {tuple(factor.shape)}."
408 )
409 elif not isinstance(factor, (float, int)):
410 raise ValueError(
411 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
412 f"but got {type(factor).__name__}."
413 )
414 hsdp_state = self.hsdp_scheduler.hsdp_state
415 if hsdp_state:
416 hsdp_state.set_gradient_scaling_factor(factor)
419def _extend_module_with_hsdp_interface(module):
420 """Dynamically extend module's class to inherit from HSDPModule, adding HSDP capabilities."""
421 origin_class = module.__class__
422 extend_class = origin_class_to_extend_class.get(origin_class, None)
423 if extend_class is None:
424 extend_class = type(f"HSDP{origin_class.__name__}", (HSDPModule, origin_class), {})
425 origin_class_to_extend_class[origin_class] = extend_class
426 module.__class__ = extend_class
429def _get_root_modules(modules: List[platform.Module]) -> List[platform.Module]:
430 """
431 Returns the modules in ``modules`` that are root modules (i.e. parent-less)
432 with respect to the set ``modules``. In other words, these are the modules
433 in ``modules`` that are not the child of any other module in ``modules``.
435 Aligned with PyTorch torch.distributed.utils._get_root_modules.
436 """
437 root_modules: List[platform.Module] = []
439 def _get_submodules(mod):
440 if platform.platform_type == PlatformType.MINDSPORE:
441 return set(c for _, c in mod.cells_and_names())
442 return set(mod.modules())
444 module_to_modules: dict[platform.Module, set] = {
445 m: _get_submodules(m) for m in modules
446 }
447 for candidate in modules:
448 is_root = True
449 for mod, submodules in module_to_modules.items():
450 if candidate is not mod and candidate in submodules:
451 is_root = False
452 break
453 if is_root:
454 root_modules.append(candidate)
455 return root_modules
458def _check_module_valid(platform_type, module):
459 """check module valid"""
460 if platform_type == PlatformType.MINDSPORE:
461 from mindspore.nn.cell import Cell
462 if not isinstance(module, Cell):
463 raise ValueError(f"module's type must be nn.cell but got {type(module)}.")
464 else:
465 from torch.nn import Module
466 if not isinstance(module, Module):
467 raise ValueError(f"module's type must be nn.Module but got {type(module)}.")
470def _validate_module_for_fully_shard(
471 module: Union[platform.Module, List[platform.Module]], platform_type
472) -> None:
473 """Validate module(s) for fully_shard. Platform-aware for single module."""
474 if isinstance(module, list):
475 if len(module) == 0:
476 raise ValueError("fully_shard does not support empty list of modules.")
477 for i, m in enumerate(module):
478 try:
479 _check_module_valid(platform_type, m)
480 except ValueError:
481 raise ValueError(
482 f"fully_shard expects nn.Module or list[nn.Module], "
483 f"but got list with {type(m).__name__} at index {i}."
484 ) from None
485 else:
486 _check_module_valid(platform_type, module)
489HsdpValidationOptions = namedtuple(
490 "HsdpValidationOptions",
491 [
492 "shard_size",
493 "threshold",
494 "optimizer_level",
495 "enable_grad_accumulation",
496 "grad_scale",
497 "reduce_dtype",
498 "comm_async",
499 "comm_fusion",
500 "bucket_size",
501 ],
502)
505def _validate_hsdp_shard_size(shard_size: int) -> None:
506 if not isinstance(shard_size, int) or (shard_size <= 0 and shard_size != -1):
507 raise ValueError(f"shard_size must be a positive integer, but got {shard_size}.")
510def _validate_hsdp_threshold(threshold: int) -> None:
511 if not isinstance(threshold, int) or threshold < 0:
512 raise ValueError(f"threshold must be a positive integer or 0, but got {threshold}.")
515def _validate_hsdp_optimizer_level(optimizer_level: str) -> None:
516 if optimizer_level not in ["level1", "level2", "level3"]:
517 raise ValueError(
518 f"Optimizer level should in ['level1', 'level2', 'level3'], but got {optimizer_level}."
519 )
522def _validate_hsdp_reduce_dtype(platform_type: PlatformType, reduce_dtype) -> None:
523 if platform_type == PlatformType.MINDSPORE:
524 from mindspore._c_expression.typing import Type
525 if reduce_dtype is not None and not isinstance(reduce_dtype, Type):
526 raise ValueError(f"reduce_dtype must be mindspore.dtype but got {reduce_dtype}.")
527 return
528 import torch
529 if reduce_dtype is not None and not isinstance(reduce_dtype, torch.dtype):
530 raise ValueError(f"reduce_dtype must be torch.dtype but got {reduce_dtype}.")
533def _check_hsdp_input_valid(platform_type, module, options: HsdpValidationOptions):
534 """check hsdp input valid"""
535 _check_module_valid(platform_type, module)
536 _validate_hsdp_shard_size(options.shard_size)
537 _validate_hsdp_threshold(options.threshold)
538 _validate_hsdp_optimizer_level(options.optimizer_level)
539 if not isinstance(options.enable_grad_accumulation, bool):
540 raise ValueError(
541 f"enable_grad_accumulation must be bool but got {options.enable_grad_accumulation}."
542 )
543 if not isinstance(options.grad_scale, float):
544 raise ValueError(f"grad_scale must be float but got {options.grad_scale}.")
545 _validate_hsdp_reduce_dtype(platform_type, options.reduce_dtype)
546 if not isinstance(options.comm_async, bool):
547 raise ValueError(f"comm_async must be bool but got {options.comm_async}.")
548 if not isinstance(options.comm_fusion, bool):
549 raise ValueError(f"comm_fusion must be bool but got {options.comm_fusion}.")
550 if not isinstance(options.bucket_size, int) or (
551 options.bucket_size < 0 and options.bucket_size != -1
552 ):
553 raise ValueError(
554 f"bucket_size must be a positive integer or 0, but got {options.bucket_size}."
555 )
558def _get_device_from_mesh(mesh: DeviceMesh):
559 """Extract and validate the torch device from the device mesh."""
560 device = None
561 device_type = mesh.device_type
562 if device_type not in ("npu", "cuda"):
563 raise AssertionError(
564 f"hyper_parallel.fully_shard support device in [torch.npu, torch.cuda], "
565 f"but got '{device_type}'"
566 )
567 if platform.platform_type == PlatformType.PYTORCH:
568 device_handle = platform.get_device_handle(device_type)
569 if device_handle is None:
570 raise ValueError(
571 f"hyper_parallel.fully_shard can't find device_handle of "
572 f"'torch.{device_type}', check the environment."
573 )
574 if device_handle.is_available():
575 import torch
576 device = torch.device(device_handle.current_device())
577 else:
578 device = device_type
579 return device
582def _normalize_replicate_params(
583 replicate_params: Optional[set[platform.Parameter]],
584) -> set[platform.Parameter]:
585 """
586 Normalize replicate_params for fully_shard
587 Args:
588 replicate_params (Optional[set[nn.Parameter]]): Set of parameters to exclude from sharding.
589 Returns:
590 set[nn.Parameter]: Set of parameters to exclude from sharding.
591 """
592 if replicate_params is None:
593 return set()
594 out = set(replicate_params)
595 for p in out:
596 if not isinstance(p, (platform.Parameter, DTensor)):
597 raise TypeError(
598 "replicate_params must contain only nn.Parameter or DTensor, "
599 f"got {type(p).__name__}."
600 )
601 return out
604def _get_modules_parameters(modules, ignored_params=None):
605 """Collect deduplicated parameters from module roots."""
606 return get_managed_modules_parameters(modules, ignored_params)
609def fully_shard(
610 module: Union[platform.Module, List[platform.Module]],
611 *,
612 mesh: Optional[DeviceMesh] = None,
613 reshard_after_forward: bool = True,
614 shard_placement_fn: None = None,
615 mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
616 offload_policy: OffloadPolicy = OffloadPolicy(),
617 ignored_params: Optional[set[platform.Parameter]] = None,
618 replicate_params: Optional[set[platform.Parameter]] = None,
619 comm_fusion: bool = False,
620 comm_fusion_zero_copy: Optional[bool] = None,
621) -> Union[platform.Module, List[platform.Module]]:
623 """
624 Apply fully_shard to a module (or list of modules) for distributed training with parameter sharding.
626 This interface provides PyTorch-compatible HSDP (Hybrid Sharded Data Parallelism)
627 functionality, enabling efficient training of large models by sharding parameters
628 across multiple devices. The module is automatically enhanced with distributed
629 capabilities including parameter sharding, gradient synchronization, and memory
630 management.
632 When a list of modules is passed, they are treated as one FSDP unit (parameters
633 grouped together). Both PyTorch and MindSpore platforms support list input.
635 Parameters:
636 module (nn.Module or List[nn.Module]):
637 The module(s) to apply fully_shard to. Modified in-place. When a list
638 is passed, parameters from all modules are grouped as one FSDP unit.
640 mesh (Optional[DeviceMesh], default=None):
641 The device mesh defining the process topology for distributed training.
642 If None, fully_shard keeps pure-DTensor modules on their original
643 distributed layout and only creates a default 1D mesh when local
644 parameters need explicit data-parallel/FSDP management.
646 reshard_after_forward (bool, default=True):
647 Whether to automatically reshard parameters after forward. When True,
648 parameters are resharded immediately after they are no longer needed,
649 freeing memory for subsequent operations. Set to False if you want to
650 keep parameters unsharded for backward pass or manual control.
652 shard_placement_fn (Callable, default=None):
653 A callable that determines how to shard each parameter. The function
654 should accept a parameter and return a Shard object specifying the
655 sharding dimension, or None to use default sharding (dimension 0)
657 mp_policy (MixedPrecisionPolicy, default=MixedPrecisionPolicy()):
658 Mixed precision training policy controlling data type conversions.
659 offload_policy (OffloadPolicy, default=OffloadPolicy()):
660 Memory offload policy for reducing device memory usage.
662 ignored_params (Optional[set[nn.Parameter]], default=None):
663 Set of parameters to exclude from fully_shard management entirely.
664 These parameters are left on the original module as regular parameters,
665 are not sharded, and do not participate in fully_shard gradient
666 synchronization. Use this for parameters that should remain outside
667 the fully_shard lifecycle.
669 comm_fusion (bool, default=False):
670 Whether enable all_gather fusion and reduce_scatter fusion.
672 replicate_params (Optional[set[nn.Parameter]], default=None):
673 Set of parameters to keep replicated while still managing them under
674 fully_shard. These parameters are not sharded, but their gradients
675 are still synchronized with DDP-style all-reduce over the current
676 fully_shard communication domain. This differs from ``ignored_params``,
677 which skips fully_shard management and gradient synchronization
678 entirely for the selected parameters.
680 comm_fusion_zero_copy (Optional[bool], default=None):
681 Whether allow the experimental zero-copy path for
682 ``comm_fusion``. When set to ``None``, fully_shard uses a backend-specific
683 default:
684 - PyTorch: enabled automatically when ``comm_fusion=True``
685 - MindSpore: disabled automatically even when ``comm_fusion=True``
686 When enabled, fully_shard may rebase sharded local parameter storage
687 into one shared flat buffer so fused all-gather can read directly from
688 contiguous memory. This path depends on optimizer compatibility with
689 view-backed parameters.
691 Returns:
692 nn.Module or List[nn.Module]: The input module(s) with HSDP capabilities added.
693 """
694 platform_type = platform.platform_type
695 _validate_module_for_fully_shard(module, platform_type)
696 if platform_type == PlatformType.MINDSPORE:
697 from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat
699 enable_mindspore_backward_compat()
701 arg_module = module
702 if isinstance(module, list):
703 modules = tuple(_get_root_modules(module))
704 else:
705 modules = (module,)
707 for mod in modules:
708 _extend_module_with_hsdp_interface(mod)
710 params = _get_modules_parameters(modules, ignored_params)
711 has_dtensor_param = any(is_dtensor_managed_param(param) for param in params)
712 replicate_params = _normalize_replicate_params(replicate_params)
714 if mesh is None and not has_dtensor_param:
715 mesh = init_device_mesh(device_type="npu", mesh_shape=(platform.get_world_size(),))
716 if mesh is not None:
717 device = _get_device_from_mesh(mesh)
718 else:
719 compat_mesh = None
720 for param in params:
721 dtensor_mesh = get_dtensor_managed_mesh(param)
722 if dtensor_mesh is not None:
723 compat_mesh = dtensor_mesh
724 break
725 if compat_mesh is None:
726 raise ValueError("fully_shard could not resolve a DTensor mesh for compatibility mode.")
727 device = _get_device_from_mesh(compat_mesh)
729 init_modules = modules
730 modules[0].hsdp_init(
731 platform_type,
732 init_modules,
733 mesh,
734 reshard_after_forward,
735 shard_placement_fn,
736 mp_policy,
737 offload_policy,
738 ignored_params,
739 replicate_params,
740 device,
741 comm_fusion,
742 comm_fusion_zero_copy,
743 )
744 # Share the same scheduler handle with other roots so mods[i].unshard()/prefetch work
745 if len(modules) > 1:
746 for mod in modules[1:]:
747 mod.hsdp_scheduler = modules[0].hsdp_scheduler
748 return arg_module
751def get_model_state_dict(model, *, options=None):
752 """Get model state dict with platform-specific implementation.
754 Delegates to the platform-specific implementation at runtime.
755 Users import from here instead of platform internals.
756 """
757 return platform.get_model_state_dict(model, options=options)
760def hsdp_sync_stream():
761 """Wait for hsdp gradient handle to be completed."""
762 platform.wait_grad_handle()