Skip to content
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

critic speedup #219

Merged
merged 2 commits into from
Jul 11, 2024
Merged

critic speedup #219

merged 2 commits into from
Jul 11, 2024

Conversation

gshennvm
Copy link
Collaborator

@gshennvm gshennvm commented Jun 24, 2024 •
edited
Loading

add some much needed cleanup to the critic and reward model inference servers.

rest of the changes are in the changelog

@gshennvm gshennvm requested review from odelalleau and terrykong June 26, 2024 21:58
odelalleau
odelalleau reviewed Jun 28, 2024
View reviewed changes
Copy link
Collaborator

@odelalleau odelalleau left a comment

Choose a reason for hiding this comment

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

Sorry, couldn't finish the review today. Monday is a holiday in Canada so I'll probably pick it up on Tuesday (unless I can work on it tonight, but not sure yet). Submitting partial review so you can already answer some of my questions.

CHANGELOG.md Outdated Show resolved Hide resolved
CHANGELOG.md Outdated Show resolved Hide resolved
examples/nlp/gpt/conf/gpt_ppo_critic.yaml Outdated Show resolved Hide resolved
examples/nlp/gpt/conf/gpt_ppo_critic.yaml Show resolved Hide resolved
examples/nlp/gpt/conf/gpt_ppo_critic.yaml Show resolved Hide resolved
nemo_aligner/algorithms/critic_server_trainer.py Outdated Show resolved Hide resolved
nemo_aligner/algorithms/critic_server_trainer.py Show resolved Hide resolved
nemo_aligner/servers/server_callables.py Outdated Show resolved Hide resolved
nemo_aligner/servers/server_callables.py Outdated Show resolved Hide resolved
nemo_aligner/servers/server_callables.py Outdated Show resolved Hide resolved
odelalleau
odelalleau reviewed Jul 3, 2024
View reviewed changes
Copy link
Collaborator

@odelalleau odelalleau left a comment

Choose a reason for hiding this comment

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

A few small follow-up comments to my previous review.

nemo_aligner/servers/server_callables.py Outdated Show resolved Hide resolved
examples/nlp/gpt/conf/gpt_ppo_critic.yaml Outdated Show resolved Hide resolved
examples/nlp/gpt/conf/inference_rm.yaml Outdated Show resolved Hide resolved
CHANGELOG.md Outdated Show resolved Hide resolved
odelalleau
odelalleau reviewed Jul 3, 2024
View reviewed changes
Copy link
Collaborator

@odelalleau odelalleau left a comment

Choose a reason for hiding this comment

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

Still not fully done but a few more comments

nemo_aligner/algorithms/critic_server_trainer.py Outdated Show resolved Hide resolved
nemo_aligner/algorithms/critic_server_trainer.py Outdated Show resolved Hide resolved
nemo_aligner/algorithms/critic_server_trainer.py Outdated Show resolved Hide resolved
nemo_aligner/utils/distributed.py Outdated Show resolved Hide resolved
nemo_aligner/utils/distributed.py Outdated Show resolved Hide resolved
odelalleau
odelalleau reviewed Jul 4, 2024
View reviewed changes
Copy link
Collaborator

@odelalleau odelalleau left a comment

Choose a reason for hiding this comment

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

All done, finally!

nemo_aligner/algorithms/reward_server.py Outdated Show resolved Hide resolved
nemo_aligner/algorithms/reward_server.py Outdated Show resolved Hide resolved
nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py Show resolved Hide resolved
nemo_aligner/models/nlp/gpt/megatron_gpt_critic.py Outdated Show resolved Hide resolved
nemo_aligner/models/nlp/gpt/megatron_gpt_reward_model.py Outdated Show resolved Hide resolved
nemo_aligner/utils/distributed.py Outdated Show resolved Hide resolved
nemo_aligner/utils/distributed.py Outdated Show resolved Hide resolved
nemo_aligner/utils/distributed.py Outdated Show resolved Hide resolved
nemo_aligner/utils/distributed.py Outdated Show resolved Hide resolved
nemo_aligner/utils/parallel_state.py Show resolved Hide resolved
@gshennvm gshennvm mentioned this pull request Jul 10, 2024
Open
odelalleau
odelalleau previously approved these changes Jul 11, 2024
View reviewed changes
Copy link
Collaborator

@odelalleau odelalleau left a comment

Choose a reason for hiding this comment

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

Please just fix the typo below and good to go, thanks for this refactor!

CHANGELOG.md Outdated Show resolved Hide resolved
@gshennvm
critic speedup
606f690 
Signed-off-by: Gerald Shen <[email protected]>
@gshennvm gshennvm force-pushed the geshen/critic_speedup branch from 6c8d698 to 606f690 Compare July 11, 2024 07:15
odelalleau
odelalleau previously approved these changes Jul 11, 2024
View reviewed changes
Copy link
Collaborator

@odelalleau odelalleau left a comment

Choose a reason for hiding this comment

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

Looking good, thanks!

@gshennvm
Copy link
Collaborator Author

thanks for the review! I checked the numerics locally against the previous main and ran a nemo generate only test. They both look good to my eye so I'm merging now.

@gshennvm gshennvm merged commit c8190c5 into main Jul 11, 2024
5 checks passed
@gshennvm gshennvm deleted the geshen/critic_speedup branch July 11, 2024 22:30
@AtsunoriFujita AtsunoriFujita mentioned this pull request Oct 18, 2024
Open
abukharin3 pushed a commit to abukharin3/NeMo-Aligner that referenced this pull request Nov 7, 2024
@gshennvm
critic speedup (NVIDIA#219)
78f5731 
* critic speedup

Signed-off-by: Gerald Shen <[email protected]>

* middle of PP should be broadcasted as well

Signed-off-by: Gerald Shen <[email protected]>

---------

Signed-off-by: Gerald Shen <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@odelalleau odelalleau odelalleau approved these changes

@terrykong terrykong Awaiting requested review from terrykong

Assignees
No one assigned
Labels
Algorithms Servers Utils
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@gshennvm @odelalleau