fix: use prompt token length for advantage group extraction and fix token mask

[yfw]

This PR fixes two multi-turn GRPO training issues:

1. The previous role-based extraction (_extract_prompt_only_messages) broke on multi-turn prompts containing assistant messages in the conversation history — it would strip them, corrupting the prompt IDs used for advantage estimation.
   Replace with extract_initial_prompt_messages() which uses the length field to identify the original prompt boundary. Applied to both sync and async GRPO paths.

2. GRPO token loss masks previously unmasked every message with role == "assistant". In multi-turn data, assistant messages can be part of the prompt history, not generated rollout output, so those tokens should not contribute to the policy loss. This PR updates masking so only assistant messages produced by generation, identified by existing generation_logprobs, are trainable. Missing generation_logprobs are still filled with zeros for downstream tensorization.

Closes #1960 and #1956

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[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.

[yuki-97]

/ok to test 628a248

[macandro96]

/ok to test 7d96ad3

[copy-pr-bot]

/ok to test 7d96ad3

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

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

[macandro96]

/ok to test d5961e5

[yuki-97]

@yfw @HeyyyyyyG could you help to take a review as well?

[HeyyyyyyG]

LGTM

[yuki-97]

/ok to test 49601f1

[macandro96]

/ok to test 18a05b0

[macandro96]

/ok to test 8e1c343

[macandro96]

@yuki-97 fixed! can you reapprove please?

[yuki-97]

LGTM, thanks @macandro96

@ZhiyuLi-Nvidia could you help to take a check at TQ path?
as I know we'll always have generation_logprobs in rollout, so I think it should be fine to remove setdefault generation_logprobs in TQ path, but better to double check.

[ZhiyuLi-Nvidia]

@ZhiyuLi-Nvidia could you help to take a check at TQ path?
as I know we'll always have generation_logprobs in rollout, so I think it should be fine to remove setdefault generation_logprobs in TQ path, but better to double check.

Thank you @yuki-97 @macandro96 ! LGTM and thank you for the change and review for data plane coverage as well.

Out of scope for this PR, but I have some concerns about the message_logs: list data type, it creates friction for both data transfer and storage with additional preprocessing/postprocessing overhead.

• The type is quite loose and too general to allow proper memory budgeting up front. It's unclear whether it's nested lists of objects or something else.
• For future GPU-based RDMA communication, a list of tensors likely implies padding overhead. With the loose type, I am not even sure how to properly handle it.

I think we should eventually move to nested jagged tensors: a single concatenated tensor paired with a sequence-length tensor to mark boundaries. I double-checked, and veRL uses nested jagged tensors as well. Curious if you agree with this direction.

[yuki-97]

thanks @ZhiyuLi-Nvidia @macandro96 , I've merged but feel free to continue discussing here or in slack.

besides, message_logs: list has both tensor and string in it now, I feel we need to see how to structure them if we eventually move to nested jagged tensors.