Adapter-sync LoRA weight sync (supersedes #90) + reload correctness fixes - #94
Draft
qywu wants to merge 6 commits into
Draft
Adapter-sync LoRA weight sync (supersedes #90) + reload correctness fixes#94qywu wants to merge 6 commits into
qywu wants to merge 6 commits into
Conversation
LoRA weight sync no longer folds the adapter into the base model and broadcasts merged full weights. Instead the engine exports the session's adapter in PEFT format and the API server loads it on every registered endpoint via SGLang /load_lora_adapter. Inference endpoints for a LoRA training server are now required to support LoRA: registration rejects endpoints that do not report enable_lora=true or whose max_lora_rank is below the substrate rank. Why: the fold materializes W + (alpha/r) B A on a trainer whose base is already resident, which OOMs at large-model scale (observed live: a Qwen3.5-35B-A3B EP=2 trainer at 70GB+/GPU dies in canonical_lora_fold during sync), and it ships full-model bytes over NCCL for a ~100MB delta. Adapter publication is what the internal large-MoE RL runs already do (merged_weight_sync: false). Notes for review: - QLoRA composites and the DSV4/GLM exact active-LoRA banks are NOT covered by the PEFT export; their guards still fail loudly. - Numerics: endpoint-side LoRA applies base + B(Ax) through SGLang's LoRA kernels rather than the trainer's exact merged-forward fold, which widens train/inference K3 relative to merged sync; ratio-clipped objectives (CISPO/PPO) absorb this, and the Triton LoRA backend is required for the on-policy LM-head contract. - Client counterpart: SamplingClients must target the adapter name; the xorl-client examples currently sample the base model under merged sync.
LoRA sessions do not participate in weight sync at all; adapters publish via /load_lora_adapter and are documented with the LoRA adapter docs. The weight-sync overview now only documents the dense full-weight transport.
execute_sync_inference_weights.build_output whitelists result keys, which silently dropped adapter_sync/adapter_path/model_id — the API layer then saw a generic success and never drove the endpoint-side /load_lora_adapter calls.
The load helper treats 'already loaded' as success, so per-step adapter re-publication under the same name silently no-ops and the sampler serves the first synced adapter forever (off-policy with no error). Unload the name first on both sync paths; the first sync's unload is a tolerated no-op.
Broly Security ScanNote Summary 1 actionable finding(s) in this PR
All actionable items are in the table below. No finding is at or above
Dismiss false positivesTick a box to dismiss the finding; untick it to bring the finding back. That is the same as replying
Note Re-scan this PR anytime with
|
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
1 participant
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
Retargets #90 at
devand adds two correctness fixes found by exercising the path live on a Qwen3.5-35B-A3B EP=2 trainer + TP2 LoRA SGLang sampler (the configuration merged-weight sync structurally cannot serve — it OOMs in the LoRA fold).Base change (from #90): LoRA sessions publish PEFT adapters via
/load_lora_adapterinstead of folding into the base and broadcasting merged full weights; registration rejects endpoints that don't reportenable_lora=trueor whosemax_lora_rankis below the substrate rank; weight-sync docs scope NCCL transport to dense full-weight only.Fix 1 — result whitelist (
request_processor.py):execute_sync_inference_weights.build_outputwhitelists result keys and silently droppedadapter_sync/adapter_path/model_id; the API layer then saw a generic success and never drove the endpoint-side loads. Observed live: engine exported the adapter, client got success, sampler haddict_keys([])loaded.Fix 2 — stale adapter on re-sync (
inference_endpoints.py): the load helper treats SGLang's "already loaded" as success (correct for its original one-shot use). Under per-step syncs that re-publish updated weights under the same name, every sync after the first silently no-ops and the sampler serves the step-1 policy forever — invisibly off-policy training with plausible-looking metrics. Both sync paths now unload the name before re-loading (the first sync's unload is a tolerated no-op).Live validation
End-to-end on 35B-A3B: adapter export (~100MB PEFT) → whitelist passthrough → API load →
Loaded LoRA adapter on http://...:30018in the log; a 64-step RL run is in progress on this stack.Endpoint compatibility contract learned in the shakedown (candidate doc additions)
--lora-target-modulesmust be the explicit projection list;allwraps lm_head, which the exact-contract fp32 decode head rejects (ValueError: ... does not support a LoRA-wrapped lm_head).g_projis not accepted, so qwen3_5 trainers targeting it produce unloadable adapters.--mem-fraction-staticwith--max-total-tokenscapped: the mamba state cache (~61MB/req) is carved from inside the static budget.qywu/wordle-upgrades): syncs carry the sessionmodel_id; LoRA runs publish the zero adapter before step-1 rollouts; generation selects the adapter by name.Open questions
Unchanged from #90: K₃ numerics of runtime-LoRA vs exact merged-forward; QLoRA and DSV4/GLM factor banks excluded (guards intact);
pools-scoped adapter loads.