[LoRA] DSA indexer LoRA targets for GLM-5.1 / DeepSeek-V3.2-family (cherry-pick #28110) - #28703
Closed
yushengsu-thu wants to merge 46 commits into
Closed
yushengsu-thu wants to merge 46 commits into
yushengsu-thu wants to merge 46 commits into
Conversation
…t-buffer registry on non-cuda devices (sgl-project#27549) (sgl-project#27711) Co-authored-by: Zaili Wang <109502517+ZailiWang@users.noreply.github.com>
…nit (sgl-project#27671) (sgl-project#27713) Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com> Co-authored-by: lmzheng <lmzheng@fb.com>
…st (sgl-project#27688) (sgl-project#27743) Co-authored-by: Ke Bao <ispobaoke@gmail.com>
…dSharedPrefixDataset.from_args for in-process callers (sgl-project#27580) (sgl-project#27754) Co-authored-by: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
…ure sizes (sgl-project#27659) (sgl-project#27756) Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
…oss capture sizes" (sgl-project#27758) (sgl-project#27760) Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
…on backend (sgl-project#26791) (sgl-project#27715) Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
…htly tests (sgl-project#27838) (sgl-project#27847) Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
…OOB on underfilled (tiny) prefill replay (sgl-project#27747) (sgl-project#27864) Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>
…e-back load_back / eviction crashes (sgl-project#27108) (sgl-project#27865) Co-authored-by: Yongji Wu <30348494+libertyeagle@users.noreply.github.com> Co-authored-by: hzh0425 <hzh0425@apache.org>
…HiCache via UnifiedTree by default. (sgl-project#27759) (sgl-project#27878) Co-authored-by: Zhangheng <hzh0425@apache.org>
…rved GPU buffers)
The base-model checksum/snapshot/compare path never saw LoRA weights -- they live in
model_runner.lora_manager, not model.named_parameters() -- so a broken LoRA push or a
stale served buffer was invisible to the weight checker.
Adds a NEW, INDEPENDENT action; the existing snapshot / reset_tensors / compare /
checksum actions are left unchanged (base-model only, byte-for-byte as before):
- `lora_checksum` action -> _compute_lora_checksum(): LoRA-only checksum, hashing just
the adapter weights and skipping the base model. Deliberately independent of the
base-model hash -- gpu_tensor_hash currently hits a CUDA illegal-memory-access on
some GLM base tensors, so a combined base+LoRA checksum would die on the base hash;
hashing the LoRA buffers themselves on GPU is fine (verified on GLM MoE/MLA).
- _lora_named_tensors(): enumerates every LoRA tensor the rollout engine holds, keyed
by the STABLE lora_name (not the per-load uuid), under two views:
lora_cpu:: -- the delivered CPU adapter (lora_manager.loras)
lora_gpu:: -- the served GPU buffer slice the forward actually reads
Covers every module family (attn / MLP / MoE / MLA / DSA-indexer) plus the separate
embedding / lm_head / added-token buffers. Fully defensive (returns {} if LoRA is off
/ internals missing; never raises on a serving engine).
- _hash_resilient(): hashes on the GPU triton path (the SAME path real serving uses)
and falls back to CPU hashlib only on a recoverable exception. (A CUDA IMA is not
recoverable, but LoRA GPU hashing is verified not to IMA.)
Verified on GLM-5.2_5layer: lora_checksum returns LoRA entries (lora_cpu + lora_gpu)
for attn-unfused, attn-fused, moe-unfused and moe-fused target sets, all on the GPU
triton path (no IMA, CPU fallback never triggered). Also verified on Qwen2.5-0.5B.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
yushengsu-thu
force-pushed
the
sglang-miles-glm-dev
branch
from
June 30, 2026 08:48
80e8ba9 to
fe43763
Compare
…CPU + served GPU buffers)" This reverts commit fe43763. The lora_checksum action was added solely to back the miles --check-lora-weight-update-equal served-buffer check, which is being reverted: counting served lora_gpu:: buffers is not a reliable serving-correctness signal. In colocate the buffer is torn down by offload/train/resume between the rollout and the sync (false positive: lora_gpu=0 even when serving LoRA correctly, train_rollout_kl=1.16e-4), and a stale resident adapter could leave it non-empty while serving the wrong weights (false pass). Value-level comparison / train_rollout_kl is the right signal. The free_lora mem-pool fix (parent commit 0a2bcf5) is unaffected and retained. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GZ6ujQSTDHBSgWxR3tYbTP
…keys left uninitialized)
Under --experts-shared-outer-loras, one side of a MoE projection is a single
shared 3D tensor (expert_dim=1) while the other side is per-expert (one 2D
tensor per expert): gate_up has shared lora_A + per-expert lora_B, down has
the opposite.
In load_lora_weight_to_buffer the per-expert branch initialized all four
staging dicts (temp_{A,B}_buffer / temp_{A,B}_cache_keys) only when
`temp_A_buffer[module] is None`. When A is the shared side, the dim()==3
branch sets temp_A_buffer to a Tensor first, so that guard is False and the
per-expert B side's temp_B_cache_keys is never created. The load then hits
`TypeError: 'NoneType' object does not support item assignment` and the
sglang rollout scheduler dies on the first forward.
Fix: initialize each side's buffer and cache_keys independently, guarded by
an assert that a projection side is not simultaneously a shared tensor and
per-expert. Validated: full GLM-5.2 (744B) attention + MoE LoRA rollout now
serves end to end.
Co-authored-by: Cursor <cursoragent@cursor.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This cherry-pick extends SGLang’s LoRA plumbing to support DSA “lightning indexer” projection modules (wq_b, wk, weights_proj) used by GLM-5.1 / DeepSeek-V3.2-family, and includes several CUDA-graph safety fixes for MoE-LoRA routing buffers and adapter reload behavior.
Changes:
- Add DSA indexer projections as first-class LoRA targets (normalization + auto-detect + target matching).
- Fix/adjust CUDA-graph safety for MoE-LoRA routing buffers (incl. DP-attention gathered-token sizing) and avoid trimming routing buffers that can cause OOB reads under capture.
- Improve LoRA memory-pool behavior for unload/reload and correct a shared-outer expert-weight loading edge case.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| python/sglang/srt/utils/common.py | Extends CLI --lora-target-modules choices to include indexer-related module names. |
| python/sglang/srt/lora/utils.py | Registers/normalizes DSA indexer LoRA target names and updates auto-detection logic to avoid ambiguous leaf-name collisions. |
| python/sglang/srt/lora/triton_ops/virtual_experts.py | Prevents unsafe trimming of MoE routing buffers that can cause unmasked OOB reads during CUDA-graph capture/replay. |
| python/sglang/srt/lora/mem_pool.py | Adds buffer-slot freeing on unload, fixes a column-parallel replicated-output edge case, and hardens shared-outer per-expert weight loading. |
| python/sglang/srt/lora/lora_manager.py | Frees memory-pool slots on adapter unload and broadens LoRA module matching to include parent-qualified targets. |
| python/sglang/srt/lora/backend/base_backend.py | Sizes MoE-LoRA per-token routing buffers for DP-gathered token counts and adjusts token→adapter mapping logic. |
| python/sglang/srt/layers/attention/dsa/dsa_indexer.py | Routes weights_proj through an eager module call when LoRA-wrapped and explicitly blocks incompatible piecewise CUDA-graph mode. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+3673
to
+3675
| "wq_b", | ||
| "wk", | ||
| "weights_proj", |
Comment on lines
+336
to
+339
| if moe_num_tokens > num_tokens and num_tokens > 0: | ||
| token_lora_mapping[num_tokens:moe_num_tokens].copy_( | ||
| token_lora_mapping[num_tokens - 1] | ||
| ) |
Comment on lines
+264
to
+268
| # Free the memory-pool buffer slot too, so a later load (colocate RL pushes a fresh uid | ||
| # every step) re-copies the new weights into a cleanly-freed slot instead of leaving a | ||
| # dangling uid_to_buffer_id entry that makes the reload skip the in-place buffer copy -> | ||
| # served (cuda-graph) buffer would keep stale weights. Cuda-graph-replay-safe. | ||
| self.memory_pool.free_lora(lora_ref.lora_id) |
Comment on lines
+881
to
+885
| Per-step LoRA weight refresh in colocate RL pushes a FRESH uid every step (unload + load). | ||
| Without freeing the slot here, unload leaves ``uid_to_buffer_id`` pointing at the unloaded | ||
| adapter, so the next load's in-place buffer copy in ``prepare_lora_batch`` is skipped (the | ||
| eviction self-heal of a dangling entry is fragile) and the SERVED (cuda-graph) buffer keeps | ||
| stale weights. Freeing the slot makes the next load re-copy the new weights into the SAME |
yueming-yuan
force-pushed
the
sglang-miles
branch
from
July 7, 2026 06:43
f8cfad3 to
fbc2623
Compare
yueming-yuan
requested review from
ByronHsu,
DarkSharpness,
HydraQYH,
ShangmingCai,
alphabetc1,
celve,
hanming-lu,
hebiao064,
hzh0425,
iforgetmyname,
key4ng,
mickqian,
ping1jing2,
sogalin,
whybeyoung,
yeahdongcn,
yhyang201,
yizhang2077 and
yuan-luo
as code owners
July 7, 2026 06:43
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
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.
Cherry-picks upstream PR #28110 (
[LoRA] Support DSA indexer LoRA targets for GLM-5.1 / DeepSeek-V3.2-family models) ontosglang-miles.Why
Enables LoRA on the DSA "lightning indexer" modules (
wq_b/wk/weights_proj) for GLM-5.1 / DeepSeek-V3.2-family models, so the rollout engine can serve indexer LoRA adapters. Required by the miles GLM-5.1 GRPO LoRA training path.Verified
Used as the rollout side of the miles GLM-5.1 6-layer GRPO LoRA e2e (4×H200, colocate, rollout tp=2): LoRA adapter loads + syncs (
load_lora_adapter_from_tensors,update_weights_from_tensor), rollout generates, full run reachesTRAIN EXIT 0and saves a PEFT adapter. cuda graph stays enabled on the v0.5.13 / DSA base.Single commit, rebased cleanly onto the current
sglang-miles.🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #28504399031
Latest PR Test (Extra): ❌ Run #28504398957