Skip to content

refactor(update-weight): LoRA adapter tensors ride the base weight stream - #2756

Open
yueming-yuan wants to merge 32 commits into
yueming/hf-weight-iterator-movefrom
yueming/lora-weight-stream
Open

refactor(update-weight): LoRA adapter tensors ride the base weight stream#2756
yueming-yuan wants to merge 32 commits into
yueming/hf-weight-iterator-movefrom
yueming/lora-weight-stream

Conversation

@yueming-yuan

@yueming-yuan yueming-yuan commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator

ci-sglang-pr: #36384

Part of #1360 Step 6 (design: lora_weight_sync_cleanup.md). Stacked on #2754. Engine counterpart: sgl-project/sglang#36384miles tracks the sglang-miles tip continuously, so this PR must merge only after that change is in sglang-miles (and sglang-miles must keep the legacy routes until this merges).

What

The second LoRA send path is deleted; adapter tensors are ordinary entries in the bucketed stream, and adapter identity moves to a control plane.

  • Iterator: iter_hf_base_weightsiter_hf_weights(weights, *, include_base=True, adapters=(), materialize=True). Adapter units join the same unit stream under {lora_name}:{hf_key} names via the new _iter_hf_adapter_units hook. The megatron base implements it as the existing PP-local export hook (_export_pp_local_lora, unchanged in bridge/direct) + the PP gather now placement-gated (runs for gather_pp placements — bridge, colocate; disappears for direct+dist) instead of unconditional. The zero-export / no-lora-names guards move into the hook.
  • Updater: one loop. _adapters_to_push() (single: miles_lora; multi: __miles_slot_{i} sorted) + eager register_lora_adapter for new tenants before begin (engine validates r ≤ max_lora_rank before any bytes move); the frame is unconditional again with the scope passed as begin_weight_update(..., sync_base=...) — reverting the refactor(update-weight): colocated transfer as a protocol on the unified updater #2753 call gating whose substance (no quant round-trip on adapter-only rounds) now lives in the session scope. --check-lora-weight-equal becomes a driver-side sha256 manifest accumulated during the stream and verified engine-side at end_weight_update (now on both transports; previously colocate-only).
  • Deletions (net −326): send_adapter from the ABC + both implementations, the is_lora_sender role from all five protocols, get_hf_lora_weights / _export_lora_named_tensors, _send_lora_adapter / _send_multi_lora_adapters, unload-before-reload (+ session.unload_lora_adapter), the colocate LoRA IPC branch and the _repack_onto_fresh_storage cuMem workaround (adapter tensors now travel in sync-time-allocated FlattenedTensorBuckets, which is exactly why the base path never had the problem).

Behavior changes

  1. Single-LoRA becomes upsert-in-place (no unload/reload window; engine-side rollback on failure).
  2. Colocate multi-LoRA works (previously NotImplementedError).
  3. LoRA runs stamp weight_version (frame is unconditional); the actor-side CI version assertion can be re-enabled in a follow-up.
  4. Adapter placement follows base placement (PP-sharded senders under direct+dist; bridge unchanged).
  5. supports_lora remains the gate for p2p/rdt/disk-delta (unchanged: unsupported).

Tests: iterator template tests retargeted to adapter units in the stream; transport-era send_adapter/orchestration test classes removed (their engine-side successors live in sglang's test_lora_stream_update.py).

Adapter units join iter_hf_weights (renamed from iter_hf_base_weights) as
'{lora_name}:{hf_key}' entries following the resolved placement; the updater
registers adapters via the new control-plane RPC, streams one unified bucket
loop, and the engines apply the stash at end_weight_update (sync_base scope,
optional sha256 manifest for --check-lora-weight-equal). Deletes the second
send path: send_adapter, is_lora_sender, the full-adapter PP gather surface,
unload-before-reload, and the colocate cuMem repack workaround.

Requires sgl-project/sglang#36384 (register RPC + session scope + stash).

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

Tip: disable this comment in your organization's Code Review settings.

@yueming-yuan yueming-yuan changed the title refactor(lora): adapter tensors ride the base weight stream refactor(update-weight): LoRA adapter tensors ride the base weight stream Aug 26, 2026
# Conflicts:
#	miles/backends/megatron_utils/update_weight/hf_weight_iterator.py
# Conflicts:
#	miles/backends/training_utils/weight_update/hf_weight_iterator/__init__.py
# Conflicts:
#	miles/backends/training_utils/weight_update/protocol.py
#	miles/backends/training_utils/weight_update/updater.py
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