Pass gracefully if token_id not found in message by i-riyad · Pull Request #3862 · NVIDIA/Megatron-LM

i-riyad · 2026-03-13T17:03:16Z

What does this PR do ?

First chunk is meant to fix the chat-template crash when message.content is not a plain string.

Before: some requests had OpenAI-style content like:

content: [{"type":"text","text":"..."}] (list of blocks), or
content: {"text":"..."} (dict)

HF/Jinja templates in Megatron expected a string, so apply_chat_template(...) would blow up with a TypeError during template rendering.
After: it normalizes list/dict/None/non-str content into a string first, so template rendering can proceed.

Second chunk in chat_completions.py fixes a hard failure from the strict assert in prefix-replacement logic.
Before this patch, the code required:

last_assistant_message["prompt_token_ids"]
last_assistant_message["generation_token_ids"]

If either key was missing (or not a list), request handling crashed with an assertion like:
AssertionError: Last assistant message must have prompt_token_ids and generation_token_ids ...

Now it:

checks types safely,
does prefix replacement only when both are valid lists,
otherwise logs a warning and continues with normal retokenized prompt.

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

copy-pr-bot · 2026-03-13T17:03:20Z

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.

github-actions · 2026-03-13T17:03:27Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

santhnm2 · 2026-03-13T17:12:54Z

 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.cuda_graphs import _CudagraphGlobalRecord


Is this change still necessary? Same for the change in transformer_layer.py

I will rebase, this will go away.

santhnm2 · 2026-03-25T19:47:17Z

/ok to test 3a6b0cf

copy-pr-bot · 2026-03-25T19:47:21Z

/ok to test 3a6b0cf

@santhnm2, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

Signed-off-by: rislam <rislam@nvidia.com>

Truncate token-id, routing index, hash, and numeric series fields in chat completion logging to reduce log volume and improve readability while preserving key diagnostics.

Improve qwen3 tool parsing for schema variants and coerce JSON-like structured arguments (arrays/objects) into correct payload types so tool calls are emitted in OpenAI-compatible format.

Apply post-parse guardrails to prevent destructive reservation calls when update tools are present and force consistent hold messaging when transferring to human agents.

Shift object/array argument coercion from qwen3 parser logic into chat_completions using request tool schemas so normalization happens at a single request-level layer. Keep parser extraction behavior and retain the parser tool-schema lookup compatibility path (dict/object tool configs) to avoid schema resolution regressions. Made-with: Cursor Signed-off-by: rislam <rislam@nvidia.com>

santhnm2

I think we need to do a follow-up PR to clean this up and generalize, but will approve for now so we can get it into main.

santhnm2 · 2026-03-26T03:24:57Z

/ok to test a164b8b

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

santhnm2 · 2026-03-26T03:46:43Z

/ok to test 416654b

svcnvidia-nemo-ci · 2026-03-26T05:16:49Z

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/23578747653

i-riyad requested review from a team as code owners March 13, 2026 17:03

svcnvidia-nemo-ci marked this pull request as draft March 13, 2026 17:03

i-riyad marked this pull request as ready for review March 13, 2026 17:03

svcnvidia-nemo-ci requested a review from a team March 13, 2026 17:04

i-riyad changed the title ~~Pass gracefully if token id not found is message~~ Pass gracefully if token_id not found in message Mar 13, 2026

santhnm2 approved these changes Mar 13, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Mar 13, 2026

kvareddy approved these changes Mar 13, 2026

View reviewed changes

santhnm2 mentioned this pull request Mar 13, 2026

Allow untokenized messages when using the prevent tokenization mode. #3864

Closed

5 tasks

shanmugamr1992 approved these changes Mar 14, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Mar 14, 2026

i-riyad force-pushed the rislam/missing-token-id-handle branch 6 times, most recently from fa102a0 to 3a6b0cf Compare March 25, 2026 07:39

i-riyad requested a review from a team as a code owner March 25, 2026 07:39

svcnvidia-nemo-ci removed the Approved All necessary approvals have been made label Mar 25, 2026

i-riyad force-pushed the rislam/missing-token-id-handle branch 3 times, most recently from a4c4971 to fe1840d Compare March 25, 2026 21:56

i-riyad force-pushed the rislam/missing-token-id-handle branch from adbe80a to 53e2731 Compare March 25, 2026 23:54

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Mar 25, 2026

i-riyad added 5 commits March 25, 2026 16:58

Pass gracefully if token id not found is message

826d523

Signed-off-by: rislam <rislam@nvidia.com>

Message normalization from lists

42a5920

Signed-off-by: rislam <rislam@nvidia.com>

redact verbose token/index arrays in inference logs

9300655

Truncate token-id, routing index, hash, and numeric series fields in chat completion logging to reduce log volume and improve readability while preserving key diagnostics.

normalize structured tool arguments in parser and endpoint

b517ef0

Improve qwen3 tool parsing for schema variants and coerce JSON-like structured arguments (arrays/objects) into correct payload types so tool calls are emitted in OpenAI-compatible format.

add tool-call guardrails for transfer and reservation flows

c4636e8

Apply post-parse guardrails to prevent destructive reservation calls when update tools are present and force consistent hold messaging when transferring to human agents.

i-riyad force-pushed the rislam/missing-token-id-handle branch from 53e2731 to 8e1f977 Compare March 25, 2026 23:58

i-riyad force-pushed the rislam/missing-token-id-handle branch from 8e1f977 to 5ccb0ed Compare March 26, 2026 00:07

santhnm2 approved these changes Mar 26, 2026

View reviewed changes

Merge branch 'main' into rislam/missing-token-id-handle

a164b8b

santhnm2 removed the request for review from a team March 26, 2026 03:24

santhnm2 enabled auto-merge March 26, 2026 03:25

svcnvidia-nemo-ci added this to the Core 0.16 milestone Mar 26, 2026

Linting

416654b

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

copy-pr-bot Bot temporarily deployed to test March 26, 2026 03:47 Inactive

santhnm2 added this pull request to the merge queue Mar 26, 2026

Merged via the queue into NVIDIA:main with commit 606afda Mar 26, 2026
62 of 63 checks passed

santhnm2 deleted the rislam/missing-token-id-handle branch March 26, 2026 05:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Pass gracefully if token_id not found in message#3862

Pass gracefully if token_id not found in message#3862
santhnm2 merged 8 commits intoNVIDIA:mainfrom
santhnm2:rislam/missing-token-id-handle

i-riyad commented Mar 13, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Mar 13, 2026

Uh oh!

github-actions Bot commented Mar 13, 2026

Uh oh!

Uh oh!

santhnm2 Mar 13, 2026

Uh oh!

i-riyad Mar 13, 2026

Uh oh!

santhnm2 commented Mar 25, 2026

Uh oh!

copy-pr-bot Bot commented Mar 25, 2026

Uh oh!

santhnm2 left a comment

Uh oh!

santhnm2 commented Mar 26, 2026

Uh oh!

santhnm2 commented Mar 26, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

i-riyad commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Mar 13, 2026

Uh oh!

github-actions Bot commented Mar 13, 2026

Uh oh!

Uh oh!

santhnm2 Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

i-riyad Mar 13, 2026

Choose a reason for hiding this comment

Uh oh!

santhnm2 commented Mar 25, 2026

Uh oh!

copy-pr-bot Bot commented Mar 25, 2026

Uh oh!

santhnm2 left a comment

Choose a reason for hiding this comment

Uh oh!

santhnm2 commented Mar 26, 2026

Uh oh!

santhnm2 commented Mar 26, 2026

Uh oh!

svcnvidia-nemo-ci commented Mar 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

i-riyad commented Mar 13, 2026 •

edited

Loading