`torch.compile` hits recompile limit on `megatron.core.fusions.fused_bias_dropout.py` due to `requires_grad mismatch` in PPO loop
- 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
Assessment
This issue has not been assessed yet.