Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / clip_grad.py: 67%
310 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""
39import functools
40import math
41import warnings
42from collections import defaultdict, namedtuple
43from typing import Dict, Iterable, List, Optional, Tuple, Union
45import torch
46import torch.distributed as dist
48from hyper_parallel.core.dtensor.dtensor import DTensor
49from hyper_parallel.core.dtensor.placement_types import Partial
51try:
52 from torch.utils._foreach_utils import (
53 _device_has_foreach_support,
54 _group_tensors_by_device_and_dtype,
55 _has_foreach_support,
56 )
57except ImportError:
58 _device_has_foreach_support = None # type: ignore[assignment]
59 _group_tensors_by_device_and_dtype = None # type: ignore[assignment]
60 _has_foreach_support = None # type: ignore[assignment]
62__all__: list[str] = ["clip_grad_norm_"]
65# (id(mesh) or None, shard_dims) -> list of local grads for norm computation
66_GradGroupKey = Tuple[Optional[int], Tuple[int, ...]]
68# (mesh_dim_index, dist.ReduceOp, needs_manual_avg)
69_PartialReduceInfo = Tuple[int, "dist.ReduceOp", bool]
71# Result of _build_grad_groups; tuple-unpacking compatible with prior 7-tuple.
72_GradGroups = namedtuple(
73 "_GradGroups",
74 "grad_groups all_grads norm_grads key_per_grad mesh_cache device "
75 "has_dtensor_grad",
76)
79# ---------------------------------------------------------------------------
80# Reduce-op mapping
81# ---------------------------------------------------------------------------
83_REDUCE_OP_AVG_SUPPORTED = hasattr(dist.ReduceOp, "AVG")
85_STR_TO_REDUCE_OP: Dict[str, "dist.ReduceOp"] = {
86 "sum": dist.ReduceOp.SUM,
87 "max": dist.ReduceOp.MAX,
88 "min": dist.ReduceOp.MIN,
89}
90if _REDUCE_OP_AVG_SUPPORTED:
91 _STR_TO_REDUCE_OP["avg"] = dist.ReduceOp.AVG
94def _str_to_reduce_op(op_str: str) -> Tuple["dist.ReduceOp", bool]:
95 """Map a ``Partial`` placement's *reduce_op* string to ``dist.ReduceOp``.
97 Returns ``(reduce_op, needs_manual_avg)`` where *needs_manual_avg*
98 is ``True`` when ``"avg"`` is requested but the backend does not
99 support ``dist.ReduceOp.AVG`` — the caller should use SUM and
100 manually divide by the group size.
101 """
102 lower = op_str.lower()
103 if lower == "avg" and not _REDUCE_OP_AVG_SUPPORTED:
104 return dist.ReduceOp.SUM, True
105 op = _STR_TO_REDUCE_OP.get(lower)
106 if op is None:
107 raise ValueError(
108 f"Unsupported Partial reduce_op: {op_str!r}. "
109 f"Supported: {sorted(set(list(_STR_TO_REDUCE_OP) + ['avg']))}"
110 )
111 return op, False
114# ---------------------------------------------------------------------------
115# Helpers
116# ---------------------------------------------------------------------------
118def _normalize_parameters(
119 parameters: Union["torch.nn.Module", torch.Tensor, Iterable[torch.Tensor]],
120) -> List[torch.Tensor]:
121 """Normalize *parameters* to a flat list of tensors.
123 * ``torch.nn.Module`` -> ``list(module.parameters())``
124 * single ``torch.Tensor`` -> ``[tensor]``
125 * iterable of tensors -> ``list(iterable)``
126 """
127 if isinstance(parameters, torch.nn.Module):
128 return list(parameters.parameters())
129 if isinstance(parameters, torch.Tensor):
130 return [parameters]
131 return list(parameters)
134def _param_device(param: torch.Tensor) -> torch.device:
135 """Return the local device of *param* (unwrap DTensor if needed)."""
136 if isinstance(param, DTensor):
137 return param._local_tensor.device # pylint: disable=protected-access
138 return param.device
141def _get_grad_obj(param: torch.nn.Parameter) -> Optional[torch.Tensor]:
142 """Return the gradient object for *param*.
144 Checks ``param.main_grad`` first (used when
145 ``MixedPrecisionPolicy.apply_grad_on_fp32_main_grad=True``),
146 falling back to ``param.grad``.
147 """
148 grad = getattr(param, "main_grad", None)
149 if grad is not None:
150 return grad
151 return param.grad
154def _get_local_grad(param: torch.nn.Parameter) -> Optional[torch.Tensor]:
155 """Return the local gradient tensor, or ``None`` if absent.
157 Supports ``main_grad`` for fp32 mixed-precision training.
158 """
159 if not param.requires_grad:
160 return None
161 grad = _get_grad_obj(param)
162 if grad is None:
163 return None
164 if isinstance(grad, DTensor):
165 return grad._local_tensor # pylint: disable=protected-access
166 return grad
169def _get_param_mesh_info(
170 param: torch.nn.Parameter,
171) -> Tuple[
172 Optional[object],
173 Tuple[int, ...],
174 Tuple[_PartialReduceInfo, ...],
175]:
176 """Derive DeviceMesh, Shard dims and Partial info from DTensorSpec.
178 Checks the *gradient's* spec first; falls back to the *parameter's*
179 spec when the gradient is a plain tensor on a DTensor parameter
180 (common after FSDP/HSDP backward where ``param.grad`` is stored as
181 the local shard tensor).
183 Returns ``(mesh, shard_dims, partial_info)`` where *partial_info*
184 is a tuple of ``(mesh_dim, dist.ReduceOp, needs_manual_avg)``
185 triples that respect the ``Partial`` placement's ``reduce_op``
186 attribute. *needs_manual_avg* is ``True`` when ``"avg"`` was
187 requested but the backend lacks ``dist.ReduceOp.AVG`` support.
188 """
189 grad = _get_grad_obj(param)
190 # Prefer grad's spec (most accurate); fall back to param's.
191 spec_source = grad if isinstance(grad, DTensor) else param
192 if not isinstance(spec_source, DTensor):
193 return None, (), ()
195 shard_dims = tuple(
196 i for i, p in enumerate(spec_source.placements)
197 if p.is_shard()
198 )
199 partial_info = tuple(
200 (i, *_str_to_reduce_op(p.reduce_op))
201 for i, p in enumerate(spec_source.placements)
202 if isinstance(p, Partial)
203 )
204 return spec_source.device_mesh, shard_dims, partial_info
207def _sum_p_norms(
208 dev_grads: List[torch.Tensor],
209 norm_type: float,
210 device: torch.device,
211 total: torch.Tensor,
212) -> None:
213 """Accumulate sum-of-p-th-powers for *dev_grads* into *total*."""
214 for g in dev_grads:
215 n = torch.linalg.vector_norm(g, norm_type)
216 total.add_(n.to(device=device) ** norm_type)
219def _foreach_p_norms(
220 grads: List[torch.Tensor],
221 norm_type: float,
222 device: torch.device,
223) -> torch.Tensor:
224 """Fast path: fuse per-tensor norms via ``_foreach_norm``.
226 Restricted to float32 tensors to preserve the same numerical
227 precision as ``vector_norm(dtype=float32)``. Non-float32 tensors
228 and backends that raise ``RuntimeError`` fall back to per-tensor
229 ``vector_norm``.
230 """
231 total = torch.tensor(0.0, device=device, dtype=torch.float32)
232 grouped = _group_tensors_by_device_and_dtype(
233 [[g.detach() for g in grads]],
234 )
235 for (dev, _), ([dev_grads], _) in grouped.items():
236 if (
237 dev_grads[0].dtype == torch.float32
238 and _has_foreach_support(dev_grads, dev)
239 ):
240 try:
241 per_norms = torch._foreach_norm( # pylint: disable=W0212
242 dev_grads, norm_type,
243 )
244 except RuntimeError:
245 per_norms = None
246 if per_norms is not None:
247 total.add_(
248 torch.stack([
249 n.to(device=device) ** norm_type
250 for n in per_norms
251 ]).sum(),
252 )
253 else:
254 _sum_p_norms(dev_grads, norm_type, device, total)
255 else:
256 _sum_p_norms(dev_grads, norm_type, device, total)
257 return total
260def _per_tensor_norms(
261 grads: List[torch.Tensor],
262 norm_type: float,
263 device: torch.device,
264) -> List[torch.Tensor]:
265 """Return per-tensor norms as a list of scalar tensors on *device*."""
266 if not grads:
267 return []
269 if _group_tensors_by_device_and_dtype is None or not hasattr(torch, "_foreach_norm"):
270 return [
271 torch.linalg.vector_norm(g.detach(), norm_type).to(device=device)
272 for g in grads
273 ]
275 norms: List[torch.Tensor] = []
276 grouped = _group_tensors_by_device_and_dtype(
277 [[g.detach() for g in grads]],
278 )
279 for (dev, _), ([dev_grads], _) in grouped.items():
280 if dev_grads and _has_foreach_support(dev_grads, dev):
281 try:
282 per_norms = torch._foreach_norm( # pylint: disable=W0212
283 dev_grads, norm_type,
284 )
285 except RuntimeError:
286 per_norms = None
287 if per_norms is not None:
288 norms.extend(
289 [n.to(device=device) for n in per_norms],
290 )
291 continue
292 norms.extend([
293 torch.linalg.vector_norm(g, norm_type).to(device=device)
294 for g in dev_grads
295 ])
296 return norms
299def _compute_local_norm( # pylint: disable=R0911
300 grads: List[torch.Tensor],
301 norm_type: float,
302 device: torch.device,
303) -> torch.Tensor:
304 """Compute the combined norm of *grads* locally in FP32.
306 When *grads* is empty, returns the **identity element** for the
307 subsequent all-reduce so that this rank contributes a neutral value
308 (aligned with FSDP1's ``_zero_scalar`` approach):
310 * ``inf`` -> 0 (neutral for MAX; norms are non-negative)
311 * ``-inf`` -> +inf (neutral for MIN)
312 * ``0`` -> 0 (neutral for SUM)
313 * finite -> 0 (neutral for SUM)
314 """
315 if not grads:
316 if norm_type == -math.inf:
317 return torch.tensor(
318 float("inf"), device=device, dtype=torch.float32,
319 )
320 return torch.tensor(0.0, device=device, dtype=torch.float32)
322 if norm_type == math.inf:
323 norms = [
324 torch.linalg.vector_norm(g.detach(), math.inf)
325 for g in grads
326 ]
327 return torch.stack(norms).max().to(device)
329 if norm_type == -math.inf:
330 norms = [
331 torch.linalg.vector_norm(g.detach(), -math.inf)
332 for g in grads
333 ]
334 return torch.stack(norms).min().to(device)
336 if norm_type == 0:
337 norms = [
338 torch.linalg.vector_norm(g.detach(), 0)
339 for g in grads
340 ]
341 return torch.stack(norms).sum().to(device)
343 # Finite p-norm: return sum of p-th powers.
344 if (
345 len(grads) > 1
346 and _group_tensors_by_device_and_dtype is not None
347 and hasattr(torch, "_foreach_norm")
348 ):
349 return _foreach_p_norms(grads, norm_type, device)
351 # Scalar fallback when foreach utilities are unavailable.
352 norms = [
353 torch.linalg.vector_norm(g.detach(), norm_type)
354 for g in grads
355 ]
356 norm_powers = [n.to(device=device) ** norm_type for n in norms]
357 return torch.stack(norm_powers).sum()
360# ---------------------------------------------------------------------------
361# Total norm aggregation with collectives
362# ---------------------------------------------------------------------------
364def _get_total_norm(
365 grad_groups: Dict[_GradGroupKey, List[torch.Tensor]],
366 norm_type: float,
367 mesh_cache: Dict[int, object],
368 device: torch.device,
369 norm_grads: List[torch.Tensor],
370 key_per_grad: List[_GradGroupKey],
371) -> torch.Tensor:
372 """Compute total gradient norm with per-group all-reduce.
374 ``norm_grads`` (parallel to ``key_per_grad``) holds the tensor whose
375 norm to take per parameter; only the finite p-norm path consumes it.
376 """
377 if norm_type == math.inf:
378 return _total_norm_inf(
379 grad_groups, norm_type, mesh_cache, device,
380 dist.ReduceOp.MAX,
381 )
383 if norm_type == -math.inf:
384 return _total_norm_inf(
385 grad_groups, norm_type, mesh_cache, device,
386 dist.ReduceOp.MIN,
387 )
389 if norm_type == 0:
390 return _total_norm_sum(
391 grad_groups, norm_type, mesh_cache, device,
392 )
394 # Finite p-norm: FSDP2-aligned sequence.
395 total_p = _total_norm_fsdp2_aligned(
396 grad_groups, norm_type, mesh_cache, device,
397 norm_grads, key_per_grad,
398 )
399 return total_p ** (1.0 / norm_type)
402def _total_norm_inf( # pylint: disable=R0913,R0917
403 grad_groups, norm_type, mesh_cache, device, reduce_op,
404):
405 """Shared logic for inf / -inf norms."""
406 group_norms: List[torch.Tensor] = []
407 for (mesh_id, shard_dims), grads in grad_groups.items():
408 local_norm = _compute_local_norm(grads, norm_type, device)
409 if mesh_id is not None:
410 mesh = mesh_cache[mesh_id]
411 for dim in shard_dims:
412 dist.all_reduce(
413 local_norm, op=reduce_op,
414 group=mesh.get_group(dim),
415 )
416 group_norms.append(local_norm)
417 if not group_norms:
418 if norm_type == -math.inf:
419 return torch.tensor(float("inf"), device=device)
420 return torch.tensor(0.0, device=device)
421 stacked = torch.stack(group_norms)
422 return stacked.max() if reduce_op == dist.ReduceOp.MAX else stacked.min()
425def _total_norm_sum(grad_groups, norm_type, mesh_cache, device):
426 """Shared logic for finite norms and L0 (all use SUM all-reduce)."""
427 total = torch.tensor(0.0, device=device)
428 for (mesh_id, shard_dims), grads in grad_groups.items():
429 local_val = _compute_local_norm(grads, norm_type, device)
430 if mesh_id is not None:
431 mesh = mesh_cache[mesh_id]
432 for dim in shard_dims:
433 dist.all_reduce(
434 local_val, op=dist.ReduceOp.SUM,
435 group=mesh.get_group(dim),
436 )
437 total.add_(local_val)
438 return total
441def _reduction_signature(grad_groups, mesh_cache):
442 """Bucket grad-group keys by the *process group(s)* they reduce over.
444 Two keys that reduce over the same set of process groups (same global
445 ranks per shard dim) must be **pooled** so their norms accumulate in
446 one stack -- matching FSDP2's single foreach-norm + single reduce and
447 keeping the loss bit-exact even when params live on several distinct
448 ``DeviceMesh`` objects that share the same DP process group (common
449 for multi-component models). Keys that reduce over *different*
450 process groups (TP+FSDP heterogeneous sharding, expert parallel) get
451 separate buckets, each reduced over its own group.
453 Returns ``(key_to_sig, sig_groups, sig_order)``:
455 * ``key_to_sig`` -- ``key -> signature`` (hashable; ``()`` = replicate
456 / no communication).
457 * ``sig_groups`` -- ``signature -> list[ProcessGroup]`` to all-reduce.
458 * ``sig_order`` -- signatures in first-seen (parameter) order, so all
459 ranks issue the same collectives in the same order.
460 """
461 key_to_sig: Dict[_GradGroupKey, Tuple] = {}
462 sig_groups: Dict[Tuple, List[object]] = {}
463 sig_order: List[Tuple] = []
464 for mesh_id, shard_dims in grad_groups:
465 if shard_dims and mesh_id is not None:
466 mesh = mesh_cache[mesh_id]
467 groups = [mesh.get_group(dim) for dim in shard_dims]
468 sig = tuple(
469 tuple(dist.get_process_group_ranks(group)) for group in groups
470 )
471 else:
472 groups = []
473 sig = ()
474 key_to_sig[(mesh_id, shard_dims)] = sig
475 if sig not in sig_groups:
476 sig_groups[sig] = groups
477 sig_order.append(sig)
478 return key_to_sig, sig_groups, sig_order
481def _total_norm_fsdp2_aligned(grad_groups, norm_type, mesh_cache, device,
482 norm_grads, key_per_grad):
483 """FSDP2-aligned norm for finite p-norms.
485 Grads are bucketed by the *process group* they reduce over (see
486 :func:`_reduction_signature`); ``norm_grads`` supplies, in global
487 parameter order, the Partial-reduced (already-global) view for Partial
488 grads and the raw local grad otherwise. Each bucket does ONE ``stack``
489 → ONE ``vector_norm`` → ``^p`` → ONE ``all_reduce SUM`` per shard dim,
490 and the buckets are summed locally.
492 Reduction rules:
494 * Sharded bucket -- reduce over its own group(s). Same-group params
495 across distinct ``DeviceMesh`` objects pool into one reduce; distinct
496 groups (TP+FSDP heterogeneous, expert parallel) stay separate -- the
497 per-group convention also used by :func:`_total_norm_inf` /
498 :func:`_total_norm_sum` and by VeOmni / torchtitan.
499 * Replicate bucket (signature ``()``) -- contribute norm² locally with
500 NO communication; reducing it over the shard group would over-count
501 it ``shard_world_size`` times (the FSDP1 convention).
503 Empty-but-present buckets still issue their all_reduce (identity ``0``
504 for SUM) so all ranks run the same collectives.
506 Bit-exact with upstream ``torch.nn.utils.clip_grad_norm_`` only for the
507 pure-FSDP, single-bucket, single-dtype case. Mixed shard + replicate is
508 mathematically correct but intentionally diverges (upstream folds the
509 replicate norm into one ``_NormPartial`` reduce and over-counts it);
510 mixed-dtype follows upstream's :func:`_per_tensor_norms` device/dtype
511 regrouping rather than strict global order.
513 Returns the global sum of p-th powers (caller takes p-th root).
514 """
515 key_to_sig, sig_groups, sig_order = _reduction_signature(
516 grad_groups, mesh_cache,
517 )
519 # Bucket norms by reduction signature, preserving global parameter order.
520 sig_grads: Dict[Tuple, List[torch.Tensor]] = {sig: [] for sig in sig_order}
521 for grad, key in zip(norm_grads, key_per_grad):
522 sig_grads[key_to_sig[key]].append(grad)
524 total_p = torch.tensor(0.0, device=device, dtype=torch.float32)
525 for sig in sig_order:
526 grads = sig_grads[sig]
527 if grads:
528 norms = _per_tensor_norms(grads, norm_type, device)
529 local_p = torch.linalg.vector_norm(
530 torch.stack(norms).to(torch.float32), norm_type,
531 ) ** norm_type
532 else:
533 local_p = torch.tensor(0.0, device=device, dtype=torch.float32)
535 for group in sig_groups[sig]:
536 dist.all_reduce(local_p, op=dist.ReduceOp.SUM, group=group)
538 total_p = total_p + local_p
540 return total_p
543def _build_coalesce_buffer(
544 param_infos: List[Tuple],
545 indices: List[int],
546) -> Tuple[List[torch.Tensor], List[int], List[bool], List[int]]:
547 """Build flat fp32 chunks for one coalesce group.
549 Returns ``(chunks, chunk_sizes, has_grad, active_indices)``.
550 Frozen params are skipped; trainable grad-free params contribute
551 zeros so the collective matches ranks that have a grad.
552 """
553 chunks: List[torch.Tensor] = []
554 chunk_sizes: List[int] = []
555 has_grad: List[bool] = []
556 active_indices: List[int] = []
558 for idx in indices:
559 param = param_infos[idx][0]
560 local_grad = param_infos[idx][1]
561 if local_grad is not None:
562 chunks.append(
563 local_grad.detach().reshape(-1).to(torch.float32),
564 )
565 chunk_sizes.append(local_grad.numel())
566 has_grad.append(True)
567 active_indices.append(idx)
568 elif param.requires_grad:
569 local_p = (
570 param._local_tensor # pylint: disable=W0212
571 if isinstance(param, DTensor) else param.data
572 )
573 numel = local_p.numel()
574 chunks.append(
575 torch.zeros(
576 numel, device=local_p.device,
577 dtype=torch.float32,
578 ),
579 )
580 chunk_sizes.append(numel)
581 has_grad.append(False)
582 active_indices.append(idx)
584 return chunks, chunk_sizes, has_grad, active_indices
587def _coalesce_partial_reduce( # pylint: disable=R0914
588 param_infos: List[Tuple],
589 mesh_cache: Dict[int, object],
590) -> Dict[int, torch.Tensor]:
591 """Coalesce Partial all-reduces: O(N) collectives → O(G).
593 Groups parameters sharing the same ``(mesh, partial_info)`` and
594 flattens their gradients (or zeros for trainable grad-free params)
595 into a single fp32 buffer. **One** ``all_reduce`` per buffer
596 replaces the previous per-parameter collective calls.
598 For TP+FSDP (all params share the same mesh / placements), this
599 turns ~200 individual all-reduces into 1 — saving 10-20 ms per
600 training step at typical HCCS/NCCL latencies.
602 Frozen params (``requires_grad=False``) are consistently grad-free
603 across all ranks and are excluded from the buffer to avoid wasting
604 bandwidth.
606 All buffers use float32 to guarantee dtype consistency across ranks
607 in mixed-precision training (grad may be fp16/bf16 while param is
608 fp32).
610 Returns a dict mapping *param_infos* index → reduced gradient view
611 (1-D fp32 slice of the coalesced buffer). Only entries for params
612 with actual gradients are included.
613 """
614 # Group by Partial coalesce key: (mesh_id, partial_info)
615 coalesce_groups: Dict[
616 Tuple, List[int],
617 ] = defaultdict(list)
618 for idx, info in enumerate(param_infos):
619 mesh, partial_info = info[2], info[3]
620 if partial_info:
621 if mesh is None:
622 raise RuntimeError(
623 "clip_grad_norm_: parameter has Partial placements "
624 "but no DeviceMesh. This is a DTensor invariant "
625 "violation."
626 )
627 pck = (id(mesh), partial_info)
628 coalesce_groups[pck].append(idx)
630 reduced: Dict[int, torch.Tensor] = {}
632 for (mesh_id, partial_info), indices in coalesce_groups.items():
633 mesh = mesh_cache[mesh_id]
634 chunks, chunk_sizes, has_grad, active_indices = (
635 _build_coalesce_buffer(param_infos, indices)
636 )
638 if not chunks:
639 continue # all params frozen, no collective needed
641 # Sanity check: same mesh → same device. Fail fast on
642 # misconfigured inputs rather than silent NCCL errors.
643 buf_device = chunks[0].device
644 for chunk in chunks[1:]:
645 if chunk.device != buf_device:
646 raise RuntimeError(
647 f"clip_grad_norm_: parameters in the same Partial "
648 f"coalesce group are on different devices "
649 f"({buf_device} vs {chunk.device}). All parameters "
650 f"sharing the same DeviceMesh must reside on the "
651 f"same local device."
652 )
654 buf = torch.cat(chunks)
656 for pdim, reduce_op, needs_avg in partial_info:
657 group = mesh.get_group(pdim)
658 dist.all_reduce(buf, op=reduce_op, group=group)
659 if needs_avg:
660 buf /= dist.get_world_size(group=group)
662 # Extract views for params with actual gradients.
663 offset = 0
664 for i, idx in enumerate(active_indices):
665 numel = chunk_sizes[i]
666 if has_grad[i]:
667 reduced[idx] = buf[offset:offset + numel]
668 offset += numel
670 return reduced
673def _build_grad_groups( # pylint: disable=R0914
674 params: List[torch.Tensor],
675) -> Tuple[
676 Dict[_GradGroupKey, List[torch.Tensor]],
677 List[torch.Tensor],
678 List[torch.Tensor],
679 List[_GradGroupKey],
680 Dict[int, object],
681 torch.device,
682 bool,
683]:
684 """Classify parameters into grad groups and pre-reduce Partial grads.
686 Group structure is derived from *parameter* DTensorSpecs (always
687 present on every rank) rather than gradients (which may be ``None``
688 on some ranks). This ensures every rank enters the same set of
689 collectives, preventing deadlocks (aligned with FSDP1 where all
690 ranks unconditionally execute the same all-reduce path).
692 Partial gradients are reduced via a **coalesced** all-reduce
693 (see ``_coalesce_partial_reduce``), turning O(N) per-parameter
694 collectives into O(G) where G is the number of distinct
695 ``(mesh, partial_info)`` groups (typically 1 for TP+FSDP).
697 Returns
698 ``(grad_groups, all_grads, norm_grads, key_per_grad, mesh_cache,
699 device, has_dtensor_grad)``. ``grad_groups`` maps each
700 ``(mesh_id, shard_dims)`` key to its grads; ``all_grads`` is the
701 flat list of raw local grads (global parameter order) scaled
702 in-place by the clip step; ``norm_grads`` is parallel to
703 ``all_grads`` but holds the Partial-reduced view for Partial grads
704 (the value whose norm is taken on the finite-p path);
705 ``key_per_grad`` is parallel to both and maps each grad back to its
706 group key.
707 """
708 # --- Phase 1: classify all parameters ---
709 param_infos: List[Tuple] = []
710 mesh_cache: Dict[int, object] = {}
711 device: Optional[torch.device] = None
713 for param in params:
714 mesh, shard_dims, partial_info = _get_param_mesh_info(param)
715 key: _GradGroupKey = (
716 id(mesh) if mesh is not None else None, shard_dims,
717 )
718 if mesh is not None:
719 mesh_cache[id(mesh)] = mesh
720 if device is None:
721 device = _param_device(param)
722 local_grad = _get_local_grad(param)
723 param_infos.append(
724 (param, local_grad, mesh, partial_info, key),
725 )
727 if device is None:
728 device = torch.device("cpu")
730 # --- Phase 2: coalesced Partial reduction (O(N) → O(G)) ---
731 reduced = _coalesce_partial_reduce(param_infos, mesh_cache)
733 # --- Phase 3: build grad_groups ---
734 grad_groups: Dict[_GradGroupKey, List[torch.Tensor]] = defaultdict(
735 list,
736 )
737 all_grads: List[torch.Tensor] = []
738 # norm_grads / key_per_grad are parallel to all_grads (see Returns): the
739 # norm-input view (Partial-reduced where coalesced, else raw local) and
740 # the group key per grad, kept in global parameter order.
741 norm_grads: List[torch.Tensor] = []
742 key_per_grad: List[_GradGroupKey] = []
743 has_dtensor_grad = False
745 for idx, info in enumerate(param_infos):
746 param, local_grad, key = info[0], info[1], info[4]
747 if local_grad is None:
748 # Ensure the key exists so the Shard norm all-reduce is
749 # entered even when this rank has no grads for the group.
750 if key not in grad_groups:
751 grad_groups[key] = []
752 continue
754 grad_obj = _get_grad_obj(param)
755 if isinstance(grad_obj, DTensor):
756 has_dtensor_grad = True
757 all_grads.append(local_grad)
758 key_per_grad.append(key)
759 if idx in reduced:
760 grad_groups[key].append(reduced[idx])
761 norm_grads.append(reduced[idx])
762 else:
763 grad_groups[key].append(local_grad)
764 norm_grads.append(local_grad)
766 return _GradGroups(
767 grad_groups, all_grads, norm_grads, key_per_grad,
768 mesh_cache, device, has_dtensor_grad,
769 )
772def _clip_grads_with_norm_(
773 all_grads: List[torch.Tensor],
774 max_norm: float,
775 total_norm: torch.Tensor,
776 foreach: Optional[bool] = None,
777) -> None:
778 """Scale gradients in-place so the total norm <= *max_norm*."""
779 clip_coef = max_norm / (total_norm + 1e-6)
780 clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
782 if _group_tensors_by_device_and_dtype is not None:
783 grouped_grads = _group_tensors_by_device_and_dtype(
784 [all_grads],
785 )
786 for (device, dtype), ([device_grads], _) in grouped_grads.items():
787 use_foreach = (
788 foreach is None and _has_foreach_support(device_grads, device)
789 ) or (foreach and _device_has_foreach_support(device))
790 if use_foreach:
791 torch._foreach_mul_( # pylint: disable=W0212
792 device_grads,
793 clip_coef_clamped.to(device=device, dtype=dtype),
794 )
795 elif foreach:
796 raise RuntimeError(
797 f"foreach=True was passed, but can't use the "
798 f"foreach API on {device.type} tensors"
799 )
800 else:
801 clip_coef_clamped_cast = clip_coef_clamped.to(device=device, dtype=dtype)
802 for g in device_grads:
803 g.mul_(clip_coef_clamped_cast)
804 else:
805 # Fallback when _foreach_utils is unavailable.
806 if foreach:
807 raise RuntimeError(
808 "foreach=True was passed, but "
809 "torch.utils._foreach_utils is not available"
810 )
811 for grad in all_grads:
812 grad.mul_(clip_coef_clamped.to(grad.device, grad.dtype))
815# ---------------------------------------------------------------------------
816# Public API
817# ---------------------------------------------------------------------------
819@torch.no_grad()
820def clip_grad_norm_(
821 parameters: Union[
822 "torch.nn.Module", torch.Tensor, Iterable[torch.Tensor],
823 ],
824 max_norm: float,
825 norm_type: float = 2.0,
826 error_if_nonfinite: bool = False,
827 foreach: Optional[bool] = None,
828) -> torch.Tensor:
829 """Compute and clip gradient norm for distributed models.
831 Drop-in replacement for the standard ``clip_grad_norm_`` that
832 correctly handles DTensor-sharded parameters by deriving
833 communication from each parameter's DTensorSpec.
835 .. warning:: This function uses collective communications. It
836 **must be called on all ranks** to avoid deadlocks. Aligned
837 with FSDP1: every rank participates in the same collectives
838 regardless of local gradient availability.
840 Communication is derived from each parameter's DTensorSpec:
842 * ``Shard`` on mesh dim *d* -- all-reduce norm statistics
843 across ``device_mesh.get_group(d)``
844 * ``Partial`` on mesh dim *d* -- all-reduce gradient values
845 using the placement's ``reduce_op`` before norm computation
846 * ``Replicate`` / plain tensor -- no communication
848 This covers FSDP, HSDP, TP+FSDP, and any combination expressible
849 via DTensor placements. PP cross-stage norm aggregation is not
850 yet handled (requires manual all-reduce across stages).
852 Args:
853 parameters: An ``nn.Module``, a single ``Tensor``, or an
854 iterable of ``Tensor`` s whose gradients to clip.
855 max_norm: Maximum allowed gradient norm.
856 norm_type: Type of the norm (default ``2.0``).
857 error_if_nonfinite: If ``True``, raise a ``RuntimeError``
858 when the total norm is non-finite. Default ``False``.
859 foreach: Use the faster foreach-based implementation for the
860 gradient clipping step. If ``None``, use the foreach
861 implementation for devices that support it and silently
862 fall back to the per-tensor implementation for others.
863 Default ``None``.
865 Returns:
866 The total (unclipped) gradient norm as a scalar tensor,
867 cast to the promoted dtype of all gradient tensors.
868 """
869 max_norm = float(max_norm)
870 norm_type = float(norm_type)
872 params = _normalize_parameters(parameters)
874 (
875 grad_groups, all_grads, norm_grads, key_per_grad,
876 mesh_cache, device, has_dtensor_grad,
877 ) = _build_grad_groups(params)
879 # -- Norm + clip (all ranks participate) --------------------------------
880 # _compute_local_norm returns identity elements for empty groups,
881 # so the subsequent all-reduce is safe and semantically neutral.
882 total_norm = _get_total_norm(
883 grad_groups, norm_type, mesh_cache, device,
884 norm_grads, key_per_grad,
885 )
887 if error_if_nonfinite and torch.logical_or(
888 total_norm.isnan(), total_norm.isinf()
889 ):
890 raise RuntimeError(
891 f"The total norm of order {norm_type} for gradients from "
892 "`parameters` is non-finite, so it cannot be clipped. To "
893 "disable this error and scale the gradients by the "
894 "non-finite norm anyway, set "
895 "`error_if_nonfinite=False`"
896 )
898 if all_grads:
899 # Disable foreach for dtensor-backed grads to avoid dispatch issues.
900 effective_foreach = False if has_dtensor_grad and foreach is None else foreach
901 _clip_grads_with_norm_(
902 all_grads, max_norm, total_norm, effective_foreach,
903 )
905 # Promote return dtype to match gradient dtypes (FSDP1 convention).
906 # When this rank has no gradients, return in the default FP32 dtype
907 # (same as FSDP1's behavior to avoid extra communication).
908 if not all_grads:
909 warnings.warn(
910 "clip_grad_norm_ called on this rank with no gradients -- "
911 "returning the local norm in the default dtype "
912 f"{total_norm.dtype}",
913 stacklevel=2,
914 )
915 return total_norm
917 total_norm_dtype = functools.reduce(
918 torch.promote_types,
919 [g.dtype for g in all_grads],
920 )
921 # Return global all-reduced norm, consistent with torchtitan's
922 # full_tensor() approach — .item() returns the correct global value.
923 return total_norm.to(total_norm_dtype)