Skip to content

perf: fuse PPO logprob entropy computation#2144

Merged
zhuzilin merged 1 commit into
mainfrom
zilin/fuse_log_prob_entropy
Jun 29, 2026
Merged

perf: fuse PPO logprob entropy computation#2144
zhuzilin merged 1 commit into
mainfrom
zilin/fuse_log_prob_entropy

Conversation

@zhuzilin

Copy link
Copy Markdown
Contributor

Replace the separate logprob and entropy paths with a fused autograd helper and remove the unused production legacy helpers.

Add CPU torch/gloo coverage plus a 1-GPU legacy Megatron parity test, and register both tests in the PR workflow.

@zhuzilin zhuzilin force-pushed the zilin/fuse_log_prob_entropy branch 3 times, most recently from 573cc83 to 25e8ba4 Compare June 29, 2026 06:37
Replace the separate logprob and entropy paths with a fused autograd helper and remove the unused production legacy helpers.

Add CPU torch/gloo coverage plus a 1-GPU legacy Megatron parity test, and register both tests in the PR workflow.
@zhuzilin zhuzilin force-pushed the zilin/fuse_log_prob_entropy branch from 25e8ba4 to 22538ac Compare June 29, 2026 08:03
@zhuzilin zhuzilin merged commit 6a0ee15 into main Jun 29, 2026
41 checks passed
@zhuzilin zhuzilin deleted the zilin/fuse_log_prob_entropy branch June 29, 2026 08:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant