[GRPO] add chunked grpo streaming over vocab

[kashif]

Summary

This PR fixes the chunked GRPO loss to compute only selected-token log-probs by streaming over the vocab dimension. This reduces peak memory for the fused-linear chunked path and preserves the existing high-level fused-linear API.

We also fixes the luspo reduction in the chunked path to match TRL exactly, and tightens the torch.compile boundary so we only compile the pure loss computation instead of compiling through the closure that calls torch.autograd.grad.

So now the two implementations are correctly implementing the trade-offs of their design.

Testing Done

• Hardware Type:
• run make test to ensure correctness
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[kashif]

cc @vaibhavjindal for your review

[vaibhavjindal]

Hey @kashif , thanks for the PR, looks good to me. Can you fix the checkstyle and we can then merge?

[kashif]

thanks @vaibhavjindal all done

[kashif]

@vaibhavjindal, just a reminder about this PR, thanks!

[Mecoli1219]

Hi @kashif . I tested it locally on H100 via pytest test/chunked_loss/test_grpo_loss.py, and it didn't pass. Can you check it out?

================================================ short test summary info ================================================
FAILED test/chunked_loss/test_grpo_loss.py::test_selective_chunk_forward_matches_reference[1-4096-256-5000-True-dtype1-0.05-0.05] - AssertionError: Number of mismatched elements: 1
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[None-sequence-luspo-True-True-True-0.1-0.2-0.2-1.0-True-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 14
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[None-sequence-luspo-True-True-True-0.1-0.2-0.2-1.0-False-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 1
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[None-sequence-luspo-True-False-False-0.1-0.2-0.2-1.0-True-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 2
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[None-sequence-luspo-True-False-False-0.1-0.2-0.2-1.0-False-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 9
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[None-sequence-luspo-False-False-True-0.1-0.2-0.2-1.0-True-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 14
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[2.0-sequence-luspo-True-False-False-0.1-0.2-0.2-1.0-True-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 82
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[2.0-sequence-luspo-True-False-False-0.1-0.2-0.2-1.0-False-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 1
FAILED test/chunked_loss/test_grpo_loss.py::test_correctness[2.0-sequence-luspo-False-False-True-0.1-0.2-0.2-1.0-True-1.0-dtype1-2e-05-0.001-8-128-1024-4096] - AssertionError: Number of mismatched elements: 4
========================= 9 failed, 993 passed, 386 skipped, 114 warnings in 601.18s (0:10:01) ==========================

[kashif]

thanks @Mecoli1219 checking

[kashif]

the issue could be that on h100 (and later), TF32 is enabled by default for fp32 matmul which yields different rounding between the chunked path and the full-matmul reference... I can check how to disable TF32 and do everything in fp32?

[Mecoli1219]

Hi @kashif . Thanks for the quick update. Overall LGTM besides the updated __init__.py file.

I tested it on H100 and only get a slightly mismatch error, could be fixed by increasing the tol slightly.

FAILED test/chunked_loss/test_grpo_loss.py::test_selective_chunk_forward_matches_reference[1-4096-256-5000-True-dtype1-0.05-0.05] - AssertionError: Number of mismatched elements: 1

E           AssertionError: Number of mismatched elements: 1
E           Mismatch at index (0, 2460): tensor1[(0, 2460)] = -0.38605499267578125, tensor2[(0, 2460)] = -0.47758522629737854

Can you also check if the new VESPO test success on your local machine? Thanks!

[kashif]

thanks @Mecoli1219 I relaxed the tol and the VESPO tests all pass on my local gpu

[Mecoli1219]

LGTM! Thanks for the contribution!

[kashif]

thanks @Mecoli1219