Skip to content

[feat] Add distributed RLFT training framework #6376

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
Open
wants to merge 101 commits into
base: main
Choose a base branch
from
Open

[feat] Add distributed RLFT training framework #6376

wants to merge 101 commits into from

Conversation

YeAnbang
Copy link
Contributor

@YeAnbang YeAnbang commented Aug 5, 2025

📌 Checklist before creating the PR

  • I have created an issue for this PR for traceability
  • The title follows the standard format: [doc/gemini/tensor/...]: A concise description
  • I have added relevant tags if possible for us to better distinguish different PRs
  • I have installed pre-commit: pip install pre-commit && pre-commit install

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

📝 What does this PR do?

Summarize your work here.
if you have any plots/diagrams/screenshots/tables, please attach them here.

💥 Checklist before requesting a review

  • I have linked my PR to an issue (instruction)
  • My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
  • I have performed a self-review of my code
  • I have added thorough tests.
  • I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

  • 🌝 Yes, I do.
  • 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

ver217 and others added 30 commits August 5, 2025 13:59
@ver217 @YeAnbang
[chat] add distributed impl (#6210)
162bb42
@ver217 @YeAnbang
[feature] fit RL style generation (#6213)
7a2d455 
* [feature] fit rl style generation

* [doc] add docstr

* [doc] add docstr
@TongLi3701 @YeAnbang
add reward related function
fa1272f
@TongLi3701 @YeAnbang
add simple grpo
40d6018
@TongLi3701 @YeAnbang
update grpo
1f07b71
@TongLi3701 @YeAnbang
polish
718c4b7
@YeAnbang
modify data loader
b7842f8
@YeAnbang
grpo consumer
5f178a7
@YeAnbang
update loss
9754a11
@YeAnbang
update reward fn
f8899dd
@YeAnbang
update example
5c75d5b
@YeAnbang
update loader
cc4cc78
@YeAnbang
add algo selection
1f15dc7
@YeAnbang
add save
88eb6e5
@YeAnbang
update select algo
246f16d
@pre-commit-ci @YeAnbang
[pre-commit.ci] auto fixes from pre-commit.com hooks
f71d422 
for more information, see https://pre-commit.ci
@pre-commit-ci @YeAnbang
[pre-commit.ci] auto fixes from pre-commit.com hooks
bc538ba 
for more information, see https://pre-commit.ci
@YeAnbang
update grpo
fe017d3
@YeAnbang
update reward fn
c8db826
@YeAnbang
update reward
a537aa1
@YeAnbang
fix reward score
a4862a2
@YeAnbang
add response length
b951d0b
@YeAnbang
detach
69a1a32
@YeAnbang
fix tp bug
b19355f
@YeAnbang
fix consumer
a2ae82a
@YeAnbang
convert to 8 generation
30c7ddd
@YeAnbang
print results
bfc4582
@YeAnbang
setup update
e224673
@YeAnbang
fix transformers backend
35dabd7
@duanjunwen @YeAnbang
[Feature] Support Distributed LogProb for GRPO Training (#6247)
4551853 
* [fix] fix qwen VocabParallelLMHead1D and gather output

* fix tp bug

* fix consumer

* [feat] Support Distributed LogProb for GRPO Training

* [fix] fix loss func

* [fix] fix log prob plugin

* [fix] fix qwen modeling param

* [fix] rm comments

* [fix] rm hard-code;fix non-dist version

* [fix] fix test file param name and benchmark tp gather output=True/False

* [fix] rm non-dist version in dist log prob

* [fix] fix comments

* [fix] fix dis log prob plugin

* [fix] fix test case

* [fix] fix qwen VocabParallelLMHead1D and gather output

* [fix] fix DistLogProb comments

* [fix] restore tp size

* [fix] fix comments

* [fix] fix comment; fix LogSoftmax usage

---------

Co-authored-by: Tong Li <[email protected]>
YeAnbang and others added 21 commits August 5, 2025 14:04
@YeAnbang
fix bug, tested
de40c73
@YeAnbang
remove debug code
9dbb0ff
@pre-commit-ci @YeAnbang
[pre-commit.ci] auto fixes from pre-commit.com hooks
72b2d98 
for more information, see https://pre-commit.ci
@TongLi3701 @YeAnbang
move out evaluation func (#6343)
6ae54a6 
Co-authored-by: Tong Li <[email protected]>
@TongLi3701 @YeAnbang
fix pp memory issue (#6344)
3a4681f 
Co-authored-by: Tong Li <[email protected]>
@YeAnbang
Manually schedule resources and support auto master address assigning
3b3c48d
@YeAnbang
modify readme
6a0b809
@YeAnbang
update readme
79a7b99
@YeAnbang
add ray timeout handling instruction
80c576f
@YeAnbang
Update README.md
73384be
@YeAnbang
fix num_update_per_episode
0f71c79
@YeAnbang
optimize pp log_softmax OOM
a960990
@YeAnbang
implement memory efficient logprob
245c8c2
@YeAnbang
fix small bug
b314da1
@TongLi3701 @YeAnbang
add dp rank for multi-dp (#6351)
685e0bd 
Co-authored-by: Tong Li <[email protected]>
@YeAnbang
[feat[ Support one-behind to reduce bubble time. Add profiling code (#…
594c2c6 
…6353)

* support n_behind, add profiling

* fix bugs

* fix visualization

* fix behind

* fix loop issue

* add profiling

* fix update

* update assert

* remove assert

---------

Co-authored-by: Tong Li <[email protected]>
@YeAnbang
fix code evaluation
352a8e0
@YeAnbang
fix style
eafbc89
@YeAnbang
add entropy (#6363)
3d9dd34
@YeAnbang
hotfix entropy calculation (#6364)
c782976
@YeAnbang @pre-commit-ci
[Fix] Add L2 Regularization (#6372)
118a66f 
* fix no L2 regularization error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
@YeAnbang YeAnbang changed the base branch from main to grpo-latest August 5, 2025 06:11
@YeAnbang YeAnbang force-pushed the grpo-latest-rebase-main branch from fb8202b to 3746f73 Compare August 5, 2025 06:42
@YeAnbang YeAnbang changed the base branch from grpo-latest to main August 5, 2025 06:48
@YeAnbang YeAnbang marked this pull request as ready for review August 5, 2025 09:40
@YeAnbang YeAnbang requested a review from a team as a code owner August 5, 2025 09:40
@YeAnbang YeAnbang requested a review from TongLi3701 August 11, 2025 06:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@TongLi3701 TongLi3701 Awaiting requested review from TongLi3701

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@YeAnbang @ver217 @TongLi3701 @duanjunwen