Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / utils / clip_grad.py: 70%
434 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"""Distributed-aware gradient clipping for parallel training.
17Communication is driven by each parameter's DTensorSpec (device_mesh +
18placements) rather than any specific parallelism strategy, so a single
19implementation covers FSDP, HSDP, TP+FSDP, and other DTensor-expressed
20parallelisms.
22Collective safety aligned with FSDP1; numerical precision aligned with
23FSDP2's ``_NormPartial`` norm computation path:
25* Gradient norms from sharded parameters are all-reduced across the
26 corresponding shard process group.
27* Non-sharded / replicated norms contribute locally without communication.
28* **All ranks participate in the same collectives** regardless of local
29 gradient availability, preventing collective-misalignment deadlocks.
31The finite p-norm is bit-exact with upstream only for the pure-FSDP,
32single-dtype case; mixed sharded + replicated is mathematically correct
33but intentionally diverges (see ``_total_norm_fsdp2_aligned``).
35Note: PP does not use DTensor layout for gradients today. Cross-stage
36norm aggregation will require an additional manual all-reduce and is
37left for future work.
38"""
39# pylint: disable=forbidden-backend-import,not-callable
40import functools
41import math
42import warnings
43from collections import defaultdict, namedtuple
44from collections.abc import Mapping
45from dataclasses import dataclass
46from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union
48import torch
49import torch.distributed as dist
51from hyper_parallel.core.dtensor.placement_types import Partial
53if TYPE_CHECKING:
54 from hyper_parallel.core.dtensor.layout import Layout
56try:
57 from torch.utils._foreach_utils import (
58 _device_has_foreach_support,
59 _group_tensors_by_device_and_dtype,
60 _has_foreach_support,
61 )
62except ImportError:
63 _device_has_foreach_support = None # type: ignore[assignment]
64 _group_tensors_by_device_and_dtype = None # type: ignore[assignment]
65 _has_foreach_support = None # type: ignore[assignment]
67__all__: list[str] = ["clip_grad_norm_"]
70@dataclass(frozen=True)
71class NormShardPlan:
72 """One norm contributor and its replica ranks for each dense parameter."""
74 representatives: dict[str, bool]
75 replica_ranks: dict[str, tuple[int, ...]]
78def _norm_coverage(param: torch.nn.Parameter) -> tuple[tuple[int, ...], tuple[tuple[int, int], ...]]:
79 """Describe the exact global rectangle represented by a local parameter."""
80 global_shape = tuple(param.shape)
81 layout = getattr(param, "layout", None) if _is_dtensor(param) else None
82 if layout is None:
83 area = tuple((0, size) for size in global_shape)
84 else:
85 from hyper_parallel.core.dtensor.layout import infer_slice_area_by_layout # pylint: disable=C0415
87 rank = dist.get_rank() if dist.is_initialized() else 0
88 if rank not in layout.rank_list:
89 raise ValueError("Norm layout does not contain the current global rank")
90 area = infer_slice_area_by_layout(
91 layout, layout.rank_list.index(rank), global_shape,
92 )
93 local = param._local_tensor if _is_dtensor(param) else param # pylint: disable=W0212
94 if tuple(end - start for start, end in area) != tuple(local.shape):
95 raise ValueError("Norm layout cannot describe the exact local parameter coverage")
96 return global_shape, tuple((int(start), int(end)) for start, end in area)
99def _norm_representative(fqn: str, shape: tuple[int, ...],
100 stage_data: dict[int, dict], rank: int,
101 local_area: tuple) -> tuple[bool, tuple[int, ...]]:
102 """Validate disjoint shard coverage and identify copies of this area."""
103 representatives = {}
104 for member, data in stage_data.items():
105 if data[fqn][0] != shape:
106 raise ValueError(f"Norm shard global shape mismatch: {fqn}")
107 area = data[fqn][1]
108 representatives.setdefault(area, []).append(member)
109 nonempty = [area for area in representatives if all(end > start for start, end in area)]
110 volume = sum(math.prod(end - start for start, end in area) for area in nonempty)
111 if volume != math.prod(shape):
112 raise ValueError(f"Norm shard coverage is incomplete or duplicated: {fqn}")
113 for index, left in enumerate(nonempty):
114 for right in nonempty[index + 1:]:
115 if all(max(a[0], b[0]) < min(a[1], b[1]) for a, b in zip(left, right)):
116 raise ValueError(f"Norm shard rectangles overlap: {fqn}")
117 copies = tuple(representatives[local_area])
118 return rank == min(copies), copies
121def build_norm_shard_plan(model: torch.nn.Module, excluded_params: Iterable[torch.nn.Parameter],
122 mesh_context: Any) -> NormShardPlan:
123 """Choose one rank per identical global shard and verify full coverage.
125 Args:
126 model: Model with final FSDP parameter layouts.
127 excluded_params: Parameters managed by an external optimizer.
128 mesh_context: Stage-local mesh and topology metadata.
129 """
130 excluded_ids = {id(param) for param in excluded_params}
131 local = {
132 fqn: _norm_coverage(param)
133 for fqn, param in model.named_parameters()
134 if param.requires_grad and id(param) not in excluded_ids
135 }
136 if dist.is_initialized():
137 gathered = [None] * dist.get_world_size()
138 dist.all_gather_object(gathered, local)
139 rank = dist.get_rank()
140 else:
141 gathered = [local]
142 rank = 0
143 mesh = getattr(mesh_context, "device_mesh", None)
144 stage_ranks = (tuple(mesh.rank_list) if mesh is not None and hasattr(mesh, "rank_list")
145 else tuple(range(len(gathered))))
146 stage_data = {member: gathered[member] for member in stage_ranks}
147 reference_names = set(local)
148 if any(set(data) != reference_names for data in stage_data.values()):
149 raise ValueError("Norm shard plan has different dense FQNs within one PP stage")
150 representatives = {}
151 replica_ranks = {}
152 for fqn in sorted(reference_names):
153 representatives[fqn], replica_ranks[fqn] = _norm_representative(
154 fqn, local[fqn][0], stage_data, rank, local[fqn][1],
155 )
156 return NormShardPlan(representatives, replica_ranks)
159def _effective_dense_grads(parameters: list[tuple[str, torch.nn.Parameter]]) -> tuple[list, dict]:
160 """Collect local gradients and reduce partial placements once."""
161 infos = []
162 mesh_cache = {}
163 for _, param in parameters:
164 mesh, shard_dims, partial_info = _get_param_mesh_info(param)
165 if mesh is not None:
166 mesh_cache[id(mesh)] = mesh
167 infos.append((param, _get_local_grad(param), mesh, partial_info,
168 (id(mesh) if mesh is not None else None, shard_dims)))
169 return infos, _coalesce_partial_reduce(infos, mesh_cache)
172def _check_dense_grad_presence(present: dict[str, bool], replica_ranks: dict[str, tuple[int, ...]]) -> None:
173 """Reject mismatched gradient presence on copies of the same shard."""
174 if not dist.is_initialized():
175 return
176 gathered = [None] * dist.get_world_size()
177 dist.all_gather_object(gathered, present)
178 for fqn, ranks in replica_ranks.items():
179 copies = [gathered[member][fqn] for member in ranks]
180 if any(value != copies[0] for value in copies[1:]):
181 raise ValueError(f"Norm shard gradient presence differs across replicas: {fqn}")
184def dense_grad_norm_sq_and_refs(model: torch.nn.Module,
185 excluded_params: Iterable[torch.nn.Parameter],
186 norm_shard_plan: NormShardPlan,
187 mesh_context: Any) -> tuple[torch.Tensor, list[torch.Tensor]]:
188 """Sum each dense global shard once and return mutable effective gradients.
190 Args:
191 model: Model with final FSDP parameter layouts.
192 excluded_params: Parameters managed by an external optimizer.
193 norm_shard_plan: Validated coverage and one representative per shard.
194 mesh_context: Stage-local mesh and topology metadata.
195 """
196 excluded_ids = {id(param) for param in excluded_params}
197 parameters = [
198 (fqn, param) for fqn, param in model.named_parameters()
199 if param.requires_grad and id(param) not in excluded_ids
200 ]
201 device = next(((_param_device(param)) for _, param in parameters), None)
202 if device is None:
203 mesh = getattr(mesh_context, "device_mesh", None)
204 device = torch.device(getattr(mesh, "device_type", "cpu"))
205 infos, reduced = _effective_dense_grads(parameters)
206 refs = []
207 local_sq = torch.zeros((), device=device, dtype=torch.float32)
208 present = {}
209 for index, (fqn, _) in enumerate(parameters):
210 grad = infos[index][1]
211 present[fqn] = grad is not None
212 if grad is None:
213 continue
214 refs.append(grad)
215 if norm_shard_plan.representatives[fqn]:
216 view = reduced.get(index, grad)
217 local_sq.add_(view.float().square().sum().to(device))
218 _check_dense_grad_presence(present, norm_shard_plan.replica_ranks)
219 for name in getattr(getattr(mesh_context, "device_mesh", None), "mesh_dim_names", ()) or ():
220 axis = mesh_context.device_mesh[name]
221 if axis.size() > 1:
222 dist.all_reduce(local_sq, op=dist.ReduceOp.SUM, group=axis.get_group())
223 return local_sq, refs
226def scale_grad_refs_(refs: Iterable[torch.Tensor], coefficient: torch.Tensor) -> None:
227 """Scale the exact dense gradients that the optimizer will consume.
229 Args:
230 refs: Mutable dense effective gradients.
231 coefficient: Shared clipping coefficient.
232 """
233 for grad in refs:
234 grad.mul_(coefficient)
237# (id(mesh) or None, shard_dims) -> list of local grads for norm computation
238_GradGroupKey = Tuple[Optional[int], Tuple[int, ...]]
240# (mesh_dim_index, dist.ReduceOp, needs_manual_avg)
241_PartialReduceInfo = Tuple[int, "dist.ReduceOp", bool]
243# Result of _build_grad_groups; tuple-unpacking compatible with prior 7-tuple.
244_GradGroups = namedtuple(
245 "_GradGroups",
246 "grad_groups all_grads norm_grads key_per_grad mesh_cache device "
247 "has_dtensor_grad",
248)
251def _is_dtensor(value: object) -> bool:
252 """Return whether *value* is a HyperParallel DTensor.
254 The DTensor module imports ``core.utils`` during initialization, so keep
255 this import lazy to avoid a package initialization cycle.
256 """
257 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415
259 return isinstance(value, DTensor)
262# ---------------------------------------------------------------------------
263# Reduce-op mapping
264# ---------------------------------------------------------------------------
266_REDUCE_OP_AVG_SUPPORTED = hasattr(dist.ReduceOp, "AVG")
268_STR_TO_REDUCE_OP: Dict[str, "dist.ReduceOp"] = {
269 "sum": dist.ReduceOp.SUM,
270 "max": dist.ReduceOp.MAX,
271 "min": dist.ReduceOp.MIN,
272}
273if _REDUCE_OP_AVG_SUPPORTED:
274 _STR_TO_REDUCE_OP["avg"] = dist.ReduceOp.AVG
277def _str_to_reduce_op(op_str: str) -> Tuple["dist.ReduceOp", bool]:
278 """Map a ``Partial`` placement's *reduce_op* string to ``dist.ReduceOp``.
280 Returns ``(reduce_op, needs_manual_avg)`` where *needs_manual_avg*
281 is ``True`` when ``"avg"`` is requested but the backend does not
282 support ``dist.ReduceOp.AVG`` — the caller should use SUM and
283 manually divide by the group size.
284 """
285 lower = op_str.lower()
286 if lower == "avg" and not _REDUCE_OP_AVG_SUPPORTED:
287 return dist.ReduceOp.SUM, True
288 op = _STR_TO_REDUCE_OP.get(lower)
289 if op is None:
290 raise ValueError(
291 f"Unsupported Partial reduce_op: {op_str!r}. "
292 f"Supported: {sorted(set(list(_STR_TO_REDUCE_OP) + ['avg']))}"
293 )
294 return op, False
297# ---------------------------------------------------------------------------
298# Helpers
299# ---------------------------------------------------------------------------
301def _normalize_parameters(
302 parameters: Union["torch.nn.Module", torch.Tensor, Iterable[torch.Tensor]],
303) -> List[torch.Tensor]:
304 """Normalize *parameters* to a flat list of tensors.
306 * ``torch.nn.Module`` -> ``list(module.parameters())``
307 * single ``torch.Tensor`` -> ``[tensor]``
308 * iterable of tensors -> ``list(iterable)``
309 """
310 if isinstance(parameters, torch.nn.Module):
311 return list(parameters.parameters())
312 if isinstance(parameters, torch.Tensor):
313 return [parameters]
314 return list(parameters)
317def _param_device(param: torch.Tensor) -> torch.device:
318 """Return the local device of *param* (unwrap DTensor if needed)."""
319 if _is_dtensor(param):
320 return param._local_tensor.device # pylint: disable=protected-access
321 return param.device
324def _get_grad_obj(param: torch.nn.Parameter) -> Optional[torch.Tensor]:
325 """Return the gradient object for *param*.
327 Checks ``param.main_grad`` first (used when
328 ``MixedPrecisionPolicy.apply_grad_on_fp32_main_grad=True``),
329 falling back to ``param.grad``.
330 """
331 grad = getattr(param, "main_grad", None)
332 if grad is not None:
333 return grad
334 return param.grad
337def _get_local_grad(param: torch.nn.Parameter) -> Optional[torch.Tensor]:
338 """Return the local gradient tensor, or ``None`` if absent.
340 Supports ``main_grad`` for fp32 mixed-precision training.
341 """
342 if not param.requires_grad:
343 return None
344 grad = _get_grad_obj(param)
345 if grad is None:
346 return None
347 if _is_dtensor(grad):
348 return grad._local_tensor # pylint: disable=protected-access
349 return grad
352def _get_param_mesh_info(
353 param: torch.nn.Parameter,
354 parameter_layout: Optional["Layout"] = None,
355) -> Tuple[
356 Optional[object],
357 Tuple[int, ...],
358 Tuple[_PartialReduceInfo, ...],
359]:
360 """Derive DeviceMesh, Shard dims and Partial info from DTensorSpec.
362 Checks the *gradient's* spec first; falls back to the *parameter's*
363 spec when the gradient is a plain tensor on a DTensor parameter
364 (common after FSDP/HSDP backward where ``param.grad`` is stored as
365 the local shard tensor).
367 Returns ``(mesh, shard_dims, partial_info)`` where *partial_info*
368 is a tuple of ``(mesh_dim, dist.ReduceOp, needs_manual_avg)``
369 triples that respect the ``Partial`` placement's ``reduce_op``
370 attribute. *needs_manual_avg* is ``True`` when ``"avg"`` was
371 requested but the backend lacks ``dist.ReduceOp.AVG`` support.
372 """
373 grad = _get_grad_obj(param)
374 # Prefer grad's spec (most accurate); fall back to param's.
375 spec_source = grad if _is_dtensor(grad) else param
376 if parameter_layout is None:
377 if not _is_dtensor(spec_source):
378 return None, (), ()
379 placements = spec_source.placements
380 mesh = spec_source.device_mesh
381 else:
382 layout = spec_source.layout if _is_dtensor(spec_source) else None
383 if layout is not None and (
384 layout.mesh_shape, layout.rank_list, layout.alias_name, layout.placements, layout.tensor_shape
385 ) != (
386 parameter_layout.mesh_shape, parameter_layout.rank_list, parameter_layout.alias_name,
387 parameter_layout.placements, parameter_layout.tensor_shape
388 ):
389 raise ValueError("Explicit parameter layout conflicts with its DTensor layout")
390 layout = layout if layout is not None else parameter_layout
391 placements = layout.placements
392 mesh = layout.mesh
394 shard_dims = tuple(
395 i for i, p in enumerate(placements)
396 if p.is_shard()
397 )
398 partial_info = tuple(
399 (i, *_str_to_reduce_op(p.reduce_op))
400 for i, p in enumerate(placements)
401 if isinstance(p, Partial)
402 )
403 return mesh, shard_dims, partial_info
406def _sum_p_norms(
407 dev_grads: List[torch.Tensor],
408 norm_type: float,
409 device: torch.device,
410 total: torch.Tensor,
411) -> None:
412 """Accumulate sum-of-p-th-powers for *dev_grads* into *total*."""
413 for g in dev_grads:
414 n = torch.linalg.vector_norm(g, norm_type)
415 total.add_(n.to(device=device) ** norm_type)
418def _foreach_p_norms(
419 grads: List[torch.Tensor],
420 norm_type: float,
421 device: torch.device,
422) -> torch.Tensor:
423 """Fast path: fuse per-tensor norms via ``_foreach_norm``.
425 Restricted to float32 tensors to preserve the same numerical
426 precision as ``vector_norm(dtype=float32)``. Non-float32 tensors
427 and backends that raise ``RuntimeError`` fall back to per-tensor
428 ``vector_norm``.
429 """
430 total = torch.tensor(0.0, device=device, dtype=torch.float32)
431 grouped = _group_tensors_by_device_and_dtype(
432 [[g.detach() for g in grads]],
433 )
434 for (dev, _), ([dev_grads], _) in grouped.items():
435 if (
436 dev_grads[0].dtype == torch.float32
437 and _has_foreach_support(dev_grads, dev)
438 ):
439 try:
440 per_norms = torch._foreach_norm( # pylint: disable=W0212
441 dev_grads, norm_type,
442 )
443 except RuntimeError:
444 per_norms = None
445 if per_norms is not None:
446 total.add_(
447 torch.stack([
448 n.to(device=device) ** norm_type
449 for n in per_norms
450 ]).sum(),
451 )
452 else:
453 _sum_p_norms(dev_grads, norm_type, device, total)
454 else:
455 _sum_p_norms(dev_grads, norm_type, device, total)
456 return total
459def _per_tensor_norms(
460 grads: List[torch.Tensor],
461 norm_type: float,
462 device: torch.device,
463) -> List[torch.Tensor]:
464 """Return per-tensor norms as a list of scalar tensors on *device*."""
465 if not grads:
466 return []
468 if _group_tensors_by_device_and_dtype is None or not hasattr(torch, "_foreach_norm"):
469 return [
470 torch.linalg.vector_norm(g.detach(), norm_type).to(device=device)
471 for g in grads
472 ]
474 norms: List[torch.Tensor] = []
475 grouped = _group_tensors_by_device_and_dtype(
476 [[g.detach() for g in grads]],
477 )
478 for (dev, _), ([dev_grads], _) in grouped.items():
479 if dev_grads and _has_foreach_support(dev_grads, dev):
480 try:
481 per_norms = torch._foreach_norm( # pylint: disable=W0212
482 dev_grads, norm_type,
483 )
484 except RuntimeError:
485 per_norms = None
486 if per_norms is not None:
487 norms.extend(
488 [n.to(device=device) for n in per_norms],
489 )
490 continue
491 norms.extend([
492 torch.linalg.vector_norm(g, norm_type).to(device=device)
493 for g in dev_grads
494 ])
495 return norms
498def _compute_local_norm( # pylint: disable=R0911
499 grads: List[torch.Tensor],
500 norm_type: float,
501 device: torch.device,
502) -> torch.Tensor:
503 """Compute the combined norm of *grads* locally in FP32.
505 When *grads* is empty, returns the **identity element** for the
506 subsequent all-reduce so that this rank contributes a neutral value
507 (aligned with FSDP1's ``_zero_scalar`` approach):
509 * ``inf`` -> 0 (neutral for MAX; norms are non-negative)
510 * ``-inf`` -> +inf (neutral for MIN)
511 * ``0`` -> 0 (neutral for SUM)
512 * finite -> 0 (neutral for SUM)
513 """
514 if not grads:
515 if norm_type == -math.inf:
516 return torch.tensor(
517 float("inf"), device=device, dtype=torch.float32,
518 )
519 return torch.tensor(0.0, device=device, dtype=torch.float32)
521 if norm_type == math.inf:
522 norms = [
523 torch.linalg.vector_norm(g.detach(), math.inf)
524 for g in grads
525 ]
526 return torch.stack(norms).max().to(device)
528 if norm_type == -math.inf:
529 norms = [
530 torch.linalg.vector_norm(g.detach(), -math.inf)
531 for g in grads
532 ]
533 return torch.stack(norms).min().to(device)
535 if norm_type == 0:
536 norms = [
537 torch.linalg.vector_norm(g.detach(), 0)
538 for g in grads
539 ]
540 return torch.stack(norms).sum().to(device)
542 # Finite p-norm: return sum of p-th powers.
543 if (
544 len(grads) > 1
545 and _group_tensors_by_device_and_dtype is not None
546 and hasattr(torch, "_foreach_norm")
547 ):
548 return _foreach_p_norms(grads, norm_type, device)
550 # Scalar fallback when foreach utilities are unavailable.
551 norms = [
552 torch.linalg.vector_norm(g.detach(), norm_type)
553 for g in grads
554 ]
555 norm_powers = [n.to(device=device) ** norm_type for n in norms]
556 return torch.stack(norm_powers).sum()
559# ---------------------------------------------------------------------------
560# Total norm aggregation with collectives
561# ---------------------------------------------------------------------------
563def _get_total_norm(
564 grad_groups: Dict[_GradGroupKey, List[torch.Tensor]],
565 norm_type: float,
566 mesh_cache: Dict[int, object],
567 device: torch.device,
568 norm_grads: List[torch.Tensor],
569 key_per_grad: List[_GradGroupKey],
570) -> torch.Tensor:
571 """Compute total gradient norm with per-group all-reduce.
573 ``norm_grads`` (parallel to ``key_per_grad``) holds the tensor whose
574 norm to take per parameter; only the finite p-norm path consumes it.
575 """
576 if norm_type == math.inf:
577 return _total_norm_inf(
578 grad_groups, norm_type, mesh_cache, device,
579 dist.ReduceOp.MAX,
580 )
582 if norm_type == -math.inf:
583 return _total_norm_inf(
584 grad_groups, norm_type, mesh_cache, device,
585 dist.ReduceOp.MIN,
586 )
588 if norm_type == 0:
589 return _total_norm_sum(
590 grad_groups, norm_type, mesh_cache, device,
591 )
593 # Finite p-norm: FSDP2-aligned sequence.
594 total_p = _total_norm_fsdp2_aligned(
595 grad_groups, norm_type, mesh_cache, device,
596 norm_grads, key_per_grad,
597 )
598 return total_p ** (1.0 / norm_type)
601def _total_norm_inf( # pylint: disable=R0913,R0917
602 grad_groups, norm_type, mesh_cache, device, reduce_op,
603):
604 """Shared logic for inf / -inf norms."""
605 group_norms: List[torch.Tensor] = []
606 for (mesh_id, shard_dims), grads in grad_groups.items():
607 local_norm = _compute_local_norm(grads, norm_type, device)
608 if mesh_id is not None:
609 mesh = mesh_cache[mesh_id]
610 for dim in shard_dims:
611 dist.all_reduce(
612 local_norm, op=reduce_op,
613 group=mesh.get_group(dim),
614 )
615 group_norms.append(local_norm)
616 if not group_norms:
617 if norm_type == -math.inf:
618 return torch.tensor(float("inf"), device=device)
619 return torch.tensor(0.0, device=device)
620 stacked = torch.stack(group_norms)
621 return stacked.max() if reduce_op == dist.ReduceOp.MAX else stacked.min()
624def _total_norm_sum(grad_groups, norm_type, mesh_cache, device):
625 """Shared logic for finite norms and L0 (all use SUM all-reduce)."""
626 total = torch.tensor(0.0, device=device)
627 for (mesh_id, shard_dims), grads in grad_groups.items():
628 local_val = _compute_local_norm(grads, norm_type, device)
629 if mesh_id is not None:
630 mesh = mesh_cache[mesh_id]
631 for dim in shard_dims:
632 dist.all_reduce(
633 local_val, op=dist.ReduceOp.SUM,
634 group=mesh.get_group(dim),
635 )
636 total.add_(local_val)
637 return total
640def _reduction_signature(grad_groups, mesh_cache):
641 """Bucket grad-group keys by the *process group(s)* they reduce over.
643 Two keys that reduce over the same set of process groups (same global
644 ranks per shard dim) must be **pooled** so their norms accumulate in
645 one stack -- matching FSDP2's single foreach-norm + single reduce and
646 keeping the loss bit-exact even when params live on several distinct
647 ``DeviceMesh`` objects that share the same DP process group (common
648 for multi-component models). Keys that reduce over *different*
649 process groups (TP+FSDP heterogeneous sharding, expert parallel) get
650 separate buckets, each reduced over its own group.
652 Returns ``(key_to_sig, sig_groups, sig_order)``:
654 * ``key_to_sig`` -- ``key -> signature`` (hashable; ``()`` = replicate
655 / no communication).
656 * ``sig_groups`` -- ``signature -> list[ProcessGroup]`` to all-reduce.
657 * ``sig_order`` -- signatures in first-seen (parameter) order, so all
658 ranks issue the same collectives in the same order.
659 """
660 key_to_sig: Dict[_GradGroupKey, Tuple] = {}
661 sig_groups: Dict[Tuple, List[object]] = {}
662 sig_order: List[Tuple] = []
663 for mesh_id, shard_dims in grad_groups:
664 if shard_dims and mesh_id is not None:
665 mesh = mesh_cache[mesh_id]
666 groups = [mesh.get_group(dim) for dim in shard_dims]
667 sig = tuple(
668 tuple(dist.get_process_group_ranks(group)) for group in groups
669 )
670 else:
671 groups = []
672 sig = ()
673 key_to_sig[(mesh_id, shard_dims)] = sig
674 if sig not in sig_groups:
675 sig_groups[sig] = groups
676 sig_order.append(sig)
677 return key_to_sig, sig_groups, sig_order
680def _total_norm_fsdp2_aligned(grad_groups, norm_type, mesh_cache, device,
681 norm_grads, key_per_grad):
682 """FSDP2-aligned norm for finite p-norms.
684 Grads are bucketed by the *process group* they reduce over (see
685 :func:`_reduction_signature`); ``norm_grads`` supplies, in global
686 parameter order, the Partial-reduced (already-global) view for Partial
687 grads and the raw local grad otherwise. Each bucket does ONE ``stack``
688 → ONE ``vector_norm`` → ``^p`` → ONE ``all_reduce SUM`` per shard dim,
689 and the buckets are summed locally.
691 Reduction rules:
693 * Sharded bucket -- reduce over its own group(s). Same-group params
694 across distinct ``DeviceMesh`` objects pool into one reduce; distinct
695 groups (TP+FSDP heterogeneous, expert parallel) stay separate -- the
696 per-group convention also used by :func:`_total_norm_inf` /
697 :func:`_total_norm_sum` and by VeOmni / torchtitan.
698 * Replicate bucket (signature ``()``) -- contribute norm² locally with
699 NO communication; reducing it over the shard group would over-count
700 it ``shard_world_size`` times (the FSDP1 convention).
702 Empty-but-present buckets still issue their all_reduce (identity ``0``
703 for SUM) so all ranks run the same collectives.
705 Bit-exact with upstream ``torch.nn.utils.clip_grad_norm_`` only for the
706 pure-FSDP, single-bucket, single-dtype case. Mixed shard + replicate is
707 mathematically correct but intentionally diverges (upstream folds the
708 replicate norm into one ``_NormPartial`` reduce and over-counts it);
709 mixed-dtype follows upstream's :func:`_per_tensor_norms` device/dtype
710 regrouping rather than strict global order.
712 Returns the global sum of p-th powers (caller takes p-th root).
713 """
714 key_to_sig, sig_groups, sig_order = _reduction_signature(
715 grad_groups, mesh_cache,
716 )
718 # Bucket norms by reduction signature, preserving global parameter order.
719 sig_grads: Dict[Tuple, List[torch.Tensor]] = {sig: [] for sig in sig_order}
720 for grad, key in zip(norm_grads, key_per_grad):
721 sig_grads[key_to_sig[key]].append(grad)
723 total_p = torch.tensor(0.0, device=device, dtype=torch.float32)
724 for sig in sig_order:
725 grads = sig_grads[sig]
726 if grads:
727 norms = _per_tensor_norms(grads, norm_type, device)
728 local_p = torch.linalg.vector_norm(
729 torch.stack(norms).to(torch.float32), norm_type,
730 ) ** norm_type
731 else:
732 local_p = torch.tensor(0.0, device=device, dtype=torch.float32)
734 for group in sig_groups[sig]:
735 dist.all_reduce(local_p, op=dist.ReduceOp.SUM, group=group)
737 total_p = total_p + local_p
739 return total_p
742def _build_coalesce_buffer(
743 param_infos: List[Tuple],
744 indices: List[int],
745) -> Tuple[List[torch.Tensor], List[int], List[bool], List[int]]:
746 """Build flat fp32 chunks for one coalesce group.
748 Returns ``(chunks, chunk_sizes, has_grad, active_indices)``.
749 Frozen params are skipped; trainable grad-free params contribute
750 zeros so the collective matches ranks that have a grad.
751 """
752 chunks: List[torch.Tensor] = []
753 chunk_sizes: List[int] = []
754 has_grad: List[bool] = []
755 active_indices: List[int] = []
757 for idx in indices:
758 param = param_infos[idx][0]
759 local_grad = param_infos[idx][1]
760 if local_grad is not None:
761 chunks.append(
762 local_grad.detach().reshape(-1).to(torch.float32),
763 )
764 chunk_sizes.append(local_grad.numel())
765 has_grad.append(True)
766 active_indices.append(idx)
767 elif param.requires_grad:
768 local_p = (
769 param._local_tensor # pylint: disable=W0212
770 if _is_dtensor(param) else param.data
771 )
772 numel = local_p.numel()
773 chunks.append(
774 torch.zeros(
775 numel, device=local_p.device,
776 dtype=torch.float32,
777 ),
778 )
779 chunk_sizes.append(numel)
780 has_grad.append(False)
781 active_indices.append(idx)
783 return chunks, chunk_sizes, has_grad, active_indices
786def _coalesce_partial_reduce( # pylint: disable=R0914
787 param_infos: List[Tuple],
788 mesh_cache: Dict[int, object],
789) -> Dict[int, torch.Tensor]:
790 """Coalesce Partial all-reduces: O(N) collectives → O(G).
792 Groups parameters sharing the same ``(mesh, partial_info)`` and
793 flattens their gradients (or zeros for trainable grad-free params)
794 into a single fp32 buffer. **One** ``all_reduce`` per buffer
795 replaces the previous per-parameter collective calls.
797 For TP+FSDP (all params share the same mesh / placements), this
798 turns ~200 individual all-reduces into 1 — saving 10-20 ms per
799 training step at typical HCCS/NCCL latencies.
801 Frozen params (``requires_grad=False``) are consistently grad-free
802 across all ranks and are excluded from the buffer to avoid wasting
803 bandwidth.
805 All buffers use float32 to guarantee dtype consistency across ranks
806 in mixed-precision training (grad may be fp16/bf16 while param is
807 fp32).
809 Returns a dict mapping *param_infos* index → reduced gradient view
810 (1-D fp32 slice of the coalesced buffer). Only entries for params
811 with actual gradients are included.
812 """
813 # Group by Partial coalesce key: (mesh_id, partial_info)
814 coalesce_groups: Dict[
815 Tuple, List[int],
816 ] = defaultdict(list)
817 for idx, info in enumerate(param_infos):
818 mesh, partial_info = info[2], info[3]
819 if partial_info:
820 if mesh is None:
821 raise RuntimeError(
822 "clip_grad_norm_: parameter has Partial placements "
823 "but no DeviceMesh. This is a DTensor invariant "
824 "violation."
825 )
826 pck = (id(mesh), partial_info)
827 coalesce_groups[pck].append(idx)
829 reduced: Dict[int, torch.Tensor] = {}
831 for (mesh_id, partial_info), indices in coalesce_groups.items():
832 mesh = mesh_cache[mesh_id]
833 chunks, chunk_sizes, has_grad, active_indices = (
834 _build_coalesce_buffer(param_infos, indices)
835 )
837 if not chunks:
838 continue # all params frozen, no collective needed
840 # Sanity check: same mesh → same device. Fail fast on
841 # misconfigured inputs rather than silent NCCL errors.
842 buf_device = chunks[0].device
843 for chunk in chunks[1:]:
844 if chunk.device != buf_device:
845 raise RuntimeError(
846 f"clip_grad_norm_: parameters in the same Partial "
847 f"coalesce group are on different devices "
848 f"({buf_device} vs {chunk.device}). All parameters "
849 f"sharing the same DeviceMesh must reside on the "
850 f"same local device."
851 )
853 buf = torch.cat(chunks)
855 for pdim, reduce_op, needs_avg in partial_info:
856 group = mesh.get_group(pdim)
857 dist.all_reduce(buf, op=reduce_op, group=group)
858 if needs_avg:
859 buf /= dist.get_world_size(group=group)
861 # Extract views for params with actual gradients.
862 offset = 0
863 for i, idx in enumerate(active_indices):
864 numel = chunk_sizes[i]
865 if has_grad[i]:
866 reduced[idx] = buf[offset:offset + numel]
867 offset += numel
869 return reduced
872def _build_grad_groups( # pylint: disable=R0914
873 params: List[torch.Tensor],
874 parameter_layouts: Optional[Mapping[torch.nn.Parameter, "Layout"]] = None,
875) -> Tuple[
876 Dict[_GradGroupKey, List[torch.Tensor]],
877 List[torch.Tensor],
878 List[torch.Tensor],
879 List[_GradGroupKey],
880 Dict[int, object],
881 torch.device,
882 bool,
883]:
884 """Classify parameters into grad groups and pre-reduce Partial grads.
886 Group structure is derived from *parameter* DTensorSpecs (always
887 present on every rank) rather than gradients (which may be ``None``
888 on some ranks). This ensures every rank enters the same set of
889 collectives, preventing deadlocks (aligned with FSDP1 where all
890 ranks unconditionally execute the same all-reduce path).
892 Partial gradients are reduced via a **coalesced** all-reduce
893 (see ``_coalesce_partial_reduce``), turning O(N) per-parameter
894 collectives into O(G) where G is the number of distinct
895 ``(mesh, partial_info)`` groups (typically 1 for TP+FSDP).
897 Returns
898 ``(grad_groups, all_grads, norm_grads, key_per_grad, mesh_cache,
899 device, has_dtensor_grad)``. ``grad_groups`` maps each
900 ``(mesh_id, shard_dims)`` key to its grads; ``all_grads`` is the
901 flat list of raw local grads (global parameter order) scaled
902 in-place by the clip step; ``norm_grads`` is parallel to
903 ``all_grads`` but holds the Partial-reduced view for Partial grads
904 (the value whose norm is taken on the finite-p path);
905 ``key_per_grad`` is parallel to both and maps each grad back to its
906 group key.
907 """
908 # --- Phase 1: classify all parameters ---
909 param_infos: List[Tuple] = []
910 mesh_cache: Dict[int, object] = {}
911 device: Optional[torch.device] = None
913 for param in params:
914 if parameter_layouts is None:
915 mesh, shard_dims, partial_info = _get_param_mesh_info(param)
916 else:
917 layout = parameter_layouts.get(param)
918 if layout is None and not _is_dtensor(param) and not _is_dtensor(_get_grad_obj(param)):
919 raise ValueError("Explicit layouts must cover every local parameter participating in clipping")
920 mesh, shard_dims, partial_info = _get_param_mesh_info(param, layout)
921 key: _GradGroupKey = (
922 id(mesh) if mesh is not None else None, shard_dims,
923 )
924 if mesh is not None:
925 mesh_cache[id(mesh)] = mesh
926 if device is None:
927 device = _param_device(param)
928 local_grad = _get_local_grad(param)
929 param_infos.append(
930 (param, local_grad, mesh, partial_info, key),
931 )
933 if device is None:
934 device = torch.device("cpu")
936 # --- Phase 2: coalesced Partial reduction (O(N) → O(G)) ---
937 reduced = _coalesce_partial_reduce(param_infos, mesh_cache)
939 # --- Phase 3: build grad_groups ---
940 grad_groups: Dict[_GradGroupKey, List[torch.Tensor]] = defaultdict(
941 list,
942 )
943 all_grads: List[torch.Tensor] = []
944 # norm_grads / key_per_grad are parallel to all_grads (see Returns): the
945 # norm-input view (Partial-reduced where coalesced, else raw local) and
946 # the group key per grad, kept in global parameter order.
947 norm_grads: List[torch.Tensor] = []
948 key_per_grad: List[_GradGroupKey] = []
949 has_dtensor_grad = False
951 for idx, info in enumerate(param_infos):
952 param, local_grad, key = info[0], info[1], info[4]
953 if local_grad is None:
954 # Ensure the key exists so the Shard norm all-reduce is
955 # entered even when this rank has no grads for the group.
956 if key not in grad_groups:
957 grad_groups[key] = []
958 continue
960 grad_obj = _get_grad_obj(param)
961 if _is_dtensor(grad_obj):
962 has_dtensor_grad = True
963 all_grads.append(local_grad)
964 key_per_grad.append(key)
965 if idx in reduced:
966 grad_groups[key].append(reduced[idx])
967 norm_grads.append(reduced[idx])
968 else:
969 grad_groups[key].append(local_grad)
970 norm_grads.append(local_grad)
972 return _GradGroups(
973 grad_groups, all_grads, norm_grads, key_per_grad,
974 mesh_cache, device, has_dtensor_grad,
975 )
978def _clip_grads_with_norm_(
979 all_grads: List[torch.Tensor],
980 max_norm: float,
981 total_norm: torch.Tensor,
982 foreach: Optional[bool] = None,
983) -> None:
984 """Scale gradients in-place so the total norm <= *max_norm*."""
985 clip_coef = max_norm / (total_norm + 1e-6)
986 clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
988 if _group_tensors_by_device_and_dtype is not None:
989 grouped_grads = _group_tensors_by_device_and_dtype(
990 [all_grads],
991 )
992 for (device, dtype), ([device_grads], _) in grouped_grads.items():
993 use_foreach = (
994 foreach is None and _has_foreach_support(device_grads, device)
995 ) or (foreach and _device_has_foreach_support(device))
996 if use_foreach:
997 torch._foreach_mul_( # pylint: disable=W0212
998 device_grads,
999 clip_coef_clamped.to(device=device, dtype=dtype),
1000 )
1001 elif foreach:
1002 raise RuntimeError(
1003 f"foreach=True was passed, but can't use the "
1004 f"foreach API on {device.type} tensors"
1005 )
1006 else:
1007 clip_coef_clamped_cast = clip_coef_clamped.to(device=device, dtype=dtype)
1008 for g in device_grads:
1009 g.mul_(clip_coef_clamped_cast)
1010 else:
1011 # Fallback when _foreach_utils is unavailable.
1012 if foreach:
1013 raise RuntimeError(
1014 "foreach=True was passed, but "
1015 "torch.utils._foreach_utils is not available"
1016 )
1017 for grad in all_grads:
1018 grad.mul_(clip_coef_clamped.to(grad.device, grad.dtype))
1021# ---------------------------------------------------------------------------
1022# Public API
1023# ---------------------------------------------------------------------------
1025@torch.no_grad()
1026def clip_grad_norm_(
1027 parameters: Union[
1028 "torch.nn.Module", torch.Tensor, Iterable[torch.Tensor],
1029 ],
1030 max_norm: float,
1031 norm_type: float = 2.0,
1032 error_if_nonfinite: bool = False,
1033 foreach: Optional[bool] = None,
1034 *,
1035 parameter_layouts: Optional[Mapping[torch.nn.Parameter, "Layout"]] = None,
1036) -> torch.Tensor:
1037 """Compute and clip gradient norm for distributed models.
1039 Drop-in replacement for the standard ``clip_grad_norm_`` that
1040 correctly handles DTensor-sharded parameters by deriving
1041 communication from each parameter's DTensorSpec.
1043 .. warning:: This function uses collective communications. It
1044 **must be called on all ranks** to avoid deadlocks. Aligned
1045 with FSDP1: every rank participates in the same collectives
1046 regardless of local gradient availability.
1048 Communication is derived from each parameter's DTensorSpec:
1050 * ``Shard`` on mesh dim *d* -- all-reduce norm statistics
1051 across ``device_mesh.get_group(d)``
1052 * ``Partial`` on mesh dim *d* -- all-reduce gradient values
1053 using the placement's ``reduce_op`` before norm computation
1054 * ``Replicate`` / plain tensor -- no communication
1056 This covers FSDP, HSDP, TP+FSDP, and any combination expressible
1057 via DTensor placements. PP cross-stage norm aggregation is not
1058 yet handled (requires manual all-reduce across stages).
1060 Args:
1061 parameters: An ``nn.Module``, a single ``Tensor``, or an
1062 iterable of ``Tensor`` s whose gradients to clip.
1063 max_norm: Maximum allowed gradient norm.
1064 norm_type: Type of the norm (default ``2.0``).
1065 error_if_nonfinite: If ``True``, raise a ``RuntimeError``
1066 when the total norm is non-finite. Default ``False``.
1067 foreach: Use the faster foreach-based implementation for the
1068 gradient clipping step. If ``None``, use the foreach
1069 implementation for devices that support it and silently
1070 fall back to the per-tensor implementation for others.
1071 Default ``None``.
1072 parameter_layouts: Optional parameter-identity mapping to final HyperParallel
1073 Layouts for local tensors. DTensor gradient/parameter layouts take
1074 precedence and must agree with explicit metadata. Local trainable
1075 parameters must be fully covered; frozen parameters are ignored.
1077 Returns:
1078 The total (unclipped) gradient norm as a scalar tensor,
1079 cast to the promoted dtype of all gradient tensors.
1080 """
1081 max_norm = float(max_norm)
1082 norm_type = float(norm_type)
1084 params = _normalize_parameters(parameters)
1085 if parameter_layouts is not None:
1086 params = [param for param in dict.fromkeys(params) if param.requires_grad]
1087 if foreach is None:
1088 foreach = False
1090 (
1091 grad_groups, all_grads, norm_grads, key_per_grad,
1092 mesh_cache, device, has_dtensor_grad,
1093 ) = (_build_grad_groups(params) if parameter_layouts is None
1094 else _build_grad_groups(params, parameter_layouts))
1096 # -- Norm + clip (all ranks participate) --------------------------------
1097 # _compute_local_norm returns identity elements for empty groups,
1098 # so the subsequent all-reduce is safe and semantically neutral.
1099 total_norm = _get_total_norm(
1100 grad_groups, norm_type, mesh_cache, device,
1101 norm_grads, key_per_grad,
1102 )
1104 if error_if_nonfinite and torch.logical_or(
1105 total_norm.isnan(), total_norm.isinf()
1106 ):
1107 raise RuntimeError(
1108 f"The total norm of order {norm_type} for gradients from "
1109 "`parameters` is non-finite, so it cannot be clipped. To "
1110 "disable this error and scale the gradients by the "
1111 "non-finite norm anyway, set "
1112 "`error_if_nonfinite=False`"
1113 )
1115 if all_grads:
1116 # Disable foreach for dtensor-backed grads to avoid dispatch issues.
1117 effective_foreach = False if has_dtensor_grad and foreach is None else foreach
1118 _clip_grads_with_norm_(
1119 all_grads, max_norm, total_norm, effective_foreach,
1120 )
1122 # Promote return dtype to match gradient dtypes (FSDP1 convention).
1123 # When this rank has no gradients, return in the default FP32 dtype
1124 # (same as FSDP1's behavior to avoid extra communication).
1125 if not all_grads:
1126 warnings.warn(
1127 "clip_grad_norm_ called on this rank with no gradients -- "
1128 "returning the local norm in the default dtype "
1129 f"{total_norm.dtype}",
1130 stacklevel=2,
1131 )
1132 return total_norm
1134 total_norm_dtype = functools.reduce(
1135 torch.promote_types,
1136 [g.dtype for g in all_grads],
1137 )
1138 # Return global all-reduced norm, consistent with torchtitan's
1139 # full_tensor() approach — .item() returns the correct global value.
1140 return total_norm.to(total_norm_dtype)