NVIDIA / NVIDIA/Megatron-LM

optimizer_cuda_graph=True crashes during graph capture when grad_norm_skip_threshold uses default float('inf')

Open
#5,889 1 comment 0 reactions 1 assignee Claimed by @Phlip79 View on GitHub
bug community-request waiting-on-customer
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Environment:**
1.Megatron-LM: core_r0.18.0 && Transformer Engine: 2.14

**Summary:**
When --optimizer-cuda-graph is enabled, training crashes during the CUDA graph capture step (warmup step) with:
torch.AcceleratorError: CUDA error: operation not permitted when stream is capturing
The error occurs at the grad_norm > grad_norm_skip_threshold comparison in ChainedOptimizer.step().

**Root Cause**
Two commits interact badly:
3d87bfc1b — introduced OptimizerCudaGraphWrapper, which wraps the entire ChainedOptimizer.step() inside a torch.cuda.graph() capture context on a non-default stream.
180131620 — added the grad_norm > grad_norm_skip_threshold check inside ChainedOptimizer.step().

During CUDA graph capture, grad_norm returned by get_grad_norm_fp32() is a 0-dim CUDA tensor (not a Python float) when Transformer Engine is installed, because multi_tensor_scale_tensor_impl is not None and total_norm.pow(...) returns a tensor:

[clip_grads.py#L139-L142](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/optimizer/clip_grads.py#L139-L142)

if multi_tensor_scale_tensor_impl is not None:
total_norm = total_norm.pow(1.0 / norm_type) # ← 0-dim CUDA tensor
else:
total_norm = total_norm.item() ** (1.0 / norm_type) # ← Python float
Then at the comparison:

[optimizer.py#L1482](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/optimizer/optimizer.py#L1482)

if grad_norm > optimizer.config.grad_norm_skip_threshold:
tensor > float('inf') produces a BoolTensor, and Python's if calls .item() → cudaDeviceSynchronize() — which is not permitted while the stream is capturing.

This happens even with the default value float('inf') where the comparison is always False.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.