From 0bf1def15b9b189e0b56f9f9d1753b83f3697dfe Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 04:19:51 +0000 Subject: [PATCH 1/5] Sync LoRA sessions as adapters; drop merged-weight LoRA sync 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. --- .../server-training/weight-sync/overview.mdx | 21 +++-- .../server/api_server/inference_endpoints.py | 81 +++++++++++++++++++ src/xorl/server/weight_sync/handler.py | 54 ++++++++++--- 3 files changed, 139 insertions(+), 17 deletions(-) diff --git a/docs/src/content/docs/server-training/weight-sync/overview.mdx b/docs/src/content/docs/server-training/weight-sync/overview.mdx index 3e6d0612..7793520f 100644 --- a/docs/src/content/docs/server-training/weight-sync/overview.mdx +++ b/docs/src/content/docs/server-training/weight-sync/overview.mdx @@ -181,11 +181,22 @@ Set `quantization` to `null` for BF16 transfer. Online quantization currently ac ## Sync with LoRA -For LoRA training, the sync merges LoRA weights into the base model before broadcasting: -- `W_full = W_base + lora_B @ lora_A × scaling` -- The merged BF16 weight is sent to inference - -The base weights on the training side are **not modified** — LoRA parameters remain separate for continued training. +For LoRA training, the sync publishes the **adapter**, never merged full weights: + +1. The engine exports the session's adapter in PEFT format + (`adapter_model.safetensors` + `adapter_config.json`) under + `/weight_sync_adapters//`. +2. The API server loads it on every registered endpoint via SGLang's + `/load_lora_adapter` (endpoints must run `--enable-lora` with a + `--max-lora-rank` covering the substrate rank; registration rejects + endpoints that do not report LoRA support). +3. Generation requests select the adapter by `model_id`. + +Merged-weight LoRA sync (`W_full = W_base + lora_B @ lora_A × scaling` +broadcast as full weights) was removed: materializing the fold is infeasible +at large-model scale (a 35B-A3B trainer OOMs with the base already resident), +and it shipped full-model bytes for a ~100MB delta. Adapter transfers are +orders of magnitude smaller and keep the sampler's base weights immutable. ## Sync with QLoRA diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index a01669c0..ff35c28a 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -4,6 +4,7 @@ import asyncio import json +import os import logging import socket from pathlib import Path @@ -478,6 +479,7 @@ async def _sync_weights_to_endpoints( group_name: str, buffer_size_mb: int, quantization: dict | None = None, + model_id: str | None = None, ) -> Dict[str, Any]: """ Internal method to sync weights to specific endpoints. @@ -511,6 +513,7 @@ async def _sync_weights_to_endpoints( buffer_size_mb=buffer_size_mb, sync_method=self.sync_inference_method, quantization=quantization, + model_id=model_id, ), ) @@ -525,6 +528,36 @@ async def _sync_weights_to_endpoints( return {"success": False, "message": f"Engine error: {output.error}"} result = output.outputs[0] if output.outputs else {} + if result.get("adapter_sync"): + # LoRA session: the engine exported a PEFT adapter instead of + # entering merged-weight collectives. Load it on every registered + # endpoint via SGLang's /load_lora_adapter (endpoints registered + # against a LoRA trainer are validated to support LoRA). + adapter_path = result.get("adapter_path", "") + lora_name = result.get("model_id") or model_id or "default" + try: + await self._load_lora_on_inference_endpoints(lora_name, adapter_path) + except HTTPException as exc: + return { + "success": False, + "message": f"adapter exported but endpoint load failed: {exc.detail}", + "transfer_time": 0, + "total_bytes": 0, + } + adapter_bytes = 0 + try: + for root, _dirs, files in os.walk(adapter_path): + adapter_bytes += sum(os.path.getsize(os.path.join(root, f)) for f in files) + except OSError: + pass + return { + "success": True, + "message": f"LoRA adapter '{lora_name}' loaded on all endpoints from {adapter_path}", + "transfer_time": result.get("transfer_time", 0), + "total_bytes": adapter_bytes, + "adapter_sync": True, + "adapter_path": adapter_path, + } return { "success": result.get("success", False), "message": result.get("message", ""), @@ -719,6 +752,35 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> endpoint=None, ) + # LoRA trainers publish adapters, never merged weights: every inference + # endpoint must therefore support LoRA (SGLang --enable-lora) so it can + # apply /load_lora_adapter syncs. Reject non-LoRA endpoints up front + # instead of failing at the first weight sync. + if (self.server_lora_config or {}).get("enable_lora"): + endpoint_lora = getattr(server_info, "enable_lora", None) if server_info is not None else None + if endpoint_lora is not True: + return AddInferenceEndpointResponse( + success=False, + message=( + "This training server runs LoRA sessions and syncs adapters via " + "/load_lora_adapter; the endpoint does not report enable_lora=true. " + "Relaunch the inference endpoint with --enable-lora (and a " + "max-lora-rank covering the substrate rank)." + ), + endpoint=None, + ) + max_rank = getattr(server_info, "max_lora_rank", None) + substrate_rank = (self.server_lora_config or {}).get("lora_rank") + if max_rank is not None and substrate_rank and int(max_rank) < int(substrate_rank): + return AddInferenceEndpointResponse( + success=False, + message=( + f"Endpoint max_lora_rank={max_rank} is below the trainer's LoRA " + f"substrate rank {substrate_rank}; raise --max-lora-rank." + ), + endpoint=None, + ) + # Determine world_size: use server_info.tp_size if available, else request.world_size world_size = request.world_size if server_info is not None and server_info.tp_size is not None and server_info.tp_size > 1: @@ -769,6 +831,7 @@ async def add_inference_endpoint(self, request: AddInferenceEndpointRequest) -> group_name=request.group_name, buffer_size_mb=request.buffer_size_mb, quantization=self._get_endpoint_quantization(), + model_id=getattr(request, "model_id", None), ) weights_synced = sync_result.get("success", False) if weights_synced: @@ -1001,6 +1064,24 @@ async def sync_inference_weights(self, request: SyncInferenceWeightsRequest) -> # Extract results result = output.outputs[0] if output.outputs else {} + if result.get("adapter_sync"): + # LoRA session: adapter exported by the engine; perform the + # endpoint-side dynamic loads here (never merged collectives). + adapter_path = result.get("adapter_path", "") + lora_name = result.get("model_id") or request.model_id or "default" + await self._load_lora_on_inference_endpoints(lora_name, adapter_path) + targeted = [ep for ep in self.inference_endpoints if request.pools is None or ep.pool in request.pools] + return SyncInferenceWeightsResponse( + success=True, + message=f"LoRA adapter '{lora_name}' loaded on {len(targeted)} endpoints from {adapter_path}", + endpoints_synced=[ + EndpointSyncResult(host=ep.host, port=ep.port, success=True, message="adapter loaded") + for ep in targeted + ], + transfer_time=float(result.get("transfer_time", 0) or 0), + total_bytes=int(result.get("total_bytes", 0) or 0), + ) + # Build endpoint sync results endpoint_results = [] for ep_result in result.get("endpoint_results", []): diff --git a/src/xorl/server/weight_sync/handler.py b/src/xorl/server/weight_sync/handler.py index e2733cdb..da99e6bf 100644 --- a/src/xorl/server/weight_sync/handler.py +++ b/src/xorl/server/weight_sync/handler.py @@ -466,7 +466,22 @@ def _clear_sync_abort(self, abort_path: str) -> None: except Exception as e: logger.debug("Rank %d: [WeightSync] failed to clear abort marker %s: %s", self.rank, abort_path, e) - def _prepare_lora_adapter_for_sync(self, model_id: Optional[str]) -> Optional[str]: + def _export_adapter_for_sync( + self, model_id: Optional[str], weight_version: Optional[str] + ) -> Optional[Dict[str, Any]]: + """Export the session's LoRA adapter for endpoint-side dynamic loading. + + LoRA weight sync publishes the ADAPTER, never merged full weights: the + old fold-and-broadcast path materialized ``W + (alpha/r) * B @ A`` on + the trainer, which is infeasible at large-model scale (a 35B-A3B + trainer OOMs in the fold with the base already resident) and forces + the sampler to receive full-model bytes for a ~100MB delta. Inference + endpoints are expected to support LoRA (SGLang ``--enable-lora``) and + load the exported PEFT adapter via ``/load_lora_adapter``; the API + layer performs those loads after this export returns. + + Returns None for full-weight trainers (dense sync path unchanged). + """ adapter_manager = getattr(self.trainer, "adapter_manager", None) if adapter_manager is None: return None @@ -475,14 +490,14 @@ def _prepare_lora_adapter_for_sync(self, model_id: Optional[str]) -> Optional[st if contains_dsv4_exact_active_lora_component(model): raise RuntimeError( "DSV4-Flash exact active-LoRA requires complete factor-only adapter publication; " - "legacy merged full-weight synchronization is not admitted. Export all 948 factors " + "the PEFT adapter export does not cover dsv4_expert_banks. Export all 948 factors " "as dsv4_expert_banks and load a fresh sampler adapter version." ) if contains_glm52_exact_active_lora_component(model): raise RuntimeError( "GLM-5.2 exact active-LoRA composites require a complete factor-only adapter synchronization " - "protocol, which this handler does not implement; legacy merged full-weight synchronization is " - "not admitted. Export all 1,700 factors and start a fresh sampler adapter lifecycle." + "protocol, which this handler does not implement. Export all 1,700 factors and start a " + "fresh sampler adapter lifecycle." ) resolved_model_id = model_id or getattr(adapter_manager, "current_adapter_id", None) or "default" @@ -492,13 +507,27 @@ def _prepare_lora_adapter_for_sync(self, model_id: Optional[str]) -> Optional[st raise KeyError(f"LoRA adapter for model_id={resolved_model_id!r} is not registered") register_adapter(resolved_model_id, lr=None) - adapter_manager.sync_weights_to_model(resolved_model_id) + 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) + + # Collective PEFT export (adapter_model.safetensors + adapter_config.json, + # the format SGLang's /load_lora_adapter consumes). All ranks participate. + self.trainer.save_lora_only(export_dir, model_id=resolved_model_id) logger.info( - "Rank %s: [WeightSync] Prepared LoRA adapter model_id=%s for merged weight extraction", + "Rank %s: [WeightSync] Exported LoRA adapter model_id=%s to %s for endpoint-side loading", self.rank, resolved_model_id, + export_dir, ) - return resolved_model_id + return { + "success": True, + "adapter_sync": True, + "model_id": resolved_model_id, + "adapter_path": export_dir, + "message": f"adapter exported to {export_dir}; endpoints load via /load_lora_adapter", + } def _mark_sync_abort(self, abort_path: str, err: Exception) -> None: try: @@ -804,11 +833,12 @@ async def handle_sync_inference_weights(self, command_dict: Dict[str, Any]) -> D ) try: - synced_model_id = self._prepare_lora_adapter_for_sync(model_id) - if synced_model_id is not None: - logger.info( - "Rank %s: [WeightSync] Syncing merged LoRA weights for model_id=%s", self.rank, synced_model_id - ) + adapter_export = self._export_adapter_for_sync(model_id, weight_version) + if adapter_export is not None: + # LoRA sessions never enter merged-weight collectives: the + # adapter is on disk and the API layer drives endpoint-side + # /load_lora_adapter calls. NCCL transport is dense-only. + return adapter_export if sparse_delta_paths: if sync_method != "sparse_delta": return { From 29b9bbcf3d98dbf3a3da5c0d60b496ee509b1c84 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 04:29:03 +0000 Subject: [PATCH 2/5] Weight-sync docs: LoRA is not a weight-sync mode 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. --- .../server-training/weight-sync/overview.mdx | 26 ++++++------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/docs/src/content/docs/server-training/weight-sync/overview.mdx b/docs/src/content/docs/server-training/weight-sync/overview.mdx index 7793520f..a869582d 100644 --- a/docs/src/content/docs/server-training/weight-sync/overview.mdx +++ b/docs/src/content/docs/server-training/weight-sync/overview.mdx @@ -179,24 +179,14 @@ POST /api/v1/set_sync_quantization Set `quantization` to `null` for BF16 transfer. Online quantization currently accepts the Slime/SGLang-compatible block-FP8 E4M3 format with FP32 inverse scales; unsupported INT4, AWQ, compressed-tensors, and fake-quant formats fail before transport starts. The receiver installs the transferred FP8 tensors and scale metadata rather than treating them as BF16 weights. -## Sync with LoRA - -For LoRA training, the sync publishes the **adapter**, never merged full weights: - -1. The engine exports the session's adapter in PEFT format - (`adapter_model.safetensors` + `adapter_config.json`) under - `/weight_sync_adapters//`. -2. The API server loads it on every registered endpoint via SGLang's - `/load_lora_adapter` (endpoints must run `--enable-lora` with a - `--max-lora-rank` covering the substrate rank; registration rejects - endpoints that do not report LoRA support). -3. Generation requests select the adapter by `model_id`. - -Merged-weight LoRA sync (`W_full = W_base + lora_B @ lora_A × scaling` -broadcast as full weights) was removed: materializing the fold is infeasible -at large-model scale (a 35B-A3B trainer OOMs with the base already resident), -and it shipped full-model bytes for a ~100MB delta. Adapter transfers are -orders of magnitude smaller and keep the sampler's base weights immutable. +## LoRA + +LoRA sessions do not use weight sync. Adapters are published to inference +endpoints in PEFT format via SGLang's `/load_lora_adapter`, and endpoints for +a LoRA training server must support LoRA (`--enable-lora`, with +`--max-lora-rank` covering the substrate rank — registration rejects +endpoints that do not). Weight sync's NCCL transport is dense full-weight +only. See the [LoRA adapter docs](/adapters/lora/) for the publication flow. ## Sync with QLoRA From b60c7f3155119b4770d32b6d79b1ba838d27494f Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 04:54:25 +0000 Subject: [PATCH 3/5] Pass adapter-sync fields through the sync result whitelist MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/xorl/server/orchestrator/request_processor.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/xorl/server/orchestrator/request_processor.py b/src/xorl/server/orchestrator/request_processor.py index 4829f3e9..e096ff8e 100644 --- a/src/xorl/server/orchestrator/request_processor.py +++ b/src/xorl/server/orchestrator/request_processor.py @@ -1461,6 +1461,11 @@ def build_output(result): "p2p_rank_summaries": result.get("p2p_rank_summaries", []), "endpoint_results": result.get("endpoint_results", []), "execution_time": result.get("execution_time", 0.0), + # Adapter-sync results (LoRA sessions): the API layer keys + # endpoint-side /load_lora_adapter calls off these fields. + "adapter_sync": result.get("adapter_sync", False), + "adapter_path": result.get("adapter_path", ""), + "model_id": result.get("model_id", ""), } ] From 83c4daebdf090303fcb6a477f5b2f74c3cec7e23 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 05:15:25 +0000 Subject: [PATCH 4/5] Unload before re-loading a synced adapter 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. --- src/xorl/server/api_server/inference_endpoints.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index ff35c28a..e0cae406 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -536,6 +536,13 @@ async def _sync_weights_to_endpoints( adapter_path = result.get("adapter_path", "") lora_name = result.get("model_id") or model_id or "default" try: + # Re-syncs publish UPDATED weights under the same name; the + # load helper treats "already loaded" as success, so a stale + # step-1 adapter would be served forever without this unload. + try: + await self._unload_lora_on_inference_endpoints(lora_name) + except Exception: + pass # first sync: nothing to unload await self._load_lora_on_inference_endpoints(lora_name, adapter_path) except HTTPException as exc: return { @@ -1069,6 +1076,10 @@ async def sync_inference_weights(self, request: SyncInferenceWeightsRequest) -> # endpoint-side dynamic loads here (never merged collectives). adapter_path = result.get("adapter_path", "") lora_name = result.get("model_id") or request.model_id or "default" + try: + await self._unload_lora_on_inference_endpoints(lora_name) + except Exception: + pass # first sync: nothing to unload await self._load_lora_on_inference_endpoints(lora_name, adapter_path) targeted = [ep for ep in self.inference_endpoints if request.pools is None or ep.pool in request.pools] return SyncInferenceWeightsResponse( From 8717722674291579f4e097ad717b7a9f4b8a4fc7 Mon Sep 17 00:00:00 2001 From: Qingyang Wu Date: Thu, 27 Aug 2026 06:27:00 +0000 Subject: [PATCH 5/5] Fix ruff import ordering in inference_endpoints --- src/xorl/server/api_server/inference_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/xorl/server/api_server/inference_endpoints.py b/src/xorl/server/api_server/inference_endpoints.py index e0cae406..40d44e32 100644 --- a/src/xorl/server/api_server/inference_endpoints.py +++ b/src/xorl/server/api_server/inference_endpoints.py @@ -4,8 +4,8 @@ import asyncio import json -import os import logging +import os import socket from pathlib import Path from typing import Any, Dict, List