GRPO crashes with --transformer-impl local when RL training CUDA graphs are disabled
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
### Bug description
Running RL (GRPO) training with `--transformer-impl local --attention-backend unfused` crashes in the reference-logprobs pass when RL training CUDA graphs are disabled:
```
AssertionError: Packed sequence is not supported by DotProductAttention. Please use TEDotProductAttention instead.
```
`get_logprobs()` (`megatron/rl/rl_utils.py`) and the GRPO `forward_step` (`train_rl.py`) unconditionally build a single-sequence `PackedSeqParams(qkv_format='thd', ...)` whenever sequence packing is disabled, for CUDA graph signature consistency. The non-TE `DotProductAttention` (used by `--transformer-impl local`) asserts `packed_seq_params is None` and has no packed/thd path, so the reference-logprobs pass crashes as soon as RL training CUDA graphs are off (the default).
### Environment
- Megatron-LM main @ 1d82259e
- Hardware without a flash-attn-compatible backend (e.g. Ascend 910B), so the RL recipe uses `--transformer-impl local --attention-backend unfused`
- RL training CUDA graphs off (default: `args.rl_training_cuda_graphs` is False)
### Steps to reproduce
Run GRPO RL training (`train_rl.py --perform-rl-step ...`) with `--transformer-impl local --attention-backend unfused`. The crash occurs in `compute_logprobs_batch` → `get_logprobs` on the reference-logprobs pass.
### Expected behavior
A single-sequence thd is identical to dense. `packed_seq_params` is only needed for CUDA graph signature consistency, so it should only be constructed when RL training CUDA graphs are actually in use; otherwise leave it `None` so the unfused dense path is used.
### Proposed fix
Gate the single-sequence `PackedSeqParams` construction on `args.rl_training_cuda_graphs` in both `get_logprobs()` and the GRPO `forward_step` (change `else:` to `elif args.rl_training_cuda_graphs:`).
Contributor guide
Research direction
Start with megatron/rl/rl_utils.py at get_logprobs() and train_rl.py at the GRPO forward_step, then trace the compute_logprobs_batch reference-logprobs path. Reproduce with --transformer-impl local --attention-backend unfused and RL training CUDA graphs disabled. Done means the reference-logprobs pass completes without the packed-sequence assertion while CUDA-graph behavior remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 84/100