Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / api.py: 82%
356 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-25 04:27 +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, SourceShardMetaInfo
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)
28platform = get_platform()
29ModuleClass = platform.Module
30TensorClass = platform.Tensor
31ParameterClass = platform.Parameter
33origin_class_to_extend_class = {}
36def _resolve_comm_fusion_zero_copy_default(
37 platform_type: PlatformType,
38 comm_fusion: bool,
39 comm_fusion_zero_copy: Optional[bool],
40) -> bool:
41 """Resolve backend-specific default for the comm_fusion zero-copy path."""
42 if comm_fusion_zero_copy is not None:
43 return comm_fusion_zero_copy
44 if not comm_fusion:
45 return False
46 if platform_type == PlatformType.PYTORCH:
47 return True
48 if platform_type == PlatformType.MINDSPORE:
49 return False
50 return False
53def _check_strict_keys(
54 module: ModuleClass, state_dict: Mapping[str, Any],
55) -> None:
56 """Raise ``RuntimeError`` if *state_dict* keys do not match *module*."""
57 expected_keys = set(module.state_dict().keys())
58 missing = expected_keys - set(state_dict.keys())
59 unexpected = set(state_dict.keys()) - expected_keys
60 error_msgs: list[str] = []
61 if missing:
62 error_msgs.append(
63 "Missing key(s): " + ", ".join(repr(k) for k in sorted(missing))
64 )
65 if unexpected:
66 error_msgs.append(
67 "Unexpected key(s): " + ", ".join(repr(k) for k in sorted(unexpected))
68 )
69 if error_msgs:
70 raise RuntimeError(
71 f"Error(s) in loading state_dict for "
72 f"{module.__class__.__name__}:\n\t"
73 + "\n\t".join(error_msgs)
74 )
77def _resolve_local_tensor(
78 key: str, val: TensorClass, target: DTensor,
79) -> TensorClass:
80 """Return the local shard tensor to be loaded into *target*."""
81 if isinstance(val, DTensor):
82 return val.to_local()
83 local_shape = tuple(target.local_shape)
84 global_shape = tuple(target.shape)
85 val_shape = tuple(val.shape)
86 if val_shape == local_shape:
87 return val
88 if val_shape == global_shape:
89 wrapped = distribute_tensor(
90 val, target.device_mesh,
91 target.layout.alias_placements if target.layout else target.placements,
92 )
93 return wrapped.to_local()
95 raise ValueError(
96 f"load '{key}': plain tensor shape {val_shape} "
97 f"matches neither local shard {local_shape} "
98 f"nor global {global_shape}."
99 )
102class _UnshardHandle:
103 """Unshard handle for user call HSDPModule.unshard(async_op=True)"""
104 def __init__(self, hsdp_state=None):
105 """
106 Initialize an async unshard handle.
108 Args:
109 hsdp_state (HSDPState, optional): The state to wait on. None means a no-op handle.
110 """
111 self._hsdp_state = hsdp_state
113 def wait(self):
114 """Block until the async unshard operation completes."""
115 if self._hsdp_state is not None:
116 self._hsdp_state.wait_for_unshard()
117 self._hsdp_state = None
120class HSDPModule:
121 """
122 The hsdp block of neural networks with hsdp interface.
124 Supported Platforms:
125 ``MindSpore`` ``torch``
126 """
128 def __init__(self):
129 """Initialize HSDPModule."""
130 self.hsdp_scheduler = None # Initialized in hsdp_init()
132 # pylint: disable=C0415
133 def hsdp_init(self, platform_type, module, mesh, reshard_after_forward,
134 shard_placement_fn, mp_policy, offload_policy, ignored_params, replicate_params, device,
135 comm_fusion, comm_fusion_zero_copy: Optional[bool] = None,
136 source_shard_infos: Optional[Mapping[ParameterClass, SourceShardMetaInfo]] = None):
137 """init hsdp2 scheduler."""
138 scheduler_class = None
139 if platform_type == PlatformType.MINDSPORE:
140 from hyper_parallel.platform.mindspore.fully_shard.scheduler import MindSporeHSDPSchedulerV2
141 scheduler_class = MindSporeHSDPSchedulerV2
142 else:
143 from hyper_parallel.platform.torch.fully_shard.scheduler import TorchHSDPSchedulerV2
144 scheduler_class = TorchHSDPSchedulerV2
146 resolved_comm_fusion_zero_copy = _resolve_comm_fusion_zero_copy_default(
147 platform_type,
148 comm_fusion,
149 comm_fusion_zero_copy,
150 )
152 self.hsdp_scheduler = scheduler_class(module,
153 mesh,
154 reshard_after_forward,
155 shard_placement_fn,
156 mp_policy,
157 offload_policy,
158 ignored_params,
159 replicate_params,
160 device,
161 comm_fusion,
162 resolved_comm_fusion_zero_copy,
163 source_shard_infos=source_shard_infos,
164 )
166 def set_requires_gradient_sync(self, requires_grad_sync):
167 r"""
168 set requires grad sync flag.
169 Args:
170 requires_grad_sync(bool): requires_grad_sync is used to control gradient sync process.
171 Raises:
172 ValueError: If `requires_grad_sync` is not bool.
173 """
174 if not isinstance(requires_grad_sync, bool):
175 raise ValueError(f"requires_grad_sync must be bool but got {requires_grad_sync}.")
176 if not hasattr(self, "hsdp_scheduler"):
177 raise ValueError("call hsdp interface first.")
179 for _, module in platform.get_cells_and_names(self):
180 if isinstance(module, HSDPModule):
181 module.hsdp_scheduler.set_requires_grad_sync(requires_grad_sync)
183 def zero_grad(self):
184 """zero accumunication grads"""
185 if not hasattr(self, "hsdp_scheduler"):
186 raise ValueError("call hsdp interface first.")
187 if platform.platform_type == PlatformType.PYTORCH:
188 return super().zero_grad()
189 for _, module in platform.get_cells_and_names(self):
190 if isinstance(module, HSDPModule):
191 module.hsdp_scheduler.zero_grad()
193 def set_modules_to_forward_prefetch(self, modules):
194 """set forward prefetch module list to prefetch all gather for unsharded parameters"""
195 if not isinstance(modules, (tuple, list)):
196 raise ValueError("modules must be HSDPModule list")
197 for module in modules:
198 if not isinstance(module, HSDPModule):
199 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
200 if not hasattr(self, "hsdp_scheduler"):
201 raise ValueError("call hsdp interface first.")
202 self.hsdp_scheduler.set_forward_prefetch_cells(modules)
204 def set_modules_to_backward_prefetch(self, modules):
205 """set backward prefetch module list to prefetch all gather for unsharded parameters"""
206 if not isinstance(modules, (tuple, list)):
207 raise ValueError("modules must be HSDPModule list")
208 for module in modules:
209 if not isinstance(module, HSDPModule):
210 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
211 if not hasattr(self, "hsdp_scheduler"):
212 raise ValueError("call fully_shard interface first.")
213 self.hsdp_scheduler.set_backward_prefetch_cells(modules)
215 def reshard(self) -> None:
216 """reshard all sharded parameters"""
217 if not self.hsdp_scheduler:
218 raise ValueError("hsdp_scheduler is None")
219 hsdp_state = self.hsdp_scheduler.hsdp_state
220 if hsdp_state:
221 hsdp_state.shard()
223 def unshard(self, async_op: bool = False):
224 """unshard all sharded parameters"""
225 if not isinstance(async_op, bool):
226 raise ValueError(f"async_op should be a bool, got {type(async_op)}")
227 if not self.hsdp_scheduler:
228 raise ValueError("hsdp_scheduler is None")
229 hsdp_state = self.hsdp_scheduler.hsdp_state
230 if hsdp_state:
231 hsdp_state.unshard(async_op) # pylint: disable=too-many-function-args
232 if async_op:
233 return _UnshardHandle(hsdp_state=hsdp_state)
234 return None
236 def load_state_dict(
237 self,
238 state_dict: Mapping[str, Any],
239 strict: bool = True,
240 assign: bool = False,
241 ):
242 """
243 Load state dict by copying directly into local shards.
245 Bypasses ``super().load_state_dict()`` because the standard PyTorch
246 implementation triggers ``copy_`` through the DTensor dispatcher, which
247 is not registered in the hyper-parallel layout system.
249 Each value in ``state_dict`` is dispatched by type:
250 - hyper DTensor: extract local shard and copy directly.
251 - plain Tensor whose shape == local shard shape: copy as-is.
252 - plain Tensor whose shape == global shape: distribute via
253 ``distribute_tensor``, then copy the local shard.
255 Args:
256 state_dict (Mapping[str, Any]): Fully-qualified parameter/buffer
257 names mapped to tensors (DTensor or plain Tensor).
258 strict (bool): If ``True`` (default), missing or unexpected keys
259 raise ``RuntimeError``, matching ``nn.Module.load_state_dict``
260 semantics.
261 assign (bool): When ``True`` *and* every value in ``state_dict`` is
262 already a hyper DTensor, defer to the standard
263 ``nn.Module.load_state_dict(assign=True)``, which replaces the
264 module's parameters/buffers with the given DTensors instead of
265 copying into existing storage. This is required when loading
266 sharded DTensors onto a meta-device model (e.g.
267 ``cpu_ram_efficient_loading``). If ``state_dict`` contains any
268 plain tensor (local-shard or global shape), ``assign`` is
269 ignored and the copy/distribute path below is used so the
270 target stays a properly sharded DTensor.
272 Raises:
273 RuntimeError: When ``strict`` is ``True`` and keys do not match.
274 ValueError: When a plain tensor shape matches neither the local
275 shard shape nor the global shape of the target DTensor.
276 """
277 if assign and state_dict and all(
278 isinstance(val, DTensor) for val in state_dict.values()
279 ):
280 return super().load_state_dict(state_dict, strict=strict, assign=True)
281 self_module = cast(ModuleClass, self)
283 target_map: dict[str, TensorClass] = {}
284 for name, p in platform.parameters_dict(self_module):
285 target_map[name] = p
286 for name, b in self_module.named_buffers():
287 target_map[name] = b
289 if strict:
290 _check_strict_keys(self_module, state_dict)
292 with platform.no_grad():
293 for key, val in state_dict.items():
294 target = target_map.get(key)
295 if target is None:
296 continue
298 if isinstance(target, DTensor):
299 val = _resolve_local_tensor(key, val, target)
300 platform.load_into_param(target, val)
302 # Trigger load_state_dict post-hooks so that HSDP internal
303 # bookkeeping (e.g. _sharded_param_data) stays in sync.
304 # Pass an IncompatibleKeys with the same attribute names as PyTorch
305 # so external hooks can safely read .missing_keys/.unexpected_keys.
306 _IK = namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"])
307 incompatible_keys = _IK([], [])
308 for _, module in platform.get_cells_and_names(self_module):
309 hooks = module._load_state_dict_post_hooks # pylint: disable=protected-access
310 for hook in hooks.values():
311 hook(module, incompatible_keys)
313 def set_is_last_backward(self, is_last_backward: bool):
314 """set is_last_backward flag"""
315 self.hsdp_scheduler.scheduler_ctx.is_last_backward = is_last_backward
317 def reset_iter_state(self, recursive: bool = True) -> None:
318 """Reset fully_shard iteration bookkeeping without clearing optimizer gradients.
320 This method must be called only after communication for the iteration has
321 completed. It does not wait in-flight collectives or clear
322 ``sharded_param.grad``/``main_grad``.
324 Args:
325 recursive: Whether to reset every HSDPModule in this module tree.
327 Raises:
328 ValueError: If ``recursive`` is not a bool.
329 """
330 if not isinstance(recursive, bool):
331 raise ValueError(f"recursive should be a bool, got {type(recursive)}")
332 self_module = cast(ModuleClass, self)
333 modules = (
334 [module for _, module in platform.get_cells_and_names(self_module)]
335 if recursive
336 else [self_module]
337 )
338 for module in modules:
339 if isinstance(module, HSDPModule):
340 module.hsdp_scheduler.reset_iter_state()
342 def set_requires_all_reduce(self, requires_all_reduce: bool, *, recurse: bool = True) -> None:
343 """set requires_all_reduce flag"""
344 if not isinstance(requires_all_reduce, bool):
345 raise ValueError(
346 f"requires_all_reduce should be a bool, got {type(requires_all_reduce)}"
347 )
348 if not recurse:
349 raise NotImplementedError(
350 "Currently impl is equal to recurse=True, "
351 "need support module_param mapping."
352 )
353 self_module = cast(ModuleClass, self)
354 for _, module in platform.get_cells_and_names(self_module):
355 if isinstance(module, HSDPModule):
356 module.hsdp_scheduler.set_requires_all_reduce(requires_all_reduce)
358 def set_reshard_after_forward(self, reshard_after_forward: bool, recurse: bool = True) -> None:
359 """set reshard_after_forward flag"""
360 if not isinstance(reshard_after_forward, bool):
361 raise ValueError(
362 f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}"
363 )
364 if not recurse:
365 raise NotImplementedError(
366 "Currently impl is equal to recurse=True, "
367 "need support module_param mapping."
368 )
369 self_module = cast(ModuleClass, self)
370 for _, module in platform.get_cells_and_names(self_module):
371 if isinstance(module, HSDPModule):
372 module.hsdp_scheduler.set_reshard_after_forward(reshard_after_forward)
374 def set_reshard_after_backward(self, reshard_after_backward: bool, recurse: bool = True) -> None:
375 """set reshard_after_backward flag"""
376 if not isinstance(reshard_after_backward, bool):
377 raise ValueError(
378 f"reshard_after_backward should be a bool, got {type(reshard_after_backward)}"
379 )
380 if not recurse:
381 raise NotImplementedError(
382 "Currently impl is equal to recurse=True, "
383 "need support module_param mapping."
384 )
385 self_module = cast(ModuleClass, self)
386 for _, module in platform.get_cells_and_names(self_module):
387 if isinstance(module, HSDPModule):
388 module.hsdp_scheduler.set_reshard_after_backward(reshard_after_backward)
390 def set_reduce_op_type(self, reduce_op_type, recurse: bool = True) -> None:
391 """
392 set reduce_op_type for all reduce operations in HSDP
393 support reduce_op_type "avg" and "sum", default is "avg"
394 """
395 self_module = cast(ModuleClass, self)
396 if recurse:
397 sub_modules = [m for _, m in platform.get_cells_and_names(self_module)]
398 else:
399 sub_modules = [self_module]
400 for module in sub_modules:
401 if isinstance(module, HSDPModule):
402 hsdp_state = module.hsdp_scheduler.hsdp_state
403 if hsdp_state:
404 hsdp_state.set_reduce_op_type(reduce_op_type)
406 def set_gradient_scaling_factor(self, factor=None):
407 """
408 Set a multiplicative scaling factor applied to gradients after
409 reduce-scatter / all-reduce and before they are written into
410 ``sharded_param.grad``.
412 ``factor`` may be ``None`` (disable scaling), a Python ``float``/``int``,
413 or a 0-dim/1-element tensor. Setting ``factor`` to ``None`` (the default
414 on construction) skips the scaling op entirely so no extra device-side
415 ``mul_`` is launched on the hot path.
417 Args:
418 factor (None | float | int | platform.Tensor): Scaling coefficient.
419 Use ``None`` to disable scaling.
421 Raises:
422 ValueError: If ``factor`` is not one of the supported types or is a
423 tensor with more than one element.
424 """
425 if factor is not None:
426 if isinstance(factor, bool):
427 raise ValueError(
428 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
429 f"but got bool {factor}."
430 )
431 if isinstance(factor, platform.Tensor):
432 if factor.numel() != 1:
433 raise ValueError(
434 f"gradient_scaling_factor tensor must have exactly 1 element, "
435 f"but got shape {tuple(factor.shape)}."
436 )
437 elif not isinstance(factor, (float, int)):
438 raise ValueError(
439 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
440 f"but got {type(factor).__name__}."
441 )
442 hsdp_state = self.hsdp_scheduler.hsdp_state
443 if hsdp_state:
444 hsdp_state.set_gradient_scaling_factor(factor)
447def _extend_module_with_hsdp_interface(module):
448 """Dynamically extend module's class to inherit from HSDPModule, adding HSDP capabilities."""
449 origin_class = module.__class__
450 extend_class = origin_class_to_extend_class.get(origin_class, None)
451 if extend_class is None:
452 extend_class = type(f"HSDP{origin_class.__name__}", (HSDPModule, origin_class), {})
453 origin_class_to_extend_class[origin_class] = extend_class
454 module.__class__ = extend_class
457def _get_root_modules(modules: List[ModuleClass]) -> List[ModuleClass]:
458 """
459 Returns the modules in ``modules`` that are root modules (i.e. parent-less)
460 with respect to the set ``modules``. In other words, these are the modules
461 in ``modules`` that are not the child of any other module in ``modules``.
463 Aligned with PyTorch torch.distributed.utils._get_root_modules.
464 """
465 root_modules: List[ModuleClass] = []
467 def _get_submodules(mod):
468 if platform.platform_type == PlatformType.MINDSPORE:
469 return set(c for _, c in mod.cells_and_names())
470 return set(mod.modules())
472 module_to_modules: dict[ModuleClass, set] = {
473 m: _get_submodules(m) for m in modules
474 }
475 for candidate in modules:
476 is_root = True
477 for mod, submodules in module_to_modules.items():
478 if candidate is not mod and candidate in submodules:
479 is_root = False
480 break
481 if is_root:
482 root_modules.append(candidate)
483 return root_modules
486def _check_module_valid(platform_type, module):
487 """check module valid"""
488 if platform_type == PlatformType.MINDSPORE:
489 from mindspore.nn.cell import Cell
490 if not isinstance(module, Cell):
491 raise ValueError(f"module's type must be nn.cell but got {type(module)}.")
492 else:
493 from torch.nn import Module
494 if not isinstance(module, Module):
495 raise ValueError(f"module's type must be nn.Module but got {type(module)}.")
498def _validate_module_for_fully_shard(
499 module: Union[ModuleClass, List[ModuleClass]], platform_type
500) -> None:
501 """Validate module(s) for fully_shard. Platform-aware for single module."""
502 if isinstance(module, list):
503 if len(module) == 0:
504 raise ValueError("fully_shard does not support empty list of modules.")
505 for i, m in enumerate(module):
506 try:
507 _check_module_valid(platform_type, m)
508 except ValueError:
509 raise ValueError(
510 f"fully_shard expects nn.Module or list[nn.Module], "
511 f"but got list with {type(m).__name__} at index {i}."
512 ) from None
513 else:
514 _check_module_valid(platform_type, module)
517HsdpValidationOptions = namedtuple(
518 "HsdpValidationOptions",
519 [
520 "shard_size",
521 "threshold",
522 "optimizer_level",
523 "enable_grad_accumulation",
524 "grad_scale",
525 "reduce_dtype",
526 "comm_async",
527 "comm_fusion",
528 "bucket_size",
529 ],
530)
533def _validate_hsdp_shard_size(shard_size: int) -> None:
534 if not isinstance(shard_size, int) or (shard_size <= 0 and shard_size != -1):
535 raise ValueError(f"shard_size must be a positive integer, but got {shard_size}.")
538def _validate_hsdp_threshold(threshold: int) -> None:
539 if not isinstance(threshold, int) or threshold < 0:
540 raise ValueError(f"threshold must be a positive integer or 0, but got {threshold}.")
543def _validate_hsdp_optimizer_level(optimizer_level: str) -> None:
544 if optimizer_level not in ["level1", "level2", "level3"]:
545 raise ValueError(
546 f"Optimizer level should in ['level1', 'level2', 'level3'], but got {optimizer_level}."
547 )
550def _validate_hsdp_reduce_dtype(platform_type: PlatformType, reduce_dtype) -> None:
551 if platform_type == PlatformType.MINDSPORE:
552 from mindspore._c_expression.typing import Type
553 if reduce_dtype is not None and not isinstance(reduce_dtype, Type):
554 raise ValueError(f"reduce_dtype must be mindspore.dtype but got {reduce_dtype}.")
555 return
556 import torch
557 if reduce_dtype is not None and not isinstance(reduce_dtype, torch.dtype):
558 raise ValueError(f"reduce_dtype must be torch.dtype but got {reduce_dtype}.")
561def _check_hsdp_input_valid(platform_type, module, options: HsdpValidationOptions):
562 """check hsdp input valid"""
563 _check_module_valid(platform_type, module)
564 _validate_hsdp_shard_size(options.shard_size)
565 _validate_hsdp_threshold(options.threshold)
566 _validate_hsdp_optimizer_level(options.optimizer_level)
567 if not isinstance(options.enable_grad_accumulation, bool):
568 raise ValueError(
569 f"enable_grad_accumulation must be bool but got {options.enable_grad_accumulation}."
570 )
571 if not isinstance(options.grad_scale, float):
572 raise ValueError(f"grad_scale must be float but got {options.grad_scale}.")
573 _validate_hsdp_reduce_dtype(platform_type, options.reduce_dtype)
574 if not isinstance(options.comm_async, bool):
575 raise ValueError(f"comm_async must be bool but got {options.comm_async}.")
576 if not isinstance(options.comm_fusion, bool):
577 raise ValueError(f"comm_fusion must be bool but got {options.comm_fusion}.")
578 if not isinstance(options.bucket_size, int) or (
579 options.bucket_size < 0 and options.bucket_size != -1
580 ):
581 raise ValueError(
582 f"bucket_size must be a positive integer or 0, but got {options.bucket_size}."
583 )
586def _get_device_from_mesh(mesh: DeviceMesh):
587 """Extract and validate the torch device from the device mesh."""
588 device = None
589 device_type = mesh.device_type
590 if device_type not in ("npu", "cuda"):
591 raise AssertionError(
592 f"hyper_parallel.fully_shard support device in [torch.npu, torch.cuda], "
593 f"but got '{device_type}'"
594 )
595 if platform.platform_type == PlatformType.PYTORCH:
596 device_handle = platform.get_device_handle(device_type)
597 if device_handle is None:
598 raise ValueError(
599 f"hyper_parallel.fully_shard can't find device_handle of "
600 f"'torch.{device_type}', check the environment."
601 )
602 if device_handle.is_available():
603 import torch
604 device = torch.device(device_handle.current_device())
605 else:
606 device = device_type
607 return device
610def _normalize_replicate_params(
611 replicate_params: Optional[set[platform.Parameter]],
612) -> set[platform.Parameter]:
613 """
614 Normalize replicate_params for fully_shard
615 Args:
616 replicate_params (Optional[set[nn.Parameter]]): Set of parameters to exclude from sharding.
617 Returns:
618 set[nn.Parameter]: Set of parameters to exclude from sharding.
619 """
620 if replicate_params is None:
621 return set()
622 out = set(replicate_params)
623 for p in out:
624 if not isinstance(p, (platform.Parameter, DTensor)):
625 raise TypeError(
626 "replicate_params must contain only nn.Parameter or DTensor, "
627 f"got {type(p).__name__}."
628 )
629 return out
632def _get_modules_parameters(modules, ignored_params=None):
633 """Collect deduplicated parameters from module roots."""
634 return get_managed_modules_parameters(modules, ignored_params)
636def _validate_managed_params_source_shard_infos(
637 managed_parameters: set[ParameterClass],
638 source_shard_infos: Optional[Mapping[ParameterClass, SourceShardMetaInfo]],
639) -> None:
640 """Validate the parameter-identity metadata consumed by one fully_shard unit."""
641 if source_shard_infos is None:
642 return
643 if not isinstance(source_shard_infos, Mapping):
644 raise ValueError("source_shard_infos must be a mapping from Parameter to SourceShardMetaInfo")
646 managed_parameters = set(managed_parameters)
648 if any(isinstance(parameter, DTensor) for parameter in managed_parameters):
649 raise ValueError(
650 "source_shard_infos cannot be provided when fully_shard manages a native DTensor parameter"
651 )
653 invalid_values = [
654 parameter
655 for parameter, metadata in source_shard_infos.items()
656 if not isinstance(metadata, SourceShardMetaInfo)
657 ]
658 if invalid_values:
659 raise ValueError("source_shard_infos values must be SourceShardMetaInfo instances")
661 invalid_origins = [
662 parameter
663 for parameter, metadata in source_shard_infos.items()
664 if metadata.origin_is_dtensor and not isinstance(parameter, DTensor)
665 ]
666 if invalid_origins:
667 raise ValueError(
668 "source_shard_infos origin_is_dtensor=True requires a native DTensor parameter"
669 )
671 missing_parameters = managed_parameters.difference(source_shard_infos)
672 if missing_parameters:
673 raise ValueError(
674 "source_shard_infos must cover every parameter managed by this fully_shard call"
675 )
676 unexpected_parameters = set(source_shard_infos).difference(managed_parameters)
677 if unexpected_parameters:
678 raise ValueError(
679 "source_shard_infos contains a parameter not managed by this fully_shard call"
680 )
683def fully_shard(
684 module: Union[ModuleClass, List[ModuleClass]],
685 *,
686 mesh: Optional[DeviceMesh] = None,
687 reshard_after_forward: bool = True,
688 shard_placement_fn: None = None,
689 mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
690 offload_policy: OffloadPolicy = OffloadPolicy(),
691 ignored_params: Optional[set[platform.Parameter]] = None,
692 replicate_params: Optional[set[platform.Parameter]] = None,
693 comm_fusion: bool = False,
694 comm_fusion_zero_copy: Optional[bool] = None,
695 source_shard_infos: Optional[Mapping[ParameterClass, SourceShardMetaInfo]] = None,
696) -> Union[ModuleClass, List[ModuleClass]]:
698 """
699 Apply fully_shard to a module (or list of modules) for distributed training with parameter sharding.
701 This interface provides PyTorch-compatible HSDP (Hybrid Sharded Data Parallelism)
702 functionality, enabling efficient training of large models by sharding parameters
703 across multiple devices. The module is automatically enhanced with distributed
704 capabilities including parameter sharding, gradient synchronization, and memory
705 management.
707 When a list of modules is passed, they are treated as one FSDP unit (parameters
708 grouped together). Both PyTorch and MindSpore platforms support list input.
710 Parameters:
711 module (nn.Module or List[nn.Module]):
712 The module(s) to apply fully_shard to. Modified in-place. When a list
713 is passed, parameters from all modules are grouped as one FSDP unit.
715 mesh (Optional[DeviceMesh], default=None):
716 The device mesh defining the process topology for distributed training.
717 If None, fully_shard creates a default 1D mesh for local parameters.
718 Native DTensor parameters require an explicit DP submesh.
720 reshard_after_forward (bool, default=True):
721 Whether to automatically reshard parameters after forward. When True,
722 parameters are resharded immediately after they are no longer needed,
723 freeing memory for subsequent operations. Set to False if you want to
724 keep parameters unsharded for backward pass or manual control.
726 shard_placement_fn (Callable, default=None):
727 A callable that determines how to shard each parameter. The function
728 should accept a parameter and return a Shard object specifying the
729 sharding dimension, or None to use default sharding (dimension 0)
731 mp_policy (MixedPrecisionPolicy, default=MixedPrecisionPolicy()):
732 Mixed precision training policy controlling data type conversions.
734 offload_policy (OffloadPolicy, default=OffloadPolicy()):
735 Memory offload policy for reducing device memory usage.
737 ignored_params (Optional[set[nn.Parameter]], default=None):
738 Set of parameters to exclude from fully_shard management entirely.
739 These parameters are left on the original module as regular parameters,
740 are not sharded, and do not participate in fully_shard gradient
741 synchronization. Use this for parameters that should remain outside
742 the fully_shard lifecycle.
744 replicate_params (Optional[set[nn.Parameter]], default=None):
745 Set of parameters to keep replicated while still managing them under
746 fully_shard. These parameters are not sharded, but their gradients
747 are still synchronized with DDP-style all-reduce over the current
748 fully_shard communication domain. This differs from ``ignored_params``,
749 which skips fully_shard management and gradient synchronization
750 entirely for the selected parameters.
752 comm_fusion (bool, default=False):
753 Whether enable all_gather fusion and reduce_scatter fusion.
755 comm_fusion_zero_copy (Optional[bool], default=None):
756 Whether allow the experimental zero-copy path for
757 ``comm_fusion``. When set to ``None``, fully_shard uses a backend-specific
758 default:
759 - PyTorch: enabled automatically when ``comm_fusion=True``
760 - MindSpore: disabled automatically even when ``comm_fusion=True``
761 When enabled, fully_shard may rebase sharded local parameter storage
762 into one shared flat buffer so fused all-gather can read directly from
763 contiguous memory. This path depends on optimizer compatibility with
764 view-backed parameters.
765 source_shard_infos (Optional[Mapping[nn.Parameter, SourceShardMetaInfo]]):
766 Source TP/EP mesh and placements for the plain-parameter dual mode.
767 This interface is currently supported by the Torch backend only.
769 Returns:
770 nn.Module or List[nn.Module]: The input module(s) with HSDP capabilities added.
771 """
772 platform_type = platform.platform_type
773 _validate_module_for_fully_shard(module, platform_type)
775 if source_shard_infos is not None and platform_type != PlatformType.PYTORCH:
776 raise NotImplementedError("source_shard_infos is currently supported only on the Torch backend")
777 if platform_type == PlatformType.MINDSPORE:
778 from hyper_parallel.platform.mindspore.autograd_compat import enable_mindspore_backward_compat
780 enable_mindspore_backward_compat()
782 arg_module = module
783 if isinstance(module, list):
784 modules = tuple(_get_root_modules(module))
785 else:
786 modules = (module,)
788 for mod in modules:
789 _extend_module_with_hsdp_interface(mod)
791 params = _get_modules_parameters(modules, ignored_params)
792 has_dtensor_param = any(isinstance(param, DTensor) for param in params)
793 if platform_type == PlatformType.PYTORCH:
794 _validate_managed_params_source_shard_infos(set(params), source_shard_infos)
796 replicate_params = _normalize_replicate_params(replicate_params)
798 if mesh is None:
799 mesh = init_device_mesh(device_type="npu", mesh_shape=(platform.get_world_size(),))
800 if has_dtensor_param:
801 raise ValueError(
802 "fully_shard does not support mesh=None with a native DTensor parameter; "
803 "pass an explicit DP submesh instead."
804 )
805 device = _get_device_from_mesh(mesh)
807 init_modules = modules
808 modules[0].hsdp_init(
809 platform_type,
810 init_modules,
811 mesh,
812 reshard_after_forward,
813 shard_placement_fn,
814 mp_policy,
815 offload_policy,
816 ignored_params,
817 replicate_params,
818 device,
819 comm_fusion,
820 comm_fusion_zero_copy,
821 source_shard_infos=source_shard_infos,
822 )
823 # Share the same scheduler handle with other roots so mods[i].unshard()/prefetch work
824 if len(modules) > 1:
825 for mod in modules[1:]:
826 mod.hsdp_scheduler = modules[0].hsdp_scheduler
827 return arg_module
830def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
831 """Get model state dict with platform-specific implementation.
833 Delegates to the platform-specific implementation at runtime.
834 Users import from here instead of platform internals.
836 Args:
837 model: The model whose state dict to retrieve.
838 options: Optional :class:`StateDictOptions` controlling gathering
839 (``full_state_dict``), offloading (``cpu_offload``) and frozen
840 parameter filtering (``ignore_frozen_params``).
842 Returns:
843 The model state dict. By default values are sharded DTensors; when
844 ``full_state_dict=True`` values are full Tensors (CPU on rank 0 only
845 when ``cpu_offload=True``).
846 """
847 return platform.get_model_state_dict(model, options=options)
850def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
851 """Set model state dict with platform-specific implementation.
853 Delegates to the platform-specific implementation at runtime. Full tensors
854 are scattered into DTensor shards matching the model's layout before the
855 in-place load.
857 Args:
858 model: The model to load state into.
859 model_state_dict: State dict to load. Values may be plain (global)
860 tensors when ``full_state_dict=True`` or sharded DTensors otherwise.
861 options: Optional :class:`StateDictOptions` controlling scattering
862 (``full_state_dict``), device placement (``cpu_offload``) and
863 strictness (``strict``).
865 Returns:
866 None.
867 """
868 return platform.set_model_state_dict(model, model_state_dict, options=options)
871def hsdp_sync_stream():
872 """Wait for hsdp gradient handle to be completed."""
873 platform.wait_grad_handle()