optimizer_cuda_graph=True crashes during graph capture when grad_norm_skip_threshold uses default float('inf')
- 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
Assessment
This issue has not been assessed yet.