Conversation
…nistic CPU tests Issue 4227 reports a suspected shape mismatch between GRPOLoss masking strategies. Local verification on CPU shows that every supported combination of wrapper input mode and masking strategy produces boolean (batch, T) masks and a finite loss, and that the failing combinations are contract questions rather than shape bugs: rlhf with tokens input cannot have an assistant mask, and sft with history input silently falls back to the assistant mask. This documents the full matrix in the GRPOLoss docstring and adds a deterministic CPU test class that covers it with the tiny trl-internal-testing checkpoint, fabricating the response tokens so no generation engine is needed. The tests fail loudly if the masking contract drifts.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4340
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
github-actions
Bot
added
Documentation
vmoens
left a comment
vmoens
left a comment
There was a problem hiding this comment.
The new CPU cases exercise the loss with supported input shapes, but the documentation rules out a supported tokens-mode RLHF input, and the scalar comparisons do not establish which tokens each strategy selects. Please address the inline comments before merging.
Suggested next steps:
- Correct the assistant-mask requirement and describe the intermediate shapes explicitly: masks and log probabilities are
(B, T), log-weights are(B, T, 1), and advantages can be(B, 1, 1)per trajectory or(B, T, 1)per token for these batched inputs. - Strengthen the existing cases with expected distribution masks and unreduced loss shapes. Cover tokens-mode RLHF with both an absent and a caller-supplied assistant mask. These checks can fit into the existing cases without adding another layer of redundant tests.
- In the history helper, remove
tokenizer_kwargs={"chat_template_name": "qwen"}or use the top-levelchat_template_name="qwen"; the history path constructs its own tokenizer kwargs, so the nested setting is ignored.
Further follow-ups: make the missing-assistant-mask error suggest supplying the mask or using history input, and investigate the original skipped vLLM integration with actual generated output. The fabricated-response tests do not establish that this integration is fixed, so validate that path before removing its skip and clarify the PR's replacement claim accordingly.
Validation on e0f3d7b: pytest test/llm/test_llm_objectives.py -k 'TestGRPOLossMaskingContract or TestLosses' passed all 13 selected tests on CPU with PyTorch 2.11.0 and Transformers 4.56.2. A separate probe confirmed caller-supplied assistant masks work in tokens mode. It also confirmed all six new cases still pass with SFT/generic swapped for tokens input and generic incorrectly using assistant masking for history input. GPU/vLLM integration was not run. I found no evidence here that warrants changing the loss math.
| * ``tokens`` input: ``"sft"`` masks the prompt positions using ``("tokens", "prompt")``, | ||
| ``"generic"`` masks padding positions using the attention mask, and ``"rlhf"`` requires | ||
| an assistant mask which is only computed in history mode, so ``"rlhf"`` with ``tokens`` | ||
| input raises a ``ValueError``. |
There was a problem hiding this comment.
[P2] Qualify the tokens-mode RLHF restriction
The tokens log-prob path preserves a caller-supplied ("masks", "all_assistant_mask"), and the RLHF distribution uses it. I verified that supplying a boolean response mask makes GRPOLoss(masking_strategy="rlhf") return a finite loss with exactly that mask. Please say this raises only when no assistant mask is available, and cover both missing and supplied masks in the tests. Tokens input itself is supported.
There was a problem hiding this comment.
Done in 9dc005c. The docstring now states that rlhf with tokens input uses a caller supplied assistant mask and raises only when it is absent, and the tokens rlhf case covers both sides: without a mask it raises, with a supplied boolean mask the distribution mask equals exactly that mask and the loss is finite.
| sft = GRPOLoss(actor_network=wrapper, masking_strategy="sft")(td) | ||
| generic = GRPOLoss(actor_network=wrapper, masking_strategy="generic")(td) | ||
| assert torch.isfinite(sft.loss_objective) | ||
| assert torch.isfinite(generic.loss_objective) | ||
| # sft excludes the prompt positions, generic includes them, so with the | ||
| # same data and advantage the two losses must differ | ||
| assert sft.loss_objective.item() != generic.loss_objective.item() |
There was a problem hiding this comment.
[P2] Assert the expected token selection
Unequal scalar losses do not establish that either strategy selects the correct tokens. I swapped SFT and generic for tokens input and made generic use assistant masking for history input; all six new cases still passed. With the balanced per-trajectory advantages, the history equality check can also compare zero losses despite different token selections. Please assert the expected distribution masks: attention with prompt positions cleared for tokens-mode SFT, attention for generic, and the assistant mask for history-mode SFT/RLHF when prompt tokens are absent. Retain the forward checks and verify the unreduced (B, T, 1) loss shape to protect the shape contract.
There was a problem hiding this comment.
Done in 9dc005c. Each case now compares the distribution mask against the expected tensor: attention with prompt positions cleared for tokens sft, attention for generic, assistant mask for history sft and rlhf. The scalar loss equality case is removed in favor of those comparisons, and the parametrized history case also checks the unreduced (B, T, 1) shape with aggregation='none'. One note: the wrapper output masks arrive as int64 in the tokens path, so the comparisons cast both sides to bool; the history masks are already boolean.
| In every supported combination the masks are boolean tensors of shape ``(batch, T)`` over | ||
| the full padded sequence and the distribution log probability has the same shape, so a per | ||
| trajectory advantage of shape ``(batch, 1, 1)`` broadcasts against them. |
There was a problem hiding this comment.
[P3] Describe the log-weight dimension used for broadcasting
The advantage broadcasts against the (B, T, 1) log-weight produced by _log_weight, rather than directly against the (B, T) log probabilities or masks. Directly combining (B, 1, 1) with (B, T) would broadcast to (B, B, T). Please state the intermediate singleton dimension and the supported (B, 1, 1) per-trajectory / (B, T, 1) per-token advantage shapes explicitly; forward requires advantage and log-weight to have the same rank.
There was a problem hiding this comment.
Done in 9dc005c. The note now states the (B, T, 1) log weight that the advantage broadcasts against, the supported (B, 1, 1) per trajectory and (B, T, 1) per token advantage shapes, and the same rank requirement enforced by forward. Also fixed the history helper to pass chat_template_name to the wrapper constructor instead of the ignored nested tokenizer_kwargs. Both follow ups (the missing mask error message and validating the skipped vLLM integration with generated output) are acknowledged as separate work; the PR description no longer claims to replace that test.
…nd tests The tokens input supports rlhf whenever the tensordict carries a caller supplied assistant mask, so the docstring now states that the ValueError only fires when the mask is absent, and describes the intermediate shapes: masks and log probabilities are (B, T), the log weight is (B, T, 1), and the advantage can be (B, 1, 1) per trajectory or (B, T, 1) per token. The tests now assert the expected distribution masks (prompt positions cleared for tokens sft, attention for generic, assistant mask for history sft and rlhf), cover tokens rlhf with an absent and with a caller supplied mask, and check the unreduced (B, T, 1) loss shape. The redundant loss equality case is replaced by the direct mask comparisons, and the history helper passes chat_template_name to the wrapper constructor instead of the ignored nested tokenizer_kwargs.
2 participants
Description
Documents the masking strategy contract of
GRPOLossin its docstring and adds a deterministic CPU test class that pins it down.Context: issue #4227 reports a suspected shape mismatch between masking strategies. I reproduced the full matrix on CPU with
trl-internal-testing/tiny-Qwen2ForCausalLM-2.5and theQwen/Qwen2.5-0.5Btokenizer, fabricating the response tokens so that no generation engine is required. Where a supported combination runs, masks are boolean tensors of shape(batch, T)over the full padded sequence, the distribution log probability has the same shape, and a per trajectory advantage of shape(batch, 1, 1)broadcasts cleanly. I could not reproduce a shape bug in the loss reduction itself.What the docstring note now states, and what the tests assert:
tokensinput:sftmasks prompt positions using("tokens", "prompt"),genericmasks padding using the attention mask, andrlhfuses a caller supplied("masks", "all_assistant_mask")when the tensordict carries one and raises aValueErroronly when it does not. The tests compare the distribution mask against those expected tensors, so a drift in any selection fails in CI.historyinput: the assistant mask is computed from the chat template.rlhfuses it directly andsftfalls back to it whenever("tokens", "prompt")is unavailable, which makessftandrlhfselect the same tokens in this mode.GRPOLoss.forwardwith a finite loss for all three strategies, and the unreduced loss keeps the(B, T, 1)token layout the advantage broadcasts against.No behavior is changed: the docstring documents current behavior, and the two open contract questions from my comment on #4227 (whether
sftin history mode should warn about the fallback, and what errorrlhfwithtokensinput should suggest) remain for the maintainers to decide.Correction to my earlier comment on #4227
In that comment I claimed the transformers wrapper history path fails at batch sizes greater than one. That was wrong, and the probe script I referenced was the cause: it built the
Historyincorrectly, stacking prompt and response as separate conversation slots instead of concatenated messages, which produced a malformed(batch, 2, T)token tensor. With a correctly builtHistory(per rowHistory.from_chats, extended along the message dimension withdim=-1, then lazy stacked) the wrapper handles batch 2 with heterogeneous prompts fine, including the assistant masks, for all three strategies. The wrapper batch path is therefore not implicated in the skipped integration test as I suggested.Test results
Motivation and Context
Fixes the documentation side of #4227 and pins the contract with deterministic CPU tests. The skipped vLLM integration test is left untouched: its path involves actual generated output, which the fabricated response cases here do not exercise, so validating that integration remains a follow-up before its skip could be lifted.
Types of changes
Checklist