Skip to content

[Pipeline RL] Tensor-parallel GRPO loss #464

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
jlamypoirier merged 12 commits into from
Jan 29, 2026
Merged

[Pipeline RL] Tensor-parallel GRPO loss #464

jlamypoirier merged 12 commits into from
Jan 29, 2026

Conversation

@jlamypoirier
Copy link
Collaborator

@jlamypoirier jlamypoirier commented Jan 28, 2026 •
edited
Loading

✨ Description

  • Tweak entropy losses to improve readability
  • Extract reusable fused_predicted_logits_from_labels from _fused_cross_entropy_base_from_labels
  • Fix broken entropy loss tests for loss masking.
  • Support arbitrary tensor dimensions in entropy losses.
  • Add TP support for GRPO
  • Add TP support for Z loss
  • Generalize and simplify entropy loss tests with support for z loss and grpo, improve distributed coverage, add fp16 and bf16 tests.

Base automatically changed from jlp_grpo_sample to jlp_pipeline_rl January 28, 2026 23:36
@jlamypoirier jlamypoirier marked this pull request as ready for review January 28, 2026 23:43
@jlamypoirier jlamypoirier changed the title Tensor-parallel GRPO loss [Pipeline RL] Tensor-parallel GRPO loss Jan 29, 2026
@jlamypoirier jlamypoirier merged commit caaa9f8 into jlp_pipeline_rl Jan 29, 2026
@jlamypoirier jlamypoirier deleted the jlp_grpo_tp branch January 29, 2026 20:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jlamypoirier