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

[tengqm]

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 @ 1d82259
• 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:).