NVIDIA / NVIDIA/Megatron-LM

`torch.compile` hits recompile limit on `megatron.core.fusions.fused_bias_dropout.py` due to `requires_grad mismatch` in PPO loop

Open
#1,888 2 comments 0 reactions 0 assignees View on GitHub
bug module: rl
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Bug Description**

When running Megatron-LM within a PPO (Reinforcement Learning) training loop on a large cluster, `torch.compile` (Dynamo) is unable to correctly cache compiled graphs for fused kernels, leading to a critical training failure.

The PPO loop alternates between inference (rollout, `requires_grad=False`) and training (update, `requires_grad=True`). This switching of `requires_grad` state causes Dynamo to repeatedly recompile the Megatron fused function `bias_dropout_add_fused_train`.

This quickly hits the `torch._dynamo.config.recompile_limit` (8), at which point Dynamo gives up. This breaks the autograd graph, and during the backward pass, PyTorch emits a critical warning that `c10d::broadcast_` does not have a registered autograd kernel, "which may lead to silently incorrect behavior."

This is not just a performance warning; it indicates that gradients are likely incorrect, causing the training to fail silently.

**Environment**

* **Project:** Megatron-LM (specifically `megatron.core`)
* **PyTorch Version:** (Please specify your PyTorch version, e.g., 2.x)
* **Context:** PPO training loop
* **Hardware:** 64-node H200 Cluster
* **Parallelism:** `PP=16`, `TP=1`, `EP=32`

**Reproduction**

Run a PPO training job (e.g., using `verl`) with `torch.compile` enabled and Megatron-LM as the backend on a multi-node setup.

**Precise Error Logs**

The following logs (with `TORCH_LOGS="recompiles,dynamo,autograd"`) confirm the root cause:

**1. Root Cause: Dynamo Recompile Limit (requires\_grad mismatch)**
This log shows Dynamo hitting its limit for the *specific Megatron function* due to the `requires_grad` flag changing on its input.

```
[rank25]:W1023 12:14:24.846000 591 torch/_dynamo/convert_frame.py:1016] [6/8] torch._dynamo hit config.recompile_limit (8)
[rank25]:W1023 12:14:24.846000 591 torch/_dynamo/convert_frame.py:1016] [6/8] function: 'bias_dropout_add_fused_train' (/workspace/Megatron-LM/megatron/core/fusions/fused_bias_dropout.py:67)
[rank25]:W1023 12:14:24.846000 591 torch/_dynamo/convert_frame.py:1016] [6/8] last reason: 6/7: tensor 'x_with_bias[0]' requires_grad mismatch. expected requires_grad=0
```

**2. Critical Symptom: Autograd Warning (`c10d::broadcast_`)**
This is the downstream effect of the broken autograd state, indicating a silent failure.

```
[rank492]:W1023 12:06:01.036000 646 torch/autograd/graph.py:829] [8/1] c10d::broadcast_: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior.
```

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.