Skip to content

Optimize rollout-related ut execution time#1463

Merged
YanhuiDua merged 3 commits intoInternLM:mainfrom
YanhuiDua:optim_rollout_ut
Feb 3, 2026
Merged

Optimize rollout-related ut execution time#1463
YanhuiDua merged 3 commits intoInternLM:mainfrom
YanhuiDua:optim_rollout_ut

Conversation

@YanhuiDua
Copy link
Collaborator

@YanhuiDua YanhuiDua commented Jan 29, 2026

优化前 tests/ray运行时间:~25min
image

优化后 tests/ray运行时间:~13min
image

@YanhuiDua YanhuiDua requested a review from CyCle1024 January 29, 2026 12:23
@CyCle1024
Copy link
Collaborator

CyCle1024 commented Jan 30, 2026

@YanhuiDua test_update_weight.py和test_rl_update_weight.py有更加优化的流程写法。
目前:

Update weight rollout: init train -> init rollout with empty init -> update_weight -> rollout
Ref rollout: new init rollout -> rollout

建议:

Ref rollout: init rollout -> rollout
Update weight rollout: init train -> rollout sleep (drop weight and kvcache) -> rollout wakeup (empty weight) -> update_weight -> rollout

建议的方案优势在于少执行一次推理引擎的初始化,缺点在于需要推理引擎在sleep的一定能正确drop weight,依赖sleep实现(对于目前lmdeploy实现这个条件是符合的)

@YanhuiDua YanhuiDua requested a review from CyCle1024 February 3, 2026 06:29
@YanhuiDua
Copy link
Collaborator Author

@YanhuiDua test_update_weight.py和test_rl_update_weight.py有更加优化的流程写法。 目前:

Update weight rollout: init train -> init rollout with empty init -> update_weight -> rollout
Ref rollout: new init rollout -> rollout

建议:

Ref rollout: init rollout -> rollout
Update weight rollout: init train -> rollout sleep (drop weight and kvcache) -> rollout wakeup (empty weight) -> update_weight -> rollout

建议的方案优势在于少执行一次推理引擎的初始化,缺点在于需要推理引擎在sleep的一定能正确drop weight,依赖sleep实现(对于目前lmdeploy实现这个条件是符合的)

done,并且删掉了llm的update_weights的单测,测试VL模型的update_weights,可覆盖该功能

@YanhuiDua YanhuiDua merged commit ac773ee into InternLM:main Feb 3, 2026
4 of 5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants