NVIDIA / NVIDIA/Megatron-LM

GRPO crashes with --transformer-impl local when RL training CUDA graphs are disabled

Open Beginner friendly
#6,708 0 comments 0 reactions 0 assignees View on GitHub
bug community-request module: rl
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.