[shanmugamr1992]

Found by the nemo-rl-testing-agent while validating NVIDIA/Megatron-LM#5700
against the NeMo-RL Megatron L1 functional suite. The break is not caused by
that PR — it reproduces on megatron-core main.

Failure

Every synchronous Megatron generation path raises:

nemo_rl/models/generation/megatron/megatron_worker.py
  tokens = result[i].prompt_tokens.tolist() + result[i].generated_tokens
AttributeError: 'NoneType' object has no attribute 'tolist'

In the L1 Megatron suite that failed grpo_megatron_generation,
grpo_megatron_generation_non_colocated, and
grpo_megatron_generation_multiturn.

Worse, grpo_megatron_generation_async and
grpo_megatron_generation_colocated_async exited 0 while every single
generation raised the same error: rollouts.py catches per-sample generation
exceptions (except Exception: print(...); break), so a completely broken
backend still produces a rollout and a green test. Those two runs trained on
empty generations and reported success.

Root cause

megatron-core #5918 made the
engine drop prompt_tokens from a finished request before serializing it unless
the caller sets the new SamplingParams.return_prompt_tokens. The intent is to
keep large prompt tensors off the engine → coordinator → client path. NeMo-RL was
reading the prompt back out of that reply, so it now gets None.

Fix

Take the prompt from the batch that was submitted rather than from the engine's
reply. NeMo-RL already holds those tokens in data["input_ids"] /
data["input_lengths"] — the same tensors it used to build the request — so this
is the cheaper option as well as the correct one, and it needs no version branch:
it behaves identically against megatron-core revisions that still echo the prompt.

Setting return_prompt_tokens=True would also make the error go away, but it
would pay the transmission cost that #5918 exists to avoid, purely to be told
something we already know.

Validation

• Full L1 Megatron suite on GB200 (4 GPUs) against megatron-core main +
  Megatron-LM#5700: 5 of 6 pass, zero swallowed generation errors (previously 3
  hard failures and 2 silently-degraded passes).
• The remaining failure, grpo_megatron_generation_multiturn, now runs to
  completion and fails a numerics check
  (median(train/token_mult_prob_error) 1.99 vs threshold 1.1) with a 100%
  truncation rate. That is a separate pre-existing problem this crash was hiding;
  it is being attributed against a main baseline and is not addressed here.
• New CPU-only regression test, test_megatron_generation_parse.py, covers a
  reply with prompt_tokens=None, a reply that does echo them (must be
  identical), and an empty generation. The end-to-end megatron tests need GPUs and
  minutes, which is why a wire-format change like this slipped through.

Draft: raised by an agent, needs a human review before merge.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.