Skip to content

Conversation

@Jack-Khuu
Copy link
Contributor

@Jack-Khuu Jack-Khuu commented Sep 26, 2025

Moves compute_logprobs from GRPO app to utils/ops.py and adds basic tests for it

Also updated the logit slicing math, to an equivalent version based on TRL

Old

    context_length = logits.shape[1] - input_ids.shape[1]
    logits = logits[:, context_length - 1 : -1].to(input_ids.device)

New

    logits = logits[:, -input_ids.size(1) - 1 : -1, :].to(input_ids.device)

python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml

wandb (main): torchforge/grpo-training/runs/gq9wjime
image

wandb (PR): torchforge/grpo-training/runs/ctirsyoa
image

pytest tests/unit_tests/util/test_compute_logprobs.py

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 26, 2025
@casteryh
Copy link
Contributor

casteryh commented Sep 27, 2025

@Jack-Khuu
Copy link
Contributor Author

Can you revive some of these tests?

Did not realize these tests still existed. yup will revive

@Jack-Khuu Jack-Khuu requested a review from joecummings October 1, 2025 22:58
Copy link
Contributor

@JenniferWang JenniferWang left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@Jack-Khuu Jack-Khuu merged commit 8d84099 into main Oct 2, 2025
5 checks passed
@Jack-Khuu Jack-Khuu deleted the compute_logprobs_helper branch October 2, 2025 20:26
photomz pushed a commit to photomz/forge that referenced this pull request Oct 25, 2025
* Basic test + Move compute_logprobs to util ops

* docstring

* Rebase + add empty response test

* Update math
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants