feat(grpo): add loss_type='grpo_sample' (per-sample/per-trajectory GRPO) - #11
Draft
wanyaworld wants to merge 2 commits into
Draft
feat(grpo): add loss_type='grpo_sample' (per-sample/per-trajectory GRPO)#11wanyaworld wants to merge 2 commits into
wanyaworld wants to merge 2 commits into
Conversation
Per-sample normalization for packed GRPO: loss = sum_s( sum_{t in s} pg_t /
max(N_s, min_tokens) ). Implemented as a weighted sum (per_token_loss *
sample_weight) where the caller bakes sample_weight[t] = mask[t]/max(N_s,min)
from cu_seqlens -> no scatter inside the fused kernel; caller normalizes the
accumulated sum by global_valid_samples. sample_weight threaded whole through
the partial (mirrors full_attention_mask); Phase-1 asserts the chunk covers
the batch (B==1 packed row). Existing loss_types unchanged.
7 CPU tests: scatter-reference equivalence, single-sample==bnpo anchor, assert,
gradient flow.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
wanyaworld
force-pushed
the
feat/grpo-sample-loss
branch
from
June 24, 2026 04:49
604e9be to
aa702d9
Compare
wanyaworld
force-pushed
the
feat/grpo-sample-loss
branch
2 times, most recently
from
June 24, 2026 06:56
0bb5fbe to
bb67463
Compare
sample_weight is now sliced per chunk (torch.chunk on dim 0) and threaded through fused_fwd_bwd / accumulate_chunk exactly like vllm_is_ratio, instead of being bound whole into the partial. So loss_type='grpo_sample' works for B > chunk_size (chunks>1), e.g. unpacked + train_micro_batch_size>1 — not only the B==1 packed-row case. The kernel accumulates per-chunk loss as before. test: chunks>1 == single-chunk == scatter reference; existing loss_types stay chunk-consistent. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
wanyaworld
force-pushed
the
feat/grpo-sample-loss
branch
from
June 24, 2026 07:25
bb67463 to
55a035b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
목적
GRPO chunked loss에
loss_type="grpo_sample"추가 — packing된 row 안의 각 trajectory(sample)를 자기 토큰수로 정규화한 합:호출측(NeMoRL)이
global_valid_samples로 나눠 trajectory 평균 완성. 기존 token-level(bnpo)을 대체(합산 아님).구현 (핵심: 커널엔 scatter 없음)
caller가 per-token weight
sample_weight[t] = mask[t]/max(N_s, min)를 cu_seqlens로 미리 구워 넘기면, 커널은 weighted-sum 한 줄로 per-sample 정규화:sample_weight는vllm_is_ratio처럼 chunk별로 슬라이스되어 배선됨 →chunks>1(B>chunk_size, 예: unpacked + micro_batch>1)도 지원.커밋
feat: grpo_sample 커널 (베이스)feat: chunks>1 per-chunk sample_weight 지원테스트
test/chunked_loss/test_grpo_sample_loss.py— CPU 8개: 독립 scatter 레퍼런스 등가, single-sample==bnpo 앵커, sample_weight assert, chunks>1 == single-chunk == reference, gradient flow.이 PR(커널)은 NeMoRL-torchtitan linkedin#92(per-sample 배선)와 반드시 같이 머지되어야 학습에 쓰입니다. 단독 머지 시
grpo_sampleloss_type만 존재하고 호출부 없음.