No way for low-overhead total norm in native PyTorch with large number of tensors

[awgu]

Context
Gradient norm clipping is a popular technique for stabilizing training, which requires computing the total norm with respect to the model's gradients. This involves a norm reduction across all of the gradients down to a single scalar.

PyTorch's clip_grad_morm_ offers both single-tensor (foreach=False) and multi-tensor (foreach=True) implementations. However, even the multi-tensor foreach implementation incurs high CPU overhead when computing the total norm over a large number of tensors.

pytorch/torch/nn/utils/clip_grad.py

Lines 81 to 96 in 3434a54

	norms: List[Tensor] = []
	for (device, _), ([device_grads], _) in grouped_grads.items(): # type: ignore[assignment]
	if (foreach is None and _has_foreach_support(device_grads, device)) or (
	foreach and _device_has_foreach_support(device)
	):
	norms.extend(torch._foreach_norm(device_grads, norm_type))
	elif foreach:
	raise RuntimeError(
	f"foreach=True was passed, but can't use the foreach API on {device.type} tensors"
	)
	else:
	norms.extend([torch.linalg.vector_norm(g, norm_type) for g in device_grads])
	total_norm = torch.linalg.vector_norm(
	torch.stack([norm.to(first_device) for norm in norms]), norm_type
	)

The foreach implementation involves:

1. torch._foreach_norm(tensors) to return a list of 0D scalars, representing the norm of each gradient
2. torch.stack to cat the 0D scalars into a 1D tensor
3. torch.linalg.vector_norm() to compute the norm of the norms, representing the final total norm

Issue
The foreach implementation incurs much unnecessary CPU overhead, making the clipping heavily CPU bound when operating on a large number of tensors. For example, for ~2000 tensors, the total norm calculation takes >18 ms on CPU, while only 1.3 ms on GPU (for a particular real workload -- larger tensors would make this slower).

Assuming N tensors, some inefficiencies in the existing implementation arise from:

• _foreach_norm must call N aten::empty({}) to construct the N 0D scalar outputs (code).
• The N 0D scalars need to be stackd for the final norm reduction. stack requires 1D tensors, so each of the N 0D scalar gets unsqueezed again in stack.

Together, this leads to an extra 2N dispatcher calls to handle the N intermediate scalars. Ideally, we can avoid materializing these N intermediates, especially as torch.Tensors, where one option is a fused kernel.

Today, torch.compile cannot address this issue in a satisfying way. Default torch.compile cannot achieve horizontal fusion, leading to slower performance than eager mode. torch.compile(mode="reduce-overhead") does reduce overhead more than eager but results in 2x memory usage, likely due to copying the gradients into CUDA graph addresses. Note that we likely cannot mark the inputs as static because gradients are computed anew every iteration at possibly different addresses, and for my use case, we cannot compile the gradient allocation with torch.compile (due to FSDP).

Here is an example script for getting profiler traces of various implementations: P1529496302

The ask is to provide some native way to make this total norm calculation not CPU-overhead bound, e.g. via a fused op.

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @crcrpar @mcarilli @janeyx99 @ezyang @chauhang @penguinwu

[janeyx99]

cc @mlazos This looks similar to compiling across foreach ops for the horizontal fusion portion. Is this something that torch.compile could fuse?

[carmocca]

Since #139467 separated the norm computation from the clipping, the following is no longer a limitation for all cases

Note that we likely cannot mark the inputs as static because gradients are computed anew every iteration at possibly different addresses

As _get_total_norm can be called for non-grad tensors too (like parameters or optimizer states)

So the huge cpu overhead of the norm calculation is highly relevant in this case, because computing the norms for those objects needs to happen after the optimizer step so there are few opportunities to overlap the work within the same training step. It would need to be overlapped with the next one.

[mlazos]

You should be able to use foreach_map to compile a fast kernel to do this. Lmk if you need help with this. There is an example implementing adam with it in the compiled optimizer tests.

[awgu]

@mlazos can you clarify which part of the problem foreach_map is addressing (e.g. compared to existing _foreach_norm)?

Also, looking at the unit test, should we be using the output of this foreach_map? Or can you explain the in-place semantic if not?

pytorch/test/inductor/test_compiled_optimizers.py

Lines 114 to 115 in 4ab967c

	if weight_decay != 0:
	foreach_map(torch.add, (grads,), alpha=weight_decay)

[awgu]

@mlazos I tried some variations of using foreach_map with torch.linalg.vector_norm, and they error:

LoweringException: SubgraphLoweringException: Buffers cannot be created while lowering a pointwise subgraph. This could be for a good reason (e.g. you're calling an op we can't codegen as a pointwise op), but it could also be a bug. Please file a bug report if you think this should be supportable.

Here is a minimal usage:

Repro

import torch

from torch._higher_order_ops import foreach_map

def compute_total_norm(tensors: list[torch.Tensor], norm_type: float):
    norms = foreach_map(torch.linalg.vector_norm, (tensors,), ord=norm_type)
    return torch.linalg.vector_norm(torch.stack(norms), ord=norm_type)


if __name__ == "__main__":
    torch.manual_seed(42)
    num_tensors = 10
    ts = [
        torch.randn((1024, 1024), device="cuda") for _ in range(num_tensors)
    ]

    @torch.compile()
    def compute_total_norm_fn():
        return compute_total_norm(ts, norm_type=2.0)

    total_norm = compute_total_norm_fn()