feat: optimize refit by preparing refit info ahead of time #638

yuki-97 · 2025-07-10T10:00:21Z

Separate the refit process changes from #613.

What does this PR do ?

Update refit process, prepare refit info ahead of time to avoid repeated sending metadata from train side to inference side each time.
Add NRL_REFIT_BUFFER_MEMORY_RATIO and update the default value from 10% to 20%.
Found the dtype of some tensors (e.g. e_score_correction_bias) will change during training, have some special handle with it, and refit_param_info_mcore is not cached for now because of this.

Test Result

convergence

dtensor	mcore	mcore w/ packing (dsv3 w/ 64 tp)

time cost
In mcore w/ packing (dsv3 w/ 64 tp)

Speedup from 420s to 105s with preparing metainfo.
Speedup from 105s to 55s with updating refit buffer memory ratio.

*The ~20s overhead is due to offload.

Refit Process Changes

Colocated

Previous

Each time in refit
1. prepare_weights_for_ipc in train side.
2. get_weights_ipc_handles in train side and update_weights_from_ipc_handles in inference side.

Now

Init
1. prepare_refit_info in train side.
Each time in refit
1. prepare_weights_for_ipc in train side.
2. get_weights_ipc_handles in train side and update_weights_from_ipc_handles in inference side.

Non-colocated

Previous

Init
1. init_collective in both train and inference side.
Each time in refit
1. prepare_info_for_collective in train side.
2. broadcast_weights_for_collective in train side and update_weights_from_collective in inference side.

Now

Init
1. init_collective in both train and inference side.
2. prepare_refit_info in both train and inference side.
Each time in refit
1. broadcast_weights_for_collective in train side and update_weights_from_collective in inference side.

ZhiyuLi-Nvidia · 2025-07-10T17:39:11Z

@yuki-666
just curious, does the refit_info correspond to the policy model (like DTensor or Mcore) or generation model (vllm one)?
At first glance, look like refit_info in the change is relevant to the policy model. We probably need the generation model (vllm one) to avoid duplicated meta data transfer in collective_rpc.
Do you know if we can extract the dtype and shape without actually converting the params (from policy model to the generation one).

yuki-97 · 2025-07-11T22:26:06Z

Do you know if we can extract the dtype and shape without actually converting the params (from policy model to the generation one).

Yup, I added it in ebb874a.

ZhiyuLi-Nvidia

Thank you @yuki-666. LGTM!

nemo_rl/models/generation/vllm_backend.py

Signed-off-by: Yuki Huang <yukih@nvidia.com>

…mcore for speedup Signed-off-by: Yuki Huang <yukih@nvidia.com>

nemo_rl/models/generation/vllm_backend.py

nemo_rl/models/policy/dtensor_policy_worker.py

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Signed-off-by: Yuki Huang <yukih@nvidia.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com> Signed-off-by: Jialei Chen <jialeic@google.com>

)" This reverts commit 8f7d71e

Signed-off-by: Yuki Huang <yukih@nvidia.com>

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com>

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com> Signed-off-by: Qidong Su <qidongs@nvidia.com>

yuki-97 requested review from ZhiyuLi-Nvidia, guyueh1 and yfw July 10, 2025 10:00

yuki-97 force-pushed the yukih/prepare-refit-info branch from 61abcf9 to de56749 Compare July 11, 2025 04:29

yuki-97 added the CI:L1 Run doctests, unit tests, and functional tests label Jul 11, 2025

yuki-97 temporarily deployed to nemo-ci July 11, 2025 04:29 — with GitHub Actions Inactive

yuki-97 force-pushed the yukih/prepare-refit-info branch from de56749 to f57e799 Compare July 11, 2025 06:20

yuki-97 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Jul 11, 2025

yuki-97 temporarily deployed to nemo-ci July 11, 2025 06:21 — with GitHub Actions Inactive

yuki-97 force-pushed the yukih/prepare-refit-info branch from f57e799 to fc4d64e Compare July 11, 2025 07:21

yuki-97 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Jul 11, 2025

yuki-97 temporarily deployed to nemo-ci July 11, 2025 07:21 — with GitHub Actions Inactive

ZhiyuLi-Nvidia previously approved these changes Jul 11, 2025

View reviewed changes

yuki-97 dismissed ZhiyuLi-Nvidia’s stale review via 9c0e833 July 13, 2025 11:00

yuki-97 force-pushed the yukih/prepare-refit-info branch from 9c0e833 to e9b22fd Compare July 13, 2025 11:08

yuki-97 added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels Jul 13, 2025

yuki-97 temporarily deployed to nemo-ci July 13, 2025 11:09 — with GitHub Actions Inactive

yuki-97 changed the title ~~feat: prepare refit info~~ feat: optimize refit by preparing refit info ahead of time Jul 14, 2025

yuki-97 marked this pull request as ready for review July 14, 2025 06:46

yuki-97 requested review from SahilJain314, parthchadha and terrykong July 14, 2025 06:47

ZhiyuLi-Nvidia reviewed Jul 14, 2025

View reviewed changes

nemo_rl/models/generation/vllm_backend.py Show resolved Hide resolved

terrykong added the r0.3.0 Release r0.3.0 label Jul 14, 2025

guyueh1 mentioned this pull request Jul 14, 2025

OOM during refit in mcore path when using larger bucket sizes #630

Open

yuki-97 added 8 commits July 15, 2025 20:39

fix rebase

26af037

Signed-off-by: Yuki Huang <yukih@nvidia.com>

add metainfo in prepare_refit_info

d8e48f3

Signed-off-by: Yuki Huang <yukih@nvidia.com>

some fix and update unit test with prepare_refit_info

eb6b910

Signed-off-by: Yuki Huang <yukih@nvidia.com>

support update dtype record during training

23c5d3a

Signed-off-by: Yuki Huang <yukih@nvidia.com>

not cache refit_param_info since dtype may change

687c868

Signed-off-by: Yuki Huang <yukih@nvidia.com>

add unit test

cabac3b

Signed-off-by: Yuki Huang <yukih@nvidia.com>

rename some vars

a88ca8d

Signed-off-by: Yuki Huang <yukih@nvidia.com>

add NRL_REFIT_BUFFER_MEMORY_RATIO, update default from 10% to 20% in …

f152f61

…mcore for speedup Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 force-pushed the yukih/prepare-refit-info branch from c685241 to f152f61 Compare July 16, 2025 03:40

yfw reviewed Jul 16, 2025

View reviewed changes

nemo_rl/models/generation/vllm_backend.py Show resolved Hide resolved

yuki-97 mentioned this pull request Jul 16, 2025

fix: maintain fp32 mlp.router.expert_bias even with bf16 enabled #674

Closed

4 tasks

yfw reviewed Jul 16, 2025

View reviewed changes

nemo_rl/models/policy/dtensor_policy_worker.py Show resolved Hide resolved

yfw previously approved these changes Jul 16, 2025

View reviewed changes

add assert

63b5e3a

Signed-off-by: Yuki Huang <yukih@nvidia.com>

yuki-97 dismissed yfw’s stale review via 63b5e3a July 16, 2025 04:06

ZhiyuLi-Nvidia approved these changes Jul 16, 2025

View reviewed changes

yfw approved these changes Jul 16, 2025

View reviewed changes

terrykong approved these changes Jul 16, 2025

View reviewed changes

terrykong added this pull request to the merge queue Jul 16, 2025

Merged via the queue into main with commit 8f7d71e Jul 16, 2025
13 of 14 checks passed

terrykong deleted the yukih/prepare-refit-info branch July 16, 2025 21:50

ZhiyuLi-Nvidia pushed a commit that referenced this pull request Jul 21, 2025

feat: optimize refit by preparing refit info ahead of time (#638)

8f7e162

Signed-off-by: Yuki Huang <yukih@nvidia.com> Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

jialei777 pushed a commit to jialei777/nemo-rl that referenced this pull request Jul 23, 2025

feat: optimize refit by preparing refit info ahead of time (NVIDIA-Ne…

00d74b7

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com> Signed-off-by: Jialei Chen <jialeic@google.com>

yaoyu-33 added a commit that referenced this pull request Jul 28, 2025

Revert "feat: optimize refit by preparing refit info ahead of time (#638

6a613b1

)" This reverts commit 8f7d71e

KiddoZhu pushed a commit that referenced this pull request Jul 28, 2025

feat: optimize refit by preparing refit info ahead of time (#638)

bf076fc

Signed-off-by: Yuki Huang <yukih@nvidia.com>

xxman-google pushed a commit to xxman-google/NeMo-RL that referenced this pull request Jul 30, 2025

feat: optimize refit by preparing refit info ahead of time (NVIDIA-Ne…

c4063ae

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com>

FannYYW pushed a commit to xxman-google/NeMo-RL that referenced this pull request Aug 5, 2025

feat: optimize refit by preparing refit info ahead of time (NVIDIA-Ne…

d42ddd8

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com>

FannYYW pushed a commit to xxman-google/NeMo-RL that referenced this pull request Aug 5, 2025

feat: optimize refit by preparing refit info ahead of time (NVIDIA-Ne…

52f0fba

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com>

soodoshll pushed a commit to soodoshll/RL that referenced this pull request Aug 13, 2025

feat: optimize refit by preparing refit info ahead of time (NVIDIA-Ne…

69d3f4b

…Mo#638) Signed-off-by: Yuki Huang <yukih@nvidia.com> Signed-off-by: Qidong Su <qidongs@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: optimize refit by preparing refit info ahead of time #638

feat: optimize refit by preparing refit info ahead of time #638

Uh oh!

yuki-97 commented Jul 10, 2025 •

edited

Loading

Uh oh!

ZhiyuLi-Nvidia commented Jul 10, 2025

Uh oh!

yuki-97 commented Jul 11, 2025

Uh oh!

ZhiyuLi-Nvidia left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

feat: optimize refit by preparing refit info ahead of time #638

feat: optimize refit by preparing refit info ahead of time #638

Uh oh!

Conversation

yuki-97 commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Test Result

Refit Process Changes

Colocated

Non-colocated

Uh oh!

ZhiyuLi-Nvidia commented Jul 10, 2025

Uh oh!

yuki-97 commented Jul 11, 2025

Uh oh!

ZhiyuLi-Nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yuki-97 commented Jul 10, 2025 •

edited

Loading