Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / api.py: 78%
332 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""hybrid shard data parallel interface"""
16from collections import namedtuple
17from typing import Any, List, Mapping, cast, Optional, Union
19import torch
20import torch.distributed as dist
21from torch import Tensor, nn
23from hyper_parallel.core.fully_shard.utils import (
24 MixedPrecisionPolicy,
25 OffloadPolicy,
26 SourceShardMetaInfo,
27 load_into_param,
28 parameters_dict,
29 wait_grad_handle,
30)
31from hyper_parallel.core.utils.communication import get_device_handle
32from hyper_parallel.core.fully_shard.hsdp_scheduler import HSDPSchedulerV2
33from hyper_parallel.core.fully_shard.state_dict_utils import (
34 get_model_state_dict as get_model_state_dict_impl,
35 set_model_state_dict as set_model_state_dict_impl,
36)
37from hyper_parallel import DeviceMesh, init_device_mesh
38from hyper_parallel.core.dtensor.dtensor import DTensor, distribute_tensor
39from hyper_parallel.core.fully_shard.hsdp_utils import (
40 get_managed_modules_parameters,
41)
43origin_class_to_extend_class = {}
46def _resolve_comm_fusion_zero_copy_default(
47 comm_fusion: bool,
48 comm_fusion_zero_copy: Optional[bool],
49) -> bool:
50 """Resolve the default for the comm_fusion zero-copy path."""
51 if comm_fusion_zero_copy is not None:
52 return comm_fusion_zero_copy
53 return comm_fusion
56def _check_strict_keys(
57 module: nn.Module, state_dict: Mapping[str, Any],
58) -> None:
59 """Raise ``RuntimeError`` if *state_dict* keys do not match *module*."""
60 expected_keys = set(module.state_dict().keys())
61 missing = expected_keys - set(state_dict.keys())
62 unexpected = set(state_dict.keys()) - expected_keys
63 error_msgs: list[str] = []
64 if missing:
65 error_msgs.append(
66 "Missing key(s): " + ", ".join(repr(k) for k in sorted(missing))
67 )
68 if unexpected:
69 error_msgs.append(
70 "Unexpected key(s): " + ", ".join(repr(k) for k in sorted(unexpected))
71 )
72 if error_msgs:
73 raise RuntimeError(
74 f"Error(s) in loading state_dict for "
75 f"{module.__class__.__name__}:\n\t"
76 + "\n\t".join(error_msgs)
77 )
80def _resolve_local_tensor(
81 key: str, val: Tensor, target: DTensor,
82) -> Tensor:
83 """Return the local shard tensor to be loaded into *target*."""
84 if isinstance(val, DTensor):
85 return val.to_local()
86 local_shape = tuple(target.local_shape)
87 global_shape = tuple(target.shape)
88 val_shape = tuple(val.shape)
89 if val_shape == local_shape:
90 return val
91 if val_shape == global_shape:
92 wrapped = distribute_tensor(
93 val, target.device_mesh,
94 target.layout.alias_placements if target.layout else target.placements,
95 )
96 return wrapped.to_local()
98 raise ValueError(
99 f"load '{key}': plain tensor shape {val_shape} "
100 f"matches neither local shard {local_shape} "
101 f"nor global {global_shape}."
102 )
105class _UnshardHandle:
106 """Unshard handle for user call HSDPModule.unshard(async_op=True)"""
107 def __init__(self, hsdp_state=None):
108 """
109 Initialize an async unshard handle.
111 Args:
112 hsdp_state (HSDPState, optional): The state to wait on. None means a no-op handle.
113 """
114 self._hsdp_state = hsdp_state
116 def wait(self):
117 """Block until the async unshard operation completes."""
118 if self._hsdp_state is not None:
119 self._hsdp_state.wait_for_unshard()
120 self._hsdp_state = None
123class HSDPModule:
124 """
125 The hsdp block of neural networks with hsdp interface.
126 """
128 def __init__(self):
129 """Initialize HSDPModule."""
130 self.hsdp_scheduler = None # Initialized in hsdp_init()
132 def _apply(self, *args: Any, **kwargs: Any) -> Any:
133 """Apply a tensor conversion while preserving HSDP parameter storage.
135 ``nn.Module._apply`` can replace the DTensor local storage while an
136 unsharded parameter buffer is still materialized on the device. HSDP
137 must reshard before delegating to the conversion and rebuild its
138 sharded-parameter bookkeeping afterwards. The bookkeeping reset is a
139 storage operation and must not become part of an autograd graph.
141 The module's original class is added after HSDPModule in the runtime
142 MRO, so static analysis cannot resolve the delegated ``_apply``.
143 """
144 if not hasattr(self, "hsdp_scheduler") or self.hsdp_scheduler is None:
145 return super()._apply(*args, **kwargs) # pylint: disable=no-member
146 self.reshard()
147 result = super()._apply(*args, **kwargs) # pylint: disable=no-member
148 hsdp_state = self.hsdp_scheduler.hsdp_state
149 if hsdp_state is None:
150 return result
151 with torch.no_grad():
152 for hsdp_param in hsdp_state.hsdp_params:
153 hsdp_param.reset_sharded_param()
154 if hsdp_state.param_group is not None:
155 for bucket in hsdp_state.param_group.all_gather_buckets:
156 # Conversion rebinds shards; an obsolete zero-copy buffer
157 # must not keep the previous device allocation resident.
158 if not bucket.is_flat_buffer_valid():
159 bucket.flat_param_buffer = None
160 return result
162 def hsdp_init(self, module, mesh, reshard_after_forward,
163 shard_placement_fn, mp_policy, offload_policy, ignored_params, replicate_params, device,
164 comm_fusion, comm_fusion_zero_copy: Optional[bool] = None,
165 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]] = None):
166 """init hsdp2 scheduler."""
167 resolved_comm_fusion_zero_copy = _resolve_comm_fusion_zero_copy_default(
168 comm_fusion,
169 comm_fusion_zero_copy,
170 )
172 self.hsdp_scheduler = HSDPSchedulerV2(module,
173 mesh,
174 reshard_after_forward,
175 shard_placement_fn,
176 mp_policy,
177 offload_policy,
178 ignored_params,
179 replicate_params,
180 device,
181 comm_fusion,
182 resolved_comm_fusion_zero_copy,
183 source_shard_infos=source_shard_infos,
184 )
186 def set_requires_gradient_sync(self, requires_grad_sync: bool, *, recurse: bool = True) -> None:
187 """Set gradient synchronization for this module and optionally its descendants."""
188 if not isinstance(requires_grad_sync, bool):
189 raise ValueError("requires_grad_sync must be bool")
190 if not isinstance(recurse, bool):
191 raise ValueError("recurse must be bool")
192 self_module = cast(nn.Module, self)
193 for _, module in self_module.named_modules() if recurse else (("", self),):
194 if isinstance(module, HSDPModule):
195 module.hsdp_scheduler.set_requires_grad_sync(requires_grad_sync)
197 def zero_grad(self, set_to_none: bool = True) -> None:
198 """Clear parameter gradients according to set_to_none."""
199 super().zero_grad(set_to_none=set_to_none)
201 def set_modules_to_forward_prefetch(self, modules):
202 """set forward 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 hsdp interface first.")
210 self.hsdp_scheduler.set_forward_prefetch_cells(modules)
212 def set_modules_to_backward_prefetch(self, modules):
213 """set backward prefetch module list to prefetch all gather for unsharded parameters"""
214 if not isinstance(modules, (tuple, list)):
215 raise ValueError("modules must be HSDPModule list")
216 for module in modules:
217 if not isinstance(module, HSDPModule):
218 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
219 if not hasattr(self, "hsdp_scheduler"):
220 raise ValueError("call fully_shard interface first.")
221 self.hsdp_scheduler.set_backward_prefetch_cells(modules)
223 def reshard(self) -> None:
224 """reshard all sharded parameters"""
225 if not self.hsdp_scheduler:
226 raise ValueError("hsdp_scheduler is None")
227 hsdp_state = self.hsdp_scheduler.hsdp_state
228 if hsdp_state:
229 hsdp_state.shard()
231 def unshard(self, async_op: bool = False):
232 """unshard all sharded parameters"""
233 if not isinstance(async_op, bool):
234 raise ValueError(f"async_op should be a bool, got {type(async_op)}")
235 if not self.hsdp_scheduler:
236 raise ValueError("hsdp_scheduler is None")
237 hsdp_state = self.hsdp_scheduler.hsdp_state
238 if hsdp_state:
239 hsdp_state.unshard(async_op) # pylint: disable=too-many-function-args
240 if async_op:
241 return _UnshardHandle(hsdp_state=hsdp_state)
242 return None
244 def load_state_dict(
245 self,
246 state_dict: Mapping[str, Any],
247 strict: bool = True,
248 assign: bool = False,
249 ):
250 """
251 Load state dict by copying directly into local shards.
253 Bypasses ``super().load_state_dict()`` because the standard PyTorch
254 implementation triggers ``copy_`` through the DTensor dispatcher, which
255 is not registered in the hyper-parallel layout system.
257 Each value in ``state_dict`` is dispatched by type:
258 - hyper DTensor: extract local shard and copy directly.
259 - plain Tensor whose shape == local shard shape: copy as-is.
260 - plain Tensor whose shape == global shape: distribute via
261 ``distribute_tensor``, then copy the local shard.
263 Args:
264 state_dict (Mapping[str, Any]): Fully-qualified parameter/buffer
265 names mapped to tensors (DTensor or plain Tensor).
266 strict (bool): If ``True`` (default), missing or unexpected keys
267 raise ``RuntimeError``, matching ``nn.Module.load_state_dict``
268 semantics.
269 assign (bool): When ``True`` *and* every value in ``state_dict`` is
270 already a hyper DTensor, defer to the standard
271 ``nn.Module.load_state_dict(assign=True)``, which replaces the
272 module's parameters/buffers with the given DTensors instead of
273 copying into existing storage. This is required when loading
274 sharded DTensors onto a meta-device model (e.g.
275 ``cpu_ram_efficient_loading``). If ``state_dict`` contains any
276 plain tensor (local-shard or global shape), ``assign`` is
277 ignored and the copy/distribute path below is used so the
278 target stays a properly sharded DTensor.
280 Raises:
281 RuntimeError: When ``strict`` is ``True`` and keys do not match.
282 ValueError: When a plain tensor shape matches neither the local
283 shard shape nor the global shape of the target DTensor.
284 """
285 if assign and state_dict and all(
286 isinstance(val, DTensor) for val in state_dict.values()
287 ):
288 return super().load_state_dict(state_dict, strict=strict, assign=True)
289 self_module = cast(nn.Module, self)
291 target_map: dict[str, Tensor] = {}
292 for name, p in parameters_dict(self_module):
293 target_map[name] = p
294 for name, b in self_module.named_buffers():
295 target_map[name] = b
297 if strict:
298 _check_strict_keys(self_module, state_dict)
300 with torch.no_grad():
301 for key, val in state_dict.items():
302 target = target_map.get(key)
303 if target is None:
304 continue
306 if isinstance(target, DTensor):
307 val = _resolve_local_tensor(key, val, target)
308 load_into_param(target, val)
310 # Trigger load_state_dict post-hooks so that HSDP internal
311 # bookkeeping (e.g. _sharded_param_data) stays in sync.
312 # Pass an IncompatibleKeys with the same attribute names as PyTorch
313 # so external hooks can safely read .missing_keys/.unexpected_keys.
314 _IK = namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"])
315 incompatible_keys = _IK([], [])
316 for _, module in self_module.named_modules():
317 hooks = module._load_state_dict_post_hooks # pylint: disable=protected-access
318 for hook in hooks.values():
319 hook(module, incompatible_keys)
321 def set_is_last_backward(self, is_last_backward: bool):
322 """set is_last_backward flag"""
323 self.hsdp_scheduler.scheduler_ctx.is_last_backward = is_last_backward
325 def reset_iter_state(self, recursive: bool = True) -> None:
326 """Reset fully_shard iteration bookkeeping without clearing optimizer gradients.
328 This method must be called only after communication for the iteration has
329 completed. It does not wait in-flight collectives or clear
330 ``sharded_param.grad``/``main_grad``.
332 Args:
333 recursive: Whether to reset every HSDPModule in this module tree.
335 Raises:
336 ValueError: If ``recursive`` is not a bool.
337 """
338 if not isinstance(recursive, bool):
339 raise ValueError(f"recursive should be a bool, got {type(recursive)}")
340 self_module = cast(nn.Module, self)
341 modules = (
342 [module for _, module in self_module.named_modules()]
343 if recursive
344 else [self_module]
345 )
346 for module in modules:
347 if isinstance(module, HSDPModule):
348 module.hsdp_scheduler.reset_iter_state()
350 def set_requires_all_reduce(self, requires_all_reduce: bool, *, recurse: bool = True) -> None:
351 """Set all-reduce for this module and optionally its descendants."""
352 if not isinstance(requires_all_reduce, bool):
353 raise ValueError("requires_all_reduce must be bool")
354 if not isinstance(recurse, bool):
355 raise ValueError("recurse must be bool")
356 self_module = cast(nn.Module, self)
357 for _, module in self_module.named_modules() if recurse else (("", self),):
358 if isinstance(module, HSDPModule):
359 module.hsdp_scheduler.set_requires_all_reduce(requires_all_reduce)
361 def set_reshard_after_forward(self, reshard_after_forward: bool, recurse: bool = True) -> None:
362 """Set forward resharding for this module and optionally its descendants."""
363 if not isinstance(reshard_after_forward, bool):
364 raise ValueError("reshard_after_forward must be bool")
365 if not isinstance(recurse, bool):
366 raise ValueError("recurse must be bool")
367 self_module = cast(nn.Module, self)
368 for _, module in self_module.named_modules() if recurse else (("", self),):
369 if isinstance(module, HSDPModule):
370 module.hsdp_scheduler.set_reshard_after_forward(reshard_after_forward)
372 def set_reshard_after_backward(self, reshard_after_backward: bool, recurse: bool = True) -> None:
373 """Set backward resharding for this module and optionally its descendants."""
374 if not isinstance(reshard_after_backward, bool):
375 raise ValueError("reshard_after_backward must be bool")
376 if not isinstance(recurse, bool):
377 raise ValueError("recurse must be bool")
378 self_module = cast(nn.Module, self)
379 for _, module in self_module.named_modules() if recurse else (("", self),):
380 if isinstance(module, HSDPModule):
381 module.hsdp_scheduler.set_reshard_after_backward(reshard_after_backward)
383 def set_reduce_op_type(self, reduce_op_type, recurse: bool = True) -> None:
384 """
385 set reduce_op_type for all reduce operations in HSDP
386 support reduce_op_type "avg" and "sum", default is "avg"
387 """
388 self_module = cast(nn.Module, self)
389 if recurse:
390 sub_modules = [m for _, m in self_module.named_modules()]
391 else:
392 sub_modules = [self_module]
393 for module in sub_modules:
394 if isinstance(module, HSDPModule):
395 hsdp_state = module.hsdp_scheduler.hsdp_state
396 if hsdp_state:
397 hsdp_state.set_reduce_op_type(reduce_op_type)
399 def set_gradient_scaling_factor(self, factor=None):
400 """
401 Set a multiplicative scaling factor applied to gradients after
402 reduce-scatter / all-reduce and before they are written into
403 ``sharded_param.grad``.
405 ``factor`` may be ``None`` (disable scaling), a Python ``float``/``int``,
406 or a 0-dim/1-element tensor. Setting ``factor`` to ``None`` (the default
407 on construction) skips the scaling op entirely so no extra device-side
408 ``mul_`` is launched on the hot path.
410 Args:
411 factor (None | float | int | torch.Tensor): Scaling coefficient.
412 Use ``None`` to disable scaling.
414 Raises:
415 ValueError: If ``factor`` is not one of the supported types or is a
416 tensor with more than one element.
417 """
418 if factor is not None:
419 if isinstance(factor, bool):
420 raise ValueError(
421 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
422 f"but got bool {factor}."
423 )
424 if isinstance(factor, Tensor):
425 if factor.numel() != 1:
426 raise ValueError(
427 f"gradient_scaling_factor tensor must have exactly 1 element, "
428 f"but got shape {tuple(factor.shape)}."
429 )
430 elif not isinstance(factor, (float, int)):
431 raise ValueError(
432 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
433 f"but got {type(factor).__name__}."
434 )
435 hsdp_state = self.hsdp_scheduler.hsdp_state
436 if hsdp_state:
437 hsdp_state.set_gradient_scaling_factor(factor)
440def _extend_module_with_hsdp_interface(module):
441 """Dynamically extend module's class to inherit from HSDPModule, adding HSDP capabilities."""
442 origin_class = module.__class__
443 extend_class = origin_class_to_extend_class.get(origin_class, None)
444 if extend_class is None:
445 extend_class = type(f"HSDP{origin_class.__name__}", (HSDPModule, origin_class), {})
446 origin_class_to_extend_class[origin_class] = extend_class
447 module.__class__ = extend_class
450def _get_root_modules(modules: List[nn.Module]) -> List[nn.Module]:
451 """
452 Returns the modules in ``modules`` that are root modules (i.e. parent-less)
453 with respect to the set ``modules``. In other words, these are the modules
454 in ``modules`` that are not the child of any other module in ``modules``.
456 Aligned with PyTorch torch.distributed.utils._get_root_modules.
457 """
458 root_modules: List[nn.Module] = []
460 def _get_submodules(mod):
461 return set(mod.modules())
463 module_to_modules: dict[nn.Module, set] = {
464 m: _get_submodules(m) for m in modules
465 }
466 for candidate in modules:
467 is_root = True
468 for mod, submodules in module_to_modules.items():
469 if candidate is not mod and candidate in submodules:
470 is_root = False
471 break
472 if is_root:
473 root_modules.append(candidate)
474 return root_modules
477def _check_module_valid(module):
478 """check module valid"""
479 if not isinstance(module, nn.Module):
480 raise ValueError(f"module's type must be nn.Module but got {type(module)}.")
483def _validate_module_for_fully_shard(
484 module: Union[nn.Module, List[nn.Module]]
485) -> None:
486 """Validate module(s) for fully_shard."""
487 if isinstance(module, list):
488 if len(module) == 0:
489 raise ValueError("fully_shard does not support empty list of modules.")
490 for i, m in enumerate(module):
491 try:
492 _check_module_valid(m)
493 except ValueError:
494 raise ValueError(
495 f"fully_shard expects nn.Module or list[nn.Module], "
496 f"but got list with {type(m).__name__} at index {i}."
497 ) from None
498 else:
499 _check_module_valid(module)
502HsdpValidationOptions = namedtuple(
503 "HsdpValidationOptions",
504 [
505 "shard_size",
506 "threshold",
507 "optimizer_level",
508 "enable_grad_accumulation",
509 "grad_scale",
510 "reduce_dtype",
511 "comm_async",
512 "comm_fusion",
513 "bucket_size",
514 ],
515)
518def _validate_hsdp_shard_size(shard_size: int) -> None:
519 if not isinstance(shard_size, int) or (shard_size <= 0 and shard_size != -1):
520 raise ValueError(f"shard_size must be a positive integer, but got {shard_size}.")
523def _validate_hsdp_threshold(threshold: int) -> None:
524 if not isinstance(threshold, int) or threshold < 0:
525 raise ValueError(f"threshold must be a positive integer or 0, but got {threshold}.")
528def _validate_hsdp_optimizer_level(optimizer_level: str) -> None:
529 if optimizer_level not in ["level1", "level2", "level3"]:
530 raise ValueError(
531 f"Optimizer level should in ['level1', 'level2', 'level3'], but got {optimizer_level}."
532 )
535def _validate_hsdp_reduce_dtype(reduce_dtype) -> None:
536 """Check ``reduce_dtype`` is a torch dtype or ``None``."""
537 if reduce_dtype is not None and not isinstance(reduce_dtype, torch.dtype):
538 raise ValueError(f"reduce_dtype must be torch.dtype but got {reduce_dtype}.")
541def _check_hsdp_input_valid(module, options: HsdpValidationOptions):
542 """check hsdp input valid"""
543 _check_module_valid(module)
544 _validate_hsdp_shard_size(options.shard_size)
545 _validate_hsdp_threshold(options.threshold)
546 _validate_hsdp_optimizer_level(options.optimizer_level)
547 if not isinstance(options.enable_grad_accumulation, bool):
548 raise ValueError(
549 f"enable_grad_accumulation must be bool but got {options.enable_grad_accumulation}."
550 )
551 if not isinstance(options.grad_scale, float):
552 raise ValueError(f"grad_scale must be float but got {options.grad_scale}.")
553 _validate_hsdp_reduce_dtype(options.reduce_dtype)
554 if not isinstance(options.comm_async, bool):
555 raise ValueError(f"comm_async must be bool but got {options.comm_async}.")
556 if not isinstance(options.comm_fusion, bool):
557 raise ValueError(f"comm_fusion must be bool but got {options.comm_fusion}.")
558 if not isinstance(options.bucket_size, int) or (
559 options.bucket_size < 0 and options.bucket_size != -1
560 ):
561 raise ValueError(
562 f"bucket_size must be a positive integer or 0, but got {options.bucket_size}."
563 )
566def _get_device_from_mesh(mesh: DeviceMesh):
567 """Extract and validate the torch device from the device mesh."""
568 device = None
569 device_type = mesh.device_type
570 if device_type not in ("npu", "cuda"):
571 raise AssertionError(
572 f"hyper_parallel.fully_shard support device in [torch.npu, torch.cuda], "
573 f"but got '{device_type}'"
574 )
575 device_handle = get_device_handle(device_type)
576 if device_handle is None:
577 raise ValueError(
578 f"hyper_parallel.fully_shard can't find device_handle of "
579 f"'torch.{device_type}', check the environment."
580 )
581 if device_handle.is_available():
582 device = torch.device(device_handle.current_device())
583 return device
586def _normalize_replicate_params(
587 replicate_params: Optional[set[nn.Parameter]],
588) -> set[nn.Parameter]:
589 """
590 Normalize replicate_params for fully_shard
591 Args:
592 replicate_params (Optional[set[nn.Parameter]]): Set of parameters to exclude from sharding.
593 Returns:
594 set[nn.Parameter]: Set of parameters to exclude from sharding.
595 """
596 if replicate_params is None:
597 return set()
598 out = set(replicate_params)
599 for p in out:
600 if not isinstance(p, (nn.Parameter, DTensor)):
601 raise TypeError(
602 "replicate_params must contain only nn.Parameter or DTensor, "
603 f"got {type(p).__name__}."
604 )
605 return out
608def _get_modules_parameters(modules, ignored_params=None):
609 """Collect deduplicated parameters from module roots."""
610 return get_managed_modules_parameters(modules, ignored_params)
613def _validate_managed_params_source_shard_infos(
614 managed_parameters: set[nn.Parameter],
615 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]],
616) -> None:
617 """Validate the parameter-identity metadata consumed by one fully_shard unit."""
618 if source_shard_infos is None:
619 return
620 if not isinstance(source_shard_infos, Mapping):
621 raise ValueError("source_shard_infos must be a mapping from Parameter to SourceShardMetaInfo")
623 managed_parameters = set(managed_parameters)
625 if any(isinstance(parameter, DTensor) for parameter in managed_parameters):
626 raise ValueError(
627 "source_shard_infos cannot be provided when fully_shard manages a native DTensor parameter"
628 )
630 invalid_values = [
631 parameter
632 for parameter, metadata in source_shard_infos.items()
633 if not isinstance(metadata, SourceShardMetaInfo)
634 ]
635 if invalid_values:
636 raise ValueError("source_shard_infos values must be SourceShardMetaInfo instances")
638 invalid_origins = [
639 parameter
640 for parameter, metadata in source_shard_infos.items()
641 if metadata.origin_is_dtensor and not isinstance(parameter, DTensor)
642 ]
643 if invalid_origins:
644 raise ValueError(
645 "source_shard_infos origin_is_dtensor=True requires a native DTensor parameter"
646 )
648 missing_parameters = managed_parameters.difference(source_shard_infos)
649 if missing_parameters:
650 raise ValueError(
651 "source_shard_infos must cover every parameter managed by this fully_shard call"
652 )
653 unexpected_parameters = set(source_shard_infos).difference(managed_parameters)
654 if unexpected_parameters:
655 raise ValueError(
656 "source_shard_infos contains a parameter not managed by this fully_shard call"
657 )
660def fully_shard(
661 module: Union[nn.Module, List[nn.Module]],
662 *,
663 mesh: Optional[DeviceMesh] = None,
664 reshard_after_forward: bool = True,
665 shard_placement_fn: None = None,
666 mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
667 offload_policy: OffloadPolicy = OffloadPolicy(),
668 ignored_params: Optional[set[nn.Parameter]] = None,
669 replicate_params: Optional[set[nn.Parameter]] = None,
670 comm_fusion: bool = False,
671 comm_fusion_zero_copy: Optional[bool] = None,
672 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]] = None,
673) -> Union[nn.Module, List[nn.Module]]:
675 """
676 Apply fully_shard to a module (or list of modules) for distributed training with parameter sharding.
678 This interface provides PyTorch-compatible HSDP (Hybrid Sharded Data Parallelism)
679 functionality, enabling efficient training of large models by sharding parameters
680 across multiple devices. The module is automatically enhanced with distributed
681 capabilities including parameter sharding, gradient synchronization, and memory
682 management.
684 When a list of modules is passed, they are treated as one FSDP unit (parameters
685 grouped together).
687 Parameters:
688 module (nn.Module or List[nn.Module]):
689 The module(s) to apply fully_shard to. Modified in-place. When a list
690 is passed, parameters from all modules are grouped as one FSDP unit.
692 mesh (Optional[DeviceMesh], default=None):
693 The device mesh defining the process topology for distributed training.
694 If None, fully_shard creates a default 1D mesh for local parameters.
695 Native DTensor parameters require an explicit DP submesh.
697 reshard_after_forward (bool, default=True):
698 Whether to automatically reshard parameters after forward. When True,
699 parameters are resharded immediately after they are no longer needed,
700 freeing memory for subsequent operations. Set to False if you want to
701 keep parameters unsharded for backward pass or manual control.
703 shard_placement_fn (Callable, default=None):
704 A callable that determines how to shard each parameter. The function
705 should accept a parameter and return a Shard object specifying the
706 sharding dimension, or None to use default sharding (dimension 0)
708 mp_policy (MixedPrecisionPolicy, default=MixedPrecisionPolicy()):
709 Mixed precision training policy controlling data type conversions.
711 offload_policy (OffloadPolicy, default=OffloadPolicy()):
712 Memory offload policy for reducing device memory usage.
714 ignored_params (Optional[set[nn.Parameter]], default=None):
715 Set of parameters to exclude from fully_shard management entirely.
716 These parameters are left on the original module as regular parameters,
717 are not sharded, and do not participate in fully_shard gradient
718 synchronization. Use this for parameters that should remain outside
719 the fully_shard lifecycle.
721 replicate_params (Optional[set[nn.Parameter]], default=None):
722 Set of parameters to keep replicated while still managing them under
723 fully_shard. These parameters are not sharded, but their gradients
724 are still synchronized with DDP-style all-reduce over the current
725 fully_shard communication domain. This differs from ``ignored_params``,
726 which skips fully_shard management and gradient synchronization
727 entirely for the selected parameters.
729 comm_fusion (bool, default=False):
730 Whether enable all_gather fusion and reduce_scatter fusion.
732 comm_fusion_zero_copy (Optional[bool], default=None):
733 Whether allow the experimental zero-copy path for
734 ``comm_fusion``. When set to ``None``, fully_shard enables it
735 automatically whenever ``comm_fusion=True``.
736 When enabled, fully_shard may rebase sharded local parameter storage
737 into one shared flat buffer so fused all-gather can read directly from
738 contiguous memory. This path depends on optimizer compatibility with
739 view-backed parameters.
740 source_shard_infos (Optional[Mapping[nn.Parameter, SourceShardMetaInfo]]):
741 Source TP/EP mesh and placements for the plain-parameter dual mode.
743 Returns:
744 nn.Module or List[nn.Module]: The input module(s) with HSDP capabilities added.
745 """
746 _validate_module_for_fully_shard(module)
748 arg_module = module
749 if isinstance(module, list):
750 modules = tuple(_get_root_modules(module))
751 else:
752 modules = (module,)
754 for mod in modules:
755 _extend_module_with_hsdp_interface(mod)
757 params = _get_modules_parameters(modules, ignored_params)
758 has_dtensor_param = any(isinstance(param, DTensor) for param in params)
759 _validate_managed_params_source_shard_infos(set(params), source_shard_infos)
761 replicate_params = _normalize_replicate_params(replicate_params)
763 if mesh is None:
764 mesh = init_device_mesh(device_type="npu", mesh_shape=(dist.get_world_size(),))
765 if has_dtensor_param:
766 raise ValueError(
767 "fully_shard does not support mesh=None with a native DTensor parameter; "
768 "pass an explicit DP submesh instead."
769 )
770 device = _get_device_from_mesh(mesh)
772 init_modules = modules
773 modules[0].hsdp_init(
774 init_modules,
775 mesh,
776 reshard_after_forward,
777 shard_placement_fn,
778 mp_policy,
779 offload_policy,
780 ignored_params,
781 replicate_params,
782 device,
783 comm_fusion,
784 comm_fusion_zero_copy,
785 source_shard_infos=source_shard_infos,
786 )
787 # Share the same scheduler handle with other roots so mods[i].unshard()/prefetch work
788 if len(modules) > 1:
789 for mod in modules[1:]:
790 mod.hsdp_scheduler = modules[0].hsdp_scheduler
791 return arg_module
794def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
795 """Get the model state dict of a fully_shard-managed model.
797 Args:
798 model: The model whose state dict to retrieve.
799 options: Optional :class:`StateDictOptions` controlling gathering
800 (``full_state_dict``), offloading (``cpu_offload``) and frozen
801 parameter filtering (``ignore_frozen_params``).
803 Returns:
804 The model state dict. By default values are sharded DTensors; when
805 ``full_state_dict=True`` values are full Tensors (CPU on rank 0 only
806 when ``cpu_offload=True``).
807 """
808 return get_model_state_dict_impl(model, options=options)
811def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
812 """Load ``model_state_dict`` into a fully_shard-managed model.
814 Full tensors are scattered into DTensor shards matching the model's layout
815 before the in-place load.
817 Args:
818 model: The model to load state into.
819 model_state_dict: State dict to load. Values may be plain (global)
820 tensors when ``full_state_dict=True`` or sharded DTensors otherwise.
821 options: Optional :class:`StateDictOptions` controlling scattering
822 (``full_state_dict``), device placement (``cpu_offload``) and
823 strictness (``strict``).
825 Returns:
826 None.
827 """
828 return set_model_state_dict_impl(model, model_state_dict, options=options)
831def hsdp_sync_stream():
832 """Wait for hsdp gradient handle to be completed."""
833 wait_grad_handle()