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