Adapter-sync LoRA weight sync (supersedes #90) + reload correctness fixes

docs/src/content/docs/server-training/weight-sync/overview.mdx

@@ -179,13 +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 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.
+## 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
 

src/xorl/server/api_server/inference_endpoints.py

@@ -5,6 +5,7 @@
 import asyncio
 import json
 import logging
+import os
 import socket
 from pathlib import Path
 from typing import Any, Dict, List
@@ -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,43 @@ 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:
+                # 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 {
+                    "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 +759,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 +838,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 +1071,28 @@ 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"
+                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(
+                    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", []):

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", ""),
                 }
             ]
 

src/xorl/server/weight_sync/handler.py

@@ -466,7 +466,22 @@
         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,30 +490,44 @@
         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"
         if not adapter_manager.has_adapter(resolved_model_id):
             register_adapter = getattr(self.trainer, "register_lora_adapter", None)
             if register_adapter is None:
                 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 @@
         )
 
         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 {