Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / fully_shard / api.py: 62%
329 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-26 05:43 +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 return result
156 def hsdp_init(self, module, mesh, reshard_after_forward,
157 shard_placement_fn, mp_policy, offload_policy, ignored_params, replicate_params, device,
158 comm_fusion, comm_fusion_zero_copy: Optional[bool] = None,
159 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]] = None):
160 """init hsdp2 scheduler."""
161 resolved_comm_fusion_zero_copy = _resolve_comm_fusion_zero_copy_default(
162 comm_fusion,
163 comm_fusion_zero_copy,
164 )
166 self.hsdp_scheduler = HSDPSchedulerV2(module,
167 mesh,
168 reshard_after_forward,
169 shard_placement_fn,
170 mp_policy,
171 offload_policy,
172 ignored_params,
173 replicate_params,
174 device,
175 comm_fusion,
176 resolved_comm_fusion_zero_copy,
177 source_shard_infos=source_shard_infos,
178 )
180 def set_requires_gradient_sync(self, requires_grad_sync):
181 r"""
182 set requires grad sync flag.
183 Args:
184 requires_grad_sync(bool): requires_grad_sync is used to control gradient sync process.
185 Raises:
186 ValueError: If `requires_grad_sync` is not bool.
187 """
188 if not isinstance(requires_grad_sync, bool):
189 raise ValueError(f"requires_grad_sync must be bool but got {requires_grad_sync}.")
190 if not hasattr(self, "hsdp_scheduler"):
191 raise ValueError("call hsdp interface first.")
193 for _, module in self.named_modules():
194 if isinstance(module, HSDPModule):
195 module.hsdp_scheduler.set_requires_grad_sync(requires_grad_sync)
197 def zero_grad(self):
198 """zero accumunication grads"""
199 if not hasattr(self, "hsdp_scheduler"):
200 raise ValueError("call hsdp interface first.")
201 return super().zero_grad()
203 def set_modules_to_forward_prefetch(self, modules):
204 """set forward prefetch module list to prefetch all gather for unsharded parameters"""
205 if not isinstance(modules, (tuple, list)):
206 raise ValueError("modules must be HSDPModule list")
207 for module in modules:
208 if not isinstance(module, HSDPModule):
209 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
210 if not hasattr(self, "hsdp_scheduler"):
211 raise ValueError("call hsdp interface first.")
212 self.hsdp_scheduler.set_forward_prefetch_cells(modules)
214 def set_modules_to_backward_prefetch(self, modules):
215 """set backward prefetch module list to prefetch all gather for unsharded parameters"""
216 if not isinstance(modules, (tuple, list)):
217 raise ValueError("modules must be HSDPModule list")
218 for module in modules:
219 if not isinstance(module, HSDPModule):
220 raise ValueError(f"modules must be HSDPModule list but got {type(module)} in list.")
221 if not hasattr(self, "hsdp_scheduler"):
222 raise ValueError("call fully_shard interface first.")
223 self.hsdp_scheduler.set_backward_prefetch_cells(modules)
225 def reshard(self) -> None:
226 """reshard all sharded parameters"""
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.shard()
233 def unshard(self, async_op: bool = False):
234 """unshard all sharded parameters"""
235 if not isinstance(async_op, bool):
236 raise ValueError(f"async_op should be a bool, got {type(async_op)}")
237 if not self.hsdp_scheduler:
238 raise ValueError("hsdp_scheduler is None")
239 hsdp_state = self.hsdp_scheduler.hsdp_state
240 if hsdp_state:
241 hsdp_state.unshard(async_op) # pylint: disable=too-many-function-args
242 if async_op:
243 return _UnshardHandle(hsdp_state=hsdp_state)
244 return None
246 def load_state_dict(
247 self,
248 state_dict: Mapping[str, Any],
249 strict: bool = True,
250 assign: bool = False,
251 ):
252 """
253 Load state dict by copying directly into local shards.
255 Bypasses ``super().load_state_dict()`` because the standard PyTorch
256 implementation triggers ``copy_`` through the DTensor dispatcher, which
257 is not registered in the hyper-parallel layout system.
259 Each value in ``state_dict`` is dispatched by type:
260 - hyper DTensor: extract local shard and copy directly.
261 - plain Tensor whose shape == local shard shape: copy as-is.
262 - plain Tensor whose shape == global shape: distribute via
263 ``distribute_tensor``, then copy the local shard.
265 Args:
266 state_dict (Mapping[str, Any]): Fully-qualified parameter/buffer
267 names mapped to tensors (DTensor or plain Tensor).
268 strict (bool): If ``True`` (default), missing or unexpected keys
269 raise ``RuntimeError``, matching ``nn.Module.load_state_dict``
270 semantics.
271 assign (bool): When ``True`` *and* every value in ``state_dict`` is
272 already a hyper DTensor, defer to the standard
273 ``nn.Module.load_state_dict(assign=True)``, which replaces the
274 module's parameters/buffers with the given DTensors instead of
275 copying into existing storage. This is required when loading
276 sharded DTensors onto a meta-device model (e.g.
277 ``cpu_ram_efficient_loading``). If ``state_dict`` contains any
278 plain tensor (local-shard or global shape), ``assign`` is
279 ignored and the copy/distribute path below is used so the
280 target stays a properly sharded DTensor.
282 Raises:
283 RuntimeError: When ``strict`` is ``True`` and keys do not match.
284 ValueError: When a plain tensor shape matches neither the local
285 shard shape nor the global shape of the target DTensor.
286 """
287 if assign and state_dict and all(
288 isinstance(val, DTensor) for val in state_dict.values()
289 ):
290 return super().load_state_dict(state_dict, strict=strict, assign=True)
291 self_module = cast(nn.Module, self)
293 target_map: dict[str, Tensor] = {}
294 for name, p in parameters_dict(self_module):
295 target_map[name] = p
296 for name, b in self_module.named_buffers():
297 target_map[name] = b
299 if strict:
300 _check_strict_keys(self_module, state_dict)
302 with torch.no_grad():
303 for key, val in state_dict.items():
304 target = target_map.get(key)
305 if target is None:
306 continue
308 if isinstance(target, DTensor):
309 val = _resolve_local_tensor(key, val, target)
310 load_into_param(target, val)
312 # Trigger load_state_dict post-hooks so that HSDP internal
313 # bookkeeping (e.g. _sharded_param_data) stays in sync.
314 # Pass an IncompatibleKeys with the same attribute names as PyTorch
315 # so external hooks can safely read .missing_keys/.unexpected_keys.
316 _IK = namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"])
317 incompatible_keys = _IK([], [])
318 for _, module in self_module.named_modules():
319 hooks = module._load_state_dict_post_hooks # pylint: disable=protected-access
320 for hook in hooks.values():
321 hook(module, incompatible_keys)
323 def set_is_last_backward(self, is_last_backward: bool):
324 """set is_last_backward flag"""
325 self.hsdp_scheduler.scheduler_ctx.is_last_backward = is_last_backward
327 def reset_iter_state(self, recursive: bool = True) -> None:
328 """Reset fully_shard iteration bookkeeping without clearing optimizer gradients.
330 This method must be called only after communication for the iteration has
331 completed. It does not wait in-flight collectives or clear
332 ``sharded_param.grad``/``main_grad``.
334 Args:
335 recursive: Whether to reset every HSDPModule in this module tree.
337 Raises:
338 ValueError: If ``recursive`` is not a bool.
339 """
340 if not isinstance(recursive, bool):
341 raise ValueError(f"recursive should be a bool, got {type(recursive)}")
342 self_module = cast(nn.Module, self)
343 modules = (
344 [module for _, module in self_module.named_modules()]
345 if recursive
346 else [self_module]
347 )
348 for module in modules:
349 if isinstance(module, HSDPModule):
350 module.hsdp_scheduler.reset_iter_state()
352 def set_requires_all_reduce(self, requires_all_reduce: bool, *, recurse: bool = True) -> None:
353 """set requires_all_reduce flag"""
354 if not isinstance(requires_all_reduce, bool):
355 raise ValueError(
356 f"requires_all_reduce should be a bool, got {type(requires_all_reduce)}"
357 )
358 if not recurse:
359 raise NotImplementedError(
360 "Currently impl is equal to recurse=True, "
361 "need support module_param mapping."
362 )
363 self_module = cast(nn.Module, self)
364 for _, module in self_module.named_modules():
365 if isinstance(module, HSDPModule):
366 module.hsdp_scheduler.set_requires_all_reduce(requires_all_reduce)
368 def set_reshard_after_forward(self, reshard_after_forward: bool, recurse: bool = True) -> None:
369 """set reshard_after_forward flag"""
370 if not isinstance(reshard_after_forward, bool):
371 raise ValueError(
372 f"reshard_after_forward should be a bool, got {type(reshard_after_forward)}"
373 )
374 if not recurse:
375 raise NotImplementedError(
376 "Currently impl is equal to recurse=True, "
377 "need support module_param mapping."
378 )
379 self_module = cast(nn.Module, self)
380 for _, module in self_module.named_modules():
381 if isinstance(module, HSDPModule):
382 module.hsdp_scheduler.set_reshard_after_forward(reshard_after_forward)
384 def set_reshard_after_backward(self, reshard_after_backward: bool, recurse: bool = True) -> None:
385 """set reshard_after_backward flag"""
386 if not isinstance(reshard_after_backward, bool):
387 raise ValueError(
388 f"reshard_after_backward should be a bool, got {type(reshard_after_backward)}"
389 )
390 if not recurse:
391 raise NotImplementedError(
392 "Currently impl is equal to recurse=True, "
393 "need support module_param mapping."
394 )
395 self_module = cast(nn.Module, self)
396 for _, module in self_module.named_modules():
397 if isinstance(module, HSDPModule):
398 module.hsdp_scheduler.set_reshard_after_backward(reshard_after_backward)
400 def set_reduce_op_type(self, reduce_op_type, recurse: bool = True) -> None:
401 """
402 set reduce_op_type for all reduce operations in HSDP
403 support reduce_op_type "avg" and "sum", default is "avg"
404 """
405 self_module = cast(nn.Module, self)
406 if recurse:
407 sub_modules = [m for _, m in self_module.named_modules()]
408 else:
409 sub_modules = [self_module]
410 for module in sub_modules:
411 if isinstance(module, HSDPModule):
412 hsdp_state = module.hsdp_scheduler.hsdp_state
413 if hsdp_state:
414 hsdp_state.set_reduce_op_type(reduce_op_type)
416 def set_gradient_scaling_factor(self, factor=None):
417 """
418 Set a multiplicative scaling factor applied to gradients after
419 reduce-scatter / all-reduce and before they are written into
420 ``sharded_param.grad``.
422 ``factor`` may be ``None`` (disable scaling), a Python ``float``/``int``,
423 or a 0-dim/1-element tensor. Setting ``factor`` to ``None`` (the default
424 on construction) skips the scaling op entirely so no extra device-side
425 ``mul_`` is launched on the hot path.
427 Args:
428 factor (None | float | int | torch.Tensor): Scaling coefficient.
429 Use ``None`` to disable scaling.
431 Raises:
432 ValueError: If ``factor`` is not one of the supported types or is a
433 tensor with more than one element.
434 """
435 if factor is not None:
436 if isinstance(factor, bool):
437 raise ValueError(
438 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
439 f"but got bool {factor}."
440 )
441 if isinstance(factor, Tensor):
442 if factor.numel() != 1:
443 raise ValueError(
444 f"gradient_scaling_factor tensor must have exactly 1 element, "
445 f"but got shape {tuple(factor.shape)}."
446 )
447 elif not isinstance(factor, (float, int)):
448 raise ValueError(
449 f"gradient_scaling_factor must be None, float, int or a 1-element Tensor, "
450 f"but got {type(factor).__name__}."
451 )
452 hsdp_state = self.hsdp_scheduler.hsdp_state
453 if hsdp_state:
454 hsdp_state.set_gradient_scaling_factor(factor)
457def _extend_module_with_hsdp_interface(module):
458 """Dynamically extend module's class to inherit from HSDPModule, adding HSDP capabilities."""
459 origin_class = module.__class__
460 extend_class = origin_class_to_extend_class.get(origin_class, None)
461 if extend_class is None:
462 extend_class = type(f"HSDP{origin_class.__name__}", (HSDPModule, origin_class), {})
463 origin_class_to_extend_class[origin_class] = extend_class
464 module.__class__ = extend_class
467def _get_root_modules(modules: List[nn.Module]) -> List[nn.Module]:
468 """
469 Returns the modules in ``modules`` that are root modules (i.e. parent-less)
470 with respect to the set ``modules``. In other words, these are the modules
471 in ``modules`` that are not the child of any other module in ``modules``.
473 Aligned with PyTorch torch.distributed.utils._get_root_modules.
474 """
475 root_modules: List[nn.Module] = []
477 def _get_submodules(mod):
478 return set(mod.modules())
480 module_to_modules: dict[nn.Module, set] = {
481 m: _get_submodules(m) for m in modules
482 }
483 for candidate in modules:
484 is_root = True
485 for mod, submodules in module_to_modules.items():
486 if candidate is not mod and candidate in submodules:
487 is_root = False
488 break
489 if is_root:
490 root_modules.append(candidate)
491 return root_modules
494def _check_module_valid(module):
495 """check module valid"""
496 if not isinstance(module, nn.Module):
497 raise ValueError(f"module's type must be nn.Module but got {type(module)}.")
500def _validate_module_for_fully_shard(
501 module: Union[nn.Module, List[nn.Module]]
502) -> None:
503 """Validate module(s) for fully_shard."""
504 if isinstance(module, list):
505 if len(module) == 0:
506 raise ValueError("fully_shard does not support empty list of modules.")
507 for i, m in enumerate(module):
508 try:
509 _check_module_valid(m)
510 except ValueError:
511 raise ValueError(
512 f"fully_shard expects nn.Module or list[nn.Module], "
513 f"but got list with {type(m).__name__} at index {i}."
514 ) from None
515 else:
516 _check_module_valid(module)
519HsdpValidationOptions = namedtuple(
520 "HsdpValidationOptions",
521 [
522 "shard_size",
523 "threshold",
524 "optimizer_level",
525 "enable_grad_accumulation",
526 "grad_scale",
527 "reduce_dtype",
528 "comm_async",
529 "comm_fusion",
530 "bucket_size",
531 ],
532)
535def _validate_hsdp_shard_size(shard_size: int) -> None:
536 if not isinstance(shard_size, int) or (shard_size <= 0 and shard_size != -1):
537 raise ValueError(f"shard_size must be a positive integer, but got {shard_size}.")
540def _validate_hsdp_threshold(threshold: int) -> None:
541 if not isinstance(threshold, int) or threshold < 0:
542 raise ValueError(f"threshold must be a positive integer or 0, but got {threshold}.")
545def _validate_hsdp_optimizer_level(optimizer_level: str) -> None:
546 if optimizer_level not in ["level1", "level2", "level3"]:
547 raise ValueError(
548 f"Optimizer level should in ['level1', 'level2', 'level3'], but got {optimizer_level}."
549 )
552def _validate_hsdp_reduce_dtype(reduce_dtype) -> None:
553 """Check ``reduce_dtype`` is a torch dtype or ``None``."""
554 if reduce_dtype is not None and not isinstance(reduce_dtype, torch.dtype):
555 raise ValueError(f"reduce_dtype must be torch.dtype but got {reduce_dtype}.")
558def _check_hsdp_input_valid(module, options: HsdpValidationOptions):
559 """check hsdp input valid"""
560 _check_module_valid(module)
561 _validate_hsdp_shard_size(options.shard_size)
562 _validate_hsdp_threshold(options.threshold)
563 _validate_hsdp_optimizer_level(options.optimizer_level)
564 if not isinstance(options.enable_grad_accumulation, bool):
565 raise ValueError(
566 f"enable_grad_accumulation must be bool but got {options.enable_grad_accumulation}."
567 )
568 if not isinstance(options.grad_scale, float):
569 raise ValueError(f"grad_scale must be float but got {options.grad_scale}.")
570 _validate_hsdp_reduce_dtype(options.reduce_dtype)
571 if not isinstance(options.comm_async, bool):
572 raise ValueError(f"comm_async must be bool but got {options.comm_async}.")
573 if not isinstance(options.comm_fusion, bool):
574 raise ValueError(f"comm_fusion must be bool but got {options.comm_fusion}.")
575 if not isinstance(options.bucket_size, int) or (
576 options.bucket_size < 0 and options.bucket_size != -1
577 ):
578 raise ValueError(
579 f"bucket_size must be a positive integer or 0, but got {options.bucket_size}."
580 )
583def _get_device_from_mesh(mesh: DeviceMesh):
584 """Extract and validate the torch device from the device mesh."""
585 device = None
586 device_type = mesh.device_type
587 if device_type not in ("npu", "cuda"):
588 raise AssertionError(
589 f"hyper_parallel.fully_shard support device in [torch.npu, torch.cuda], "
590 f"but got '{device_type}'"
591 )
592 device_handle = get_device_handle(device_type)
593 if device_handle is None:
594 raise ValueError(
595 f"hyper_parallel.fully_shard can't find device_handle of "
596 f"'torch.{device_type}', check the environment."
597 )
598 if device_handle.is_available():
599 device = torch.device(device_handle.current_device())
600 return device
603def _normalize_replicate_params(
604 replicate_params: Optional[set[nn.Parameter]],
605) -> set[nn.Parameter]:
606 """
607 Normalize replicate_params for fully_shard
608 Args:
609 replicate_params (Optional[set[nn.Parameter]]): Set of parameters to exclude from sharding.
610 Returns:
611 set[nn.Parameter]: Set of parameters to exclude from sharding.
612 """
613 if replicate_params is None:
614 return set()
615 out = set(replicate_params)
616 for p in out:
617 if not isinstance(p, (nn.Parameter, DTensor)):
618 raise TypeError(
619 "replicate_params must contain only nn.Parameter or DTensor, "
620 f"got {type(p).__name__}."
621 )
622 return out
625def _get_modules_parameters(modules, ignored_params=None):
626 """Collect deduplicated parameters from module roots."""
627 return get_managed_modules_parameters(modules, ignored_params)
630def _validate_managed_params_source_shard_infos(
631 managed_parameters: set[nn.Parameter],
632 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]],
633) -> None:
634 """Validate the parameter-identity metadata consumed by one fully_shard unit."""
635 if source_shard_infos is None:
636 return
637 if not isinstance(source_shard_infos, Mapping):
638 raise ValueError("source_shard_infos must be a mapping from Parameter to SourceShardMetaInfo")
640 managed_parameters = set(managed_parameters)
642 if any(isinstance(parameter, DTensor) for parameter in managed_parameters):
643 raise ValueError(
644 "source_shard_infos cannot be provided when fully_shard manages a native DTensor parameter"
645 )
647 invalid_values = [
648 parameter
649 for parameter, metadata in source_shard_infos.items()
650 if not isinstance(metadata, SourceShardMetaInfo)
651 ]
652 if invalid_values:
653 raise ValueError("source_shard_infos values must be SourceShardMetaInfo instances")
655 invalid_origins = [
656 parameter
657 for parameter, metadata in source_shard_infos.items()
658 if metadata.origin_is_dtensor and not isinstance(parameter, DTensor)
659 ]
660 if invalid_origins:
661 raise ValueError(
662 "source_shard_infos origin_is_dtensor=True requires a native DTensor parameter"
663 )
665 missing_parameters = managed_parameters.difference(source_shard_infos)
666 if missing_parameters:
667 raise ValueError(
668 "source_shard_infos must cover every parameter managed by this fully_shard call"
669 )
670 unexpected_parameters = set(source_shard_infos).difference(managed_parameters)
671 if unexpected_parameters:
672 raise ValueError(
673 "source_shard_infos contains a parameter not managed by this fully_shard call"
674 )
677def fully_shard(
678 module: Union[nn.Module, List[nn.Module]],
679 *,
680 mesh: Optional[DeviceMesh] = None,
681 reshard_after_forward: bool = True,
682 shard_placement_fn: None = None,
683 mp_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
684 offload_policy: OffloadPolicy = OffloadPolicy(),
685 ignored_params: Optional[set[nn.Parameter]] = None,
686 replicate_params: Optional[set[nn.Parameter]] = None,
687 comm_fusion: bool = False,
688 comm_fusion_zero_copy: Optional[bool] = None,
689 source_shard_infos: Optional[Mapping[nn.Parameter, SourceShardMetaInfo]] = None,
690) -> Union[nn.Module, List[nn.Module]]:
692 """
693 Apply fully_shard to a module (or list of modules) for distributed training with parameter sharding.
695 This interface provides PyTorch-compatible HSDP (Hybrid Sharded Data Parallelism)
696 functionality, enabling efficient training of large models by sharding parameters
697 across multiple devices. The module is automatically enhanced with distributed
698 capabilities including parameter sharding, gradient synchronization, and memory
699 management.
701 When a list of modules is passed, they are treated as one FSDP unit (parameters
702 grouped together).
704 Parameters:
705 module (nn.Module or List[nn.Module]):
706 The module(s) to apply fully_shard to. Modified in-place. When a list
707 is passed, parameters from all modules are grouped as one FSDP unit.
709 mesh (Optional[DeviceMesh], default=None):
710 The device mesh defining the process topology for distributed training.
711 If None, fully_shard creates a default 1D mesh for local parameters.
712 Native DTensor parameters require an explicit DP submesh.
714 reshard_after_forward (bool, default=True):
715 Whether to automatically reshard parameters after forward. When True,
716 parameters are resharded immediately after they are no longer needed,
717 freeing memory for subsequent operations. Set to False if you want to
718 keep parameters unsharded for backward pass or manual control.
720 shard_placement_fn (Callable, default=None):
721 A callable that determines how to shard each parameter. The function
722 should accept a parameter and return a Shard object specifying the
723 sharding dimension, or None to use default sharding (dimension 0)
725 mp_policy (MixedPrecisionPolicy, default=MixedPrecisionPolicy()):
726 Mixed precision training policy controlling data type conversions.
728 offload_policy (OffloadPolicy, default=OffloadPolicy()):
729 Memory offload policy for reducing device memory usage.
731 ignored_params (Optional[set[nn.Parameter]], default=None):
732 Set of parameters to exclude from fully_shard management entirely.
733 These parameters are left on the original module as regular parameters,
734 are not sharded, and do not participate in fully_shard gradient
735 synchronization. Use this for parameters that should remain outside
736 the fully_shard lifecycle.
738 replicate_params (Optional[set[nn.Parameter]], default=None):
739 Set of parameters to keep replicated while still managing them under
740 fully_shard. These parameters are not sharded, but their gradients
741 are still synchronized with DDP-style all-reduce over the current
742 fully_shard communication domain. This differs from ``ignored_params``,
743 which skips fully_shard management and gradient synchronization
744 entirely for the selected parameters.
746 comm_fusion (bool, default=False):
747 Whether enable all_gather fusion and reduce_scatter fusion.
749 comm_fusion_zero_copy (Optional[bool], default=None):
750 Whether allow the experimental zero-copy path for
751 ``comm_fusion``. When set to ``None``, fully_shard enables it
752 automatically whenever ``comm_fusion=True``.
753 When enabled, fully_shard may rebase sharded local parameter storage
754 into one shared flat buffer so fused all-gather can read directly from
755 contiguous memory. This path depends on optimizer compatibility with
756 view-backed parameters.
757 source_shard_infos (Optional[Mapping[nn.Parameter, SourceShardMetaInfo]]):
758 Source TP/EP mesh and placements for the plain-parameter dual mode.
760 Returns:
761 nn.Module or List[nn.Module]: The input module(s) with HSDP capabilities added.
762 """
763 _validate_module_for_fully_shard(module)
765 arg_module = module
766 if isinstance(module, list):
767 modules = tuple(_get_root_modules(module))
768 else:
769 modules = (module,)
771 for mod in modules:
772 _extend_module_with_hsdp_interface(mod)
774 params = _get_modules_parameters(modules, ignored_params)
775 has_dtensor_param = any(isinstance(param, DTensor) for param in params)
776 _validate_managed_params_source_shard_infos(set(params), source_shard_infos)
778 replicate_params = _normalize_replicate_params(replicate_params)
780 if mesh is None:
781 mesh = init_device_mesh(device_type="npu", mesh_shape=(dist.get_world_size(),))
782 if has_dtensor_param:
783 raise ValueError(
784 "fully_shard does not support mesh=None with a native DTensor parameter; "
785 "pass an explicit DP submesh instead."
786 )
787 device = _get_device_from_mesh(mesh)
789 init_modules = modules
790 modules[0].hsdp_init(
791 init_modules,
792 mesh,
793 reshard_after_forward,
794 shard_placement_fn,
795 mp_policy,
796 offload_policy,
797 ignored_params,
798 replicate_params,
799 device,
800 comm_fusion,
801 comm_fusion_zero_copy,
802 source_shard_infos=source_shard_infos,
803 )
804 # Share the same scheduler handle with other roots so mods[i].unshard()/prefetch work
805 if len(modules) > 1:
806 for mod in modules[1:]:
807 mod.hsdp_scheduler = modules[0].hsdp_scheduler
808 return arg_module
811def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
812 """Get the model state dict of a fully_shard-managed model.
814 Args:
815 model: The model whose state dict to retrieve.
816 options: Optional :class:`StateDictOptions` controlling gathering
817 (``full_state_dict``), offloading (``cpu_offload``) and frozen
818 parameter filtering (``ignore_frozen_params``).
820 Returns:
821 The model state dict. By default values are sharded DTensors; when
822 ``full_state_dict=True`` values are full Tensors (CPU on rank 0 only
823 when ``cpu_offload=True``).
824 """
825 return get_model_state_dict_impl(model, options=options)
828def set_model_state_dict(model: Any, model_state_dict: dict[str, Any], *, options: Any = None) -> None:
829 """Load ``model_state_dict`` into a fully_shard-managed model.
831 Full tensors are scattered into DTensor shards matching the model's layout
832 before the in-place load.
834 Args:
835 model: The model to load state into.
836 model_state_dict: State dict to load. Values may be plain (global)
837 tensors when ``full_state_dict=True`` or sharded DTensors otherwise.
838 options: Optional :class:`StateDictOptions` controlling scattering
839 (``full_state_dict``), device placement (``cpu_offload``) and
840 strictness (``strict``).
842 Returns:
843 None.
844 """
845 return set_model_state_dict_impl(model, model_state_dict, options=options)
848def hsdp_sync_stream():
849 """Wait for hsdp gradient handle to be completed."""
850 wait_grad_handle()