feat: add checkpoint-engine refit interface and integrate NIXL#2608
feat: add checkpoint-engine refit interface and integrate NIXL#2608HollowMan6 wants to merge 6 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
ebf023f to
e5829e0
Compare
github-actions
Bot
added
the
Documentation
|
Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
6aeadb6 to
cad807c
Compare
6968a6c to
cb745fa
Compare
HollowMan6
added
CI:Lfast
bd5271d to
c251ec6
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ok to test 4876676 |
|
/claude review |
|
/claude review |
|
/ok to test 1ffd083 |
|
/ok to test ed66052 |
|
/ok to test 7c66f90 |
| if self.checkpoint_engine.cleanup_after_load: | ||
| gc.collect() | ||
| torch.cuda.empty_cache() |
There was a problem hiding this comment.
Looks like cleanup_after_load is set false by default
Have we check how much memory reserved if without clean up?
There was a problem hiding this comment.
The backend default is actually cleanup_after_load=true, I set it to false in the performance recipe explicitly just to get the possible best number for refit time.
I'll check how much memory reserved if without clean up.
There was a problem hiding this comment.
Looks like maybe we can just remove this option, this doesn't really clean up much memory
There was a problem hiding this comment.
THis option is now removed, and release_after_refit is introduced for this #2608 (comment)
| try: | ||
| for idx, logprob_dict in enumerate(generation.logprobs): | ||
| if logprob_dict: | ||
| if logprob_dict and idx < len(generated_tokens): |
There was a problem hiding this comment.
do you know why we need this change? Seems not quite relevant to refit.
Is this to fix something?
There was a problem hiding this comment.
Ah I think this is because I mixed up something when I rebased and resolve conflicts of this PR previously, will clean this up
|
/ok to test 705af30 |
| backend_name: UCX | ||
| backend_init_params: | ||
| engine_config: MAX_RMA_RAILS=8 | ||
| device_list: "mlx5_0,mlx5_1,mlx5_2,mlx5_4,mlx5_5,mlx5_6,mlx5_7,mlx5_8" |
There was a problem hiding this comment.
Do we need manually update the device_list when changing the scale, or the current default setting works for most of situation?
There was a problem hiding this comment.
device_list restricts local UCX devices, so node-count changes on homogeneous hardware do not require updating it, we only need to update the list when HCA names or node topology change. Without device_list, UCX discovers devices automatically. NIXL 1.3 defaults MAX_RMA_RAILS to 2, so auto-discovery works but will not match the validated eight-rail performance.
There was a problem hiding this comment.
Also we are skipping mlx5_3 for a reason and shall we add explanations?
There was a problem hiding this comment.
yes this is intentional, this topology is intended for the dfw cluster, mlx5_3 maps to ethernet, while all others map to IB RDMA data fabric. This is cluster specific so I guess we don't need to document cluster specific configs unless you think this is needed.
| checkpoint_engine: | ||
| enabled: true | ||
| backend: nixl | ||
| update_weights_bucket_megabytes: 2048 |
There was a problem hiding this comment.
Emm does this need to be static? can we derive it based on existing gpu memory?
like NRL_REFIT_BUFFER_MEMORY_RATIO. In previous exps the large update_weights_bucket_megabytes the better. If that helpful we can make it adaptive.
There was a problem hiding this comment.
replaced update_weights_bucket_megabytes for update_weights_bucket_memory_ratio
|
Awesome @HollowMan6. Generally looks great to me! |
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
|
/ok to test ce460e5 |
2 participants
What does this PR do ?
Adds checkpoint-engine based non-colocated GRPO refit with a NIXL backend for policy-to-vLLM weight transfer.
Issues
List issues that this PR closes (syntax):
Usage
Enable checkpoint-engine refit for non-colocated generation:
Before your PR is "Ready for review"
Pre checks:
Additional Information