fix: keep lora_A/lora_B paired in LoRA IPC weight buckets - #82
Merged
Conversation
sglang merges LoRA adapters per IPC payload, so splitting lora_A and lora_B across update_weight_buffer_size buckets silently skips layers. Group tensors by layer prefix before bucketing and add a GPU test.
Rockdu
reviewed
Jul 29, 2026
Avoid preparing all lora_A/lora_B params on GPU before bucketing, which spiked peak memory by the full adapter size. Bucket using raw state-dict refs first, then _prepare_lora_param only for the bucket being flushed.
Stream layer groups into buckets and call _prepare_lora_param on append, instead of building a fully prepared bucket list before flush. Keeps peak GPU memory bounded to the current bucket while preserving A/B pairing.
Rockdu
reviewed
Jul 29, 2026
| return layer_groups, unmapped_keys, num_lora_keys | ||
|
|
||
|
|
||
| def bucket_lora_layer_groups( |
Collaborator
There was a problem hiding this comment.
This function has become redundant now and can be removed
Remove bucket_lora_layer_groups now that bucketing is inlined in update_weights. Validate each collected layer group matches [layer.lora_A, layer.lora_B] before IPC sync.
Rockdu
approved these changes
Jul 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
lora_Aandlora_Balways ship in the same IPC payloadupdate_weight_buffer_sizeis smalltest_lora_ipc_bucket_keeps_ab_togetherto verify pairing under a tight bufferTest plan
python -m pytest tests/fast-gpu/test_lora_weight_sync.py::test_lora_ipc_bucket_keeps_ab_together -vpython -m pytest tests/fast-gpu/test_lora_weight_sync.py -v