Sync LoRA sessions as adapters; drop merged-weight LoRA sync - #90
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.
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
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| train_config = getattr(self.trainer, "train_config", {}) or {} | ||
| base_dir = str(train_config.get("output_dir") or "outputs") if isinstance(train_config, dict) else "outputs" | ||
| version_token = _safe_abort_token(weight_version) if weight_version else "latest" | ||
| export_dir = os.path.join(base_dir, "weight_sync_adapters", resolved_model_id, version_token) |
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.
|
Superseded by #94, which retargets this change at |
2 participants
Summary
LoRA weight sync no longer folds the adapter into the base model and broadcasts merged full weights. The engine exports the session's adapter in PEFT format (
adapter_model.safetensors+adapter_config.json, the format SGLang consumes) and the API layer loads it on every registered endpoint via/load_lora_adapter. When LoRA is used, inference endpoints are expected to support LoRA: registration now rejects endpoints that don't reportenable_lora=true, or whosemax_lora_rankis below the trainer's substrate rank.Why remove merged sync
The design was documented in
docs/.../weight-sync/overview.mdx§ Sync with LoRA ("the sync merges LoRA weights into the base model before broadcasting") and implemented inhandler.py::_prepare_lora_adapter_for_sync. Two problems, both observed live:W + (α/r)·B@Aon a trainer whose base is already resident. A Qwen3.5-35B-A3B EP=2 trainer at 70GB+/GPU dies incanonical_lora_foldduring the very first sync (CUDA out of memory ... 512.00 MiBwith <65MB free). Every LoRA run above ~8B on ≤2 GPUs is structurally blocked.Adapter publication is also what the internal large-MoE RL runs already do (
merged_weight_sync: falsein the q30-k3zero run configs).Changes
weight_sync/handler.py:_prepare_lora_adapter_for_sync(fold) →_export_adapter_for_sync(collectivesave_lora_onlyPEFT export to<output_dir>/weight_sync_adapters/<model_id>/<version>); LoRA sessions return before any NCCL transport.api_server/inference_endpoints.py: both sync paths (registration auto-sync and/sync_inference_weights) detect the adapter export and drive the existing_load_lora_on_inference_endpointsbroadcast;model_idis threaded throughSyncWeightsData; registration gains the LoRA capability gate.Open questions for review
Wx + B(Ax)through SGLang's LoRA kernels instead of the trainer's exact merged-forward fold — this widens train↔inference K₃ vs merged sync. Ratio-clipped objectives absorb it; the zero-K₃ contract would need the Triton-LoRA-backend alignment treated explicitly.poolsfiltering needs a targeted variant.model_id); xorl-client examples currently sample the base model under merged-sync assumptions (I can PR that side once this lands).Testing
ast-parse; not yet exercised against a live server (draft). The failure mode this removes is reproducible: any 35B-A3B LoRA trainer on 2 GPUs, firstsync_inference_weights.