From eaf5f55d15e7a83c97721bc7826c2629562cc490 Mon Sep 17 00:00:00 2001 From: shanmugamr1992 Date: Mon, 27 Jul 2026 11:31:47 -0700 Subject: [PATCH] fix: stop relying on megatron-core echoing prompt tokens back megatron-core now drops `prompt_tokens` from a finished inference request before serializing it unless `SamplingParams.return_prompt_tokens` is set (NVIDIA/Megatron-LM#5918), so the reply the inference client hands back carries `prompt_tokens=None` and packing the result raised `AttributeError: 'NoneType' object has no attribute 'tolist'`. Take the prompt from the batch that was submitted instead of from the engine's reply. NeMo-RL already holds those tokens, so this avoids shipping the prompt back over ZMQ purely to echo it, and it works against both the old and new megatron-core behaviour without a version branch. Three L1 Megatron generation tests failed outright on this; two async ones exited 0 because the rollout loop catches per-sample generation errors, so they were silently generating nothing. Signed-off-by: shanmugamr1992 --- .../generation/megatron/megatron_worker.py | 17 +- .../test_megatron_generation_parse.py | 152 ++++++++++++++++++ 2 files changed, 163 insertions(+), 6 deletions(-) create mode 100644 tests/unit/models/generation/test_megatron_generation_parse.py diff --git a/nemo_rl/models/generation/megatron/megatron_worker.py b/nemo_rl/models/generation/megatron/megatron_worker.py index 367c96cb1e9..3615a4a41d7 100644 --- a/nemo_rl/models/generation/megatron/megatron_worker.py +++ b/nemo_rl/models/generation/megatron/megatron_worker.py @@ -524,13 +524,18 @@ def _parse_result_to_batched_data_dict( batch_size, dtype=torch.long, device=input_ids.device ) for i in range(batch_size): - tokens = result[i].prompt_tokens.tolist() + result[i].generated_tokens - seq_len = len(tokens) - output_ids_padded[i, :seq_len] = torch.tensor( - tokens, dtype=torch.long, device=input_ids.device - ) + # Take the prompt from the request we submitted rather than from the + # engine's reply: mcore only echoes prompt_tokens back when + # SamplingParams.return_prompt_tokens is set, and asking for them would + # ship the whole prompt over ZMQ for data we already hold. prompt_len = input_lengths[i].item() - generation_lengths[i] = seq_len - prompt_len + generated_tokens = result[i].generated_tokens + seq_len = prompt_len + len(generated_tokens) + output_ids_padded[i, :prompt_len] = input_ids[i, :prompt_len] + output_ids_padded[i, prompt_len:seq_len] = torch.tensor( + generated_tokens, dtype=torch.long, device=input_ids.device + ) + generation_lengths[i] = len(generated_tokens) unpadded_sequence_lengths[i] = seq_len gen_logprobs = result[i].generated_log_probs logprobs_padded[i, prompt_len : prompt_len + len(gen_logprobs)] = ( diff --git a/tests/unit/models/generation/test_megatron_generation_parse.py b/tests/unit/models/generation/test_megatron_generation_parse.py new file mode 100644 index 00000000000..12fc5bcc105 --- /dev/null +++ b/tests/unit/models/generation/test_megatron_generation_parse.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for packing mcore inference replies into a GenerationOutputSpec. + +The end-to-end megatron generation tests need GPUs and several minutes, which +makes them a poor guard for the wire format of the engine's reply. That format +does change: megatron-core stopped echoing `prompt_tokens` back unless +`SamplingParams.return_prompt_tokens` is set, and every synchronous L1 functional +test began failing with `AttributeError: 'NoneType' object has no attribute +'tolist'` -- while two async ones swallowed the error per sample and still +reported success. These tests pin the packing logic without a GPU. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.generation.megatron.megatron_worker import MegatronGenerationMixin + +PAD = 0 + + +class FakeInferenceReply: + """The subset of mcore's DynamicInferenceRequest the parse path reads. + + `prompt_tokens` defaults to None to match what the inference client actually + hands back: the engine drops it before serializing unless the caller opts in. + """ + + def __init__( + self, + generated_tokens: list[int], + generated_log_probs: list[float], + prompt_tokens: torch.Tensor | None = None, + ) -> None: + self.prompt_tokens = prompt_tokens + self.generated_tokens = generated_tokens + self.generated_log_probs = generated_log_probs + + +def parse(data: BatchedDataDict, result: list[FakeInferenceReply]) -> BatchedDataDict: + """Calls the mixin method unbound, so no engine or distributed setup is needed.""" + worker = SimpleNamespace(tokenizer=SimpleNamespace(pad_token_id=PAD)) + return MegatronGenerationMixin._parse_result_to_batched_data_dict( + worker, data, result + ) + + +@pytest.fixture +def two_sample_batch() -> BatchedDataDict: + """Two right-padded prompts of different lengths (3 and 2 real tokens).""" + return BatchedDataDict( + { + "input_ids": torch.tensor( + [[5, 6, 7, PAD], [8, 9, PAD, PAD]], dtype=torch.long + ), + "input_lengths": torch.tensor([3, 2], dtype=torch.long), + } + ) + + +@pytest.mark.mcore +def test_packs_replies_that_omit_prompt_tokens(two_sample_batch): + """The regression: a reply with prompt_tokens=None must still pack correctly.""" + result = [ + FakeInferenceReply(generated_tokens=[11, 12], generated_log_probs=[-0.1, -0.2]), + FakeInferenceReply(generated_tokens=[13], generated_log_probs=[-0.3]), + ] + + out = parse(two_sample_batch, result) + + # Row width is the padded prompt length (4) plus the longest generation (2). + torch.testing.assert_close( + out["output_ids"], + torch.tensor( + [[5, 6, 7, 11, 12, PAD], [8, 9, 13, PAD, PAD, PAD]], dtype=torch.long + ), + ) + # Generation is placed at the true prompt length, not the padded length, so + # sample 1's tokens must land at index 2 rather than 4. + torch.testing.assert_close(out["generation_lengths"], torch.tensor([2, 1])) + torch.testing.assert_close(out["unpadded_sequence_lengths"], torch.tensor([5, 3])) + + expected_logprobs = torch.zeros(2, 6) + expected_logprobs[0, 3:5] = torch.tensor([-0.1, -0.2]) + expected_logprobs[1, 2:3] = torch.tensor([-0.3]) + torch.testing.assert_close(out["logprobs"], expected_logprobs) + + +@pytest.mark.mcore +def test_ignores_echoed_prompt_tokens(two_sample_batch): + """An engine that does echo the prompt back must not change the output. + + The prompt is taken from the submitted batch either way, so the two mcore + behaviours are indistinguishable downstream and no version branch is needed. + """ + without_echo = [ + FakeInferenceReply(generated_tokens=[11, 12], generated_log_probs=[-0.1, -0.2]), + FakeInferenceReply(generated_tokens=[13], generated_log_probs=[-0.3]), + ] + with_echo = [ + FakeInferenceReply( + generated_tokens=[11, 12], + generated_log_probs=[-0.1, -0.2], + prompt_tokens=torch.tensor([5, 6, 7], dtype=torch.long), + ), + FakeInferenceReply( + generated_tokens=[13], + generated_log_probs=[-0.3], + prompt_tokens=torch.tensor([8, 9], dtype=torch.long), + ), + ] + + baseline = parse(two_sample_batch, without_echo) + echoed = parse(two_sample_batch, with_echo) + + for key in ( + "output_ids", + "logprobs", + "generation_lengths", + "unpadded_sequence_lengths", + ): + torch.testing.assert_close(echoed[key], baseline[key]) + + +@pytest.mark.mcore +def test_handles_a_reply_with_no_generated_tokens(two_sample_batch): + """A request that produced nothing must yield the prompt and a zero length.""" + result = [ + FakeInferenceReply(generated_tokens=[], generated_log_probs=[]), + FakeInferenceReply(generated_tokens=[13], generated_log_probs=[-0.3]), + ] + + out = parse(two_sample_batch, result) + + torch.testing.assert_close(out["generation_lengths"], torch.tensor([0, 1])) + torch.testing.assert_close(out["unpadded_sequence_lengths"], torch.tensor([3, 3])) + assert out["output_ids"][0].tolist() == [5, 6, 7, PAD, PAD]