Pair q_a_proj / kv_a_proj_with_mqa in the bridge weight iterator - #1361
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a pairing mechanism (_pair_q_lora_tensors) for q_a_proj and kv_a_proj_with_mqa tensors to ensure they are processed together for SGLang's deepseek loader. However, the reviewer points out that the existing chunk_named_params_by_size can still split these paired tensors across chunk boundaries if the remaining capacity of a chunk is too small. To prevent weight corruption, the reviewer suggests yielding the paired tensors as grouped tuples and implementing a custom chunker (_chunk_paired_named_params_by_size) that treats the groups as atomic units.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
make sense! Rewired to use the existing `AtomicUpdateGroup` infrastructure (`get_atomic_update_groups` already declares `q_lora_a_proj` for MLA), so this no longer hard-codes the q-lora suffixes and any future model-specific group (e.g. deepseek-v4 once that lands here) will automatically apply on the bridge path the same way it already does for the direct iterator (#1264) and the distributed mixin |
0a9e424 to
2a4e37f
Compare
When ``q_lora_rank`` is set (MLA models like Kimi K2.5 / DeepSeek-V3), SGLang's deepseek loader fuses ``q_a_proj`` and ``kv_a_proj_with_mqa`` into ``fused_qkv_a_proj_with_mqa`` only when both halves are visible in the same ``load_weights`` call. The bridge's HF-name stream + the existing byte-counting chunker can put the two halves in different chunks, leaving the fusion partially applied. Reuse the existing ``AtomicUpdateGroup`` infrastructure (``get_atomic_update_groups``, already declares ``q_lora_a_proj`` for MLA) the way ``hf_weight_iterator_direct.py`` and ``update_weight_from_distributed/mixin.py`` already do. Because those non-streaming consumers (``get_named_update_units`` / ``get_named_value_update_units``) require the full param-name list upfront — which would mean materializing 1T+ base weights — add a small *streaming* buffer + group-aware chunker local to the bridge iterator. The data model is unchanged: ``AtomicUpdateGroup`` is the single source of truth for both paths, so any future model-specific group (e.g. deepseek-v4) automatically applies to the bridge path too. The thread carries ``(hf_name, weight, megatron_name)`` triples through ``_postprocess_and_quantize`` and the lora/base filter so the buffer can suffix-match on the megatron name (same key the existing groups are declared against); the chunker then drops the megatron name and yields ``(hf_name, weight)`` pairs as before. Mirrors the approach slime adopted in THUDM/slime#1753.
2a4e37f to
df7a4ab
Compare
The persistent cached_a_proj and unconditional .detach().clone() were a workaround for q_a_proj / kv_a_proj_with_mqa landing in different chunks during chunked weight updates. The proper fix is on the sender side: pair the two halves before chunking so they always arrive in the same load_weights call. That fix is in radixark/miles#1361 (mirrors THUDM/slime#1532). With sender-side pairing in place, the cached_a_proj only ever needs to live for one load_weights call, matching the original upstream behavior.
2 participants
Summary
q_lora_rankis set (MLA models like Kimi K2.5 / DeepSeek-V3), SGLang's deepseek loader fusesq_a_projandkv_a_proj_with_mqainto a singlefused_qkv_a_proj_with_mqaweight, but only when both halves arrive in the sameload_weightscall.chunk_named_params_by_sizestreams the bridge's HF-name output in iterator order, which can placeq_a_projandkv_a_proj_with_mqafor the same layer in different chunks — leaving the fusion partially applied and corrupting that layer's weights on the rollout engine.hf_weight_iterator_bridge.py, mirroring the approach slime adopted in THUDM/slime#1753: hold each half until its partner arrives, then yield both adjacent so the chunker keeps them together. The pairing covers both the base path (q_a_proj.weight) and the LoRA path (q_a_proj.lora_{A,B}.weight).cached_a_projchange from support kimi 2.5/6 lora (logprob diff exist) sgl-project/sglang#25141 on thesglang-milesbranch).Notes
hf_weight_iterator_direct.pypath already keeps the pair together via the existingAtomicUpdateGroup(_get_q_lora_atomic_update_groups), so it's only the bridge path that needed the fix.hf_weight_iterator_bridge.py; it's gated onargs.q_lora_rank is not None, so non-MLA models are unaffected.Test plan
train/train_rollout_logprob_abs_diffstays bounded (~0.03) across steps, without the persistent-cache patch indeepseek_weight_loader.py.q_lora_rank is None.