refactor(update-weight): rewrite LoRA weight update path

miles/backends/megatron_utils/hf_export.py

@@ -71,7 +71,7 @@ def export_hf_model_direct(
     weight_map: dict[str, str] = {}
     total_size = 0
     shard_index = 0
-    for hf_named_tensors in iterator.iter_hf_base_weights(megatron_local_weights):
+    for hf_named_tensors in iterator.iter_hf_weights(megatron_local_weights):
         if not is_writer:
             continue
         shard_index += 1

miles/backends/megatron_utils/update_weight/hf_weight_iterator.py

@@ -16,6 +16,7 @@
     WeightUpdatePlacement,
     resolve_placement,
 )
+from miles.utils.lora import is_lora_weight_name
 
 
 class MegatronHfWeightIteratorBase(HfWeightIteratorBase):
@@ -24,13 +25,28 @@ class MegatronHfWeightIteratorBase(HfWeightIteratorBase):
     def _hf_atomic_update_groups(self):
         return get_hf_atomic_update_groups(self.model_name, q_lora_rank=self.args.q_lora_rank)
 
-    def _export_lora_named_tensors(self, adapter):
-        # Both megatron exporters gather TP/EP but not PP.
-        return _gather_pp_full_adapter(self._export_pp_local_lora(adapter))
+    def _iter_hf_adapter_units(self, lora_name, adapter, *, materialize):
+        """Both megatron exporters are PP-local after gathering TP/EP; the PP
+        gather runs only where the resolved placement asks for it."""
+        named_tensors = self._export_pp_local_lora(adapter)
+        if self.placement.gather_pp:
+            named_tensors = _gather_pp_full_adapter(named_tensors)
+        if not materialize:
+            return
+        if not named_tensors:
+            raise RuntimeError(
+                f"LoRA weight sync failed: the adapter export produced zero tensors"
+                f"{f' for adapter {adapter!r}' if adapter is not None else ''}. "
+                "This usually means the Megatron-Bridge or SGLang version is incompatible."
+            )
+        if not any(is_lora_weight_name(name) for name, _tensor in named_tensors):
+            raise RuntimeError("LoRA weight sync failed: the adapter export contains no lora_A/lora_B names.")
+        for hf_name, tensor in named_tensors:
+            yield [(f"{lora_name}:{hf_name}", tensor)]
 
     @abstractmethod
     def _export_pp_local_lora(self, adapter) -> list[tuple[str, torch.Tensor]]:
-        """The adapter's HF-named tensors, TP/EP gathered, PP-local."""
+        """Backend hook: the adapter's HF-named tensors, TP/EP gathered, PP-local."""
 
 
 def get_hf_weight_iterator(

miles/backends/sglang_utils/sglang_engine.py

@@ -565,11 +565,11 @@ def get_weight_version(self):
                 return response.json()["weight_version"]
         response.raise_for_status()
 
-    def unload_lora_adapter(self, lora_name: str):
-        """Unload LoRA adapter."""
+    def register_lora_adapter(self, lora_name: str, config_dict: dict, pinned: bool = True):
+        """Create-or-refresh a LoRA adapter's identity and config (weights zeroed)."""
         return self._make_request(
-            "unload_lora_adapter",
-            {"lora_name": lora_name},
+            "register_lora_adapter",
+            {"lora_name": lora_name, "config_dict": config_dict, "pinned": pinned},
         )
 
     def get_scheduler_actors(self) -> list:
@@ -695,13 +695,15 @@ def continue_generation(self):
         response.raise_for_status()
         return response
 
-    def begin_weight_update(self, selector: str = "all"):
-        """Open a weight-update session on the engine (restores packed weights for loading)."""
-        return self._make_request("begin_weight_update", {"selector": selector})
+    def begin_weight_update(self, selector: str = "all", sync_base: bool = True):
+        """Open a weight-update session on the engine. sync_base=False declares an
+        adapter-only session (no quant unpack; base tensors rejected)."""
+        return self._make_request("begin_weight_update", {"selector": selector, "sync_base": sync_base})
 
-    def end_weight_update(self):
-        """Close the weight-update session (post-load + quant post-process on the full model)."""
-        return self._make_request("end_weight_update", {})
+    def end_weight_update(self, expected_lora_checksums=None):
+        """Close the weight-update session: re-finalize base weights (sync_base
+        sessions) and apply the streamed LoRA stash."""
+        return self._make_request("end_weight_update", {"expected_lora_checksums": expected_lora_checksums})
 
     def update_weight_version(self, weight_version: str):
         return self._make_request(

miles/backends/training_utils/weight_update/broadcast.py

@@ -12,11 +12,7 @@
 from miles.backends.training_utils.parallel import ParallelState
 from miles.backends.training_utils.weight_update.hf_weight_iterator import WeightUpdatePlacement
 from miles.backends.training_utils.weight_update.protocol import WeightTransferProtocol
-from miles.backends.training_utils.weight_update.session import (
-    check_weight_sync_results,
-    unload_lora_adapter,
-    weight_update_selector,
-)
+from miles.backends.training_utils.weight_update.session import weight_update_selector
 from miles.backends.training_utils.weight_update.transfer import derive_replica_position
 from miles.utils.distributed_utils import init_process_group
 
@@ -33,7 +29,6 @@ def __init__(self, args: Namespace) -> None:
         super().__init__(args)
         self._model_update_groups = None
         self._selector = weight_update_selector(args)
-        self._lora_loaded = False
 
     def connect(
         self,
@@ -56,7 +51,6 @@ def connect(
         replica_rank, _ = derive_replica_position(parallel_state, placement)
         self.is_sender = replica_rank == 0
         shard = 0 if placement.gather_pp else parallel_state.pp.rank
-        self.is_lora_sender = self.is_sender and shard == 0
         if self.is_sender:
             self._group_name = f"miles-pp_{shard}"
             disconnect_rollout_engines_from_distributed(
@@ -86,53 +80,6 @@ def send_bucket(self, bucket: list[tuple[str, torch.Tensor]], weight_version: in
             # release must run after both successful and failed broadcasts.
             ray.get(self.rollout_engine_lock.release.remote())
 
-    def send_adapter(
-        self, named_tensors: list[tuple[str, torch.Tensor]], *, lora_name: str, lora_config: dict, upsert: bool
-    ) -> None:
-        """Send adapter metadata over Ray, then broadcast the tensors (src=0).
-
-        Reuses the base broadcast group (``self._model_update_groups`` /
-        ``self._group_name``); base and adapter syncs are strictly sequential, so
-        sharing the NCCL communicator is safe. No CUDA IPC, so it works across
-        nodes: the engine allocates buffers from the metadata and broadcast-receives
-        in order. ``upsert`` maps to the engine's in-place insert-or-overwrite RPC
-        (multi-LoRA slots); without it a stale adapter is unloaded first, since the
-        engine rejects a duplicate name.
-        """
-        if not upsert and self._lora_loaded:
-            unload_lora_adapter(self.rollout_engines, lora_name)
-
-        names = [name for name, _ in named_tensors]
-        dtypes = [param.dtype for _, param in named_tensors]
-        shapes = [list(param.shape) for _, param in named_tensors]
-
-        refs = [
-            engine.load_lora_adapter_from_distributed.remote(
-                lora_name=lora_name,
-                config_dict=lora_config,
-                names=names,
-                dtypes=dtypes,
-                shapes=shapes,
-                group_name=self._group_name,
-                **({"upsert": True} if upsert else {}),
-            )
-            for engine in self.rollout_engines
-        ]
-        # NCCL needs contiguous buffers (lora_B slices are strided); the list keeps them
-        # alive until the async broadcasts complete.
-        contiguous_tensors = [
-            param.data if param.data.is_contiguous() else param.data.contiguous() for _, param in named_tensors
-        ]
-        handles = [
-            dist.broadcast(tensor, 0, group=self._model_update_groups, async_op=True) for tensor in contiguous_tensors
-        ]
-        for handle in handles:
-            handle.wait()
-
-        check_weight_sync_results(ray.get(refs), is_lora=True)
-        if not upsert:
-            self._lora_loaded = True
-
 
 def connect_rollout_engines_from_distributed(
     args: Namespace,

miles/backends/training_utils/weight_update/colocate.py

@@ -1,4 +1,3 @@
-import hashlib
 import logging
 from argparse import Namespace
 from collections.abc import Sequence
@@ -155,7 +154,6 @@ def connect(
         # Every engine-covered rank sends: the per-engine gather_object is
         # collective among the group's members.
         self.is_sender = self._ipc_gather_group is not None
-        self.is_lora_sender = self.is_sender
 
         # A LoRA sync must re-push the frozen base unless the engines keep it
         # across pauses (CPU backup, persistent GPU copy, or remote engines);
@@ -190,66 +188,6 @@ def send_bucket(self, bucket: list[tuple[str, torch.Tensor]], weight_version: in
         check_weight_sync_results(ray.get(refs or []), is_lora=False)
         del long_lived_tensors
 
-    def send_adapter(
-        self, named_tensors: list[tuple[str, torch.Tensor]], *, lora_name: str, lora_config: dict, upsert: bool
-    ) -> None:
-        if upsert:
-            raise NotImplementedError("multi-LoRA weight sync is not supported for colocated engines")
-        if self.use_distribute and self._is_distributed_src_rank:
-            raise NotImplementedError("LoRA weight sync is not yet supported for distributed (non-colocated) engines")
-
-        refs, long_lived_tensors = _send_to_colocated_engine(
-            hf_named_tensors=named_tensors,
-            ipc_engine=self._ipc_engine,
-            ipc_gather_src=self._ipc_gather_src,
-            ipc_gather_group=self._ipc_gather_group,
-            selector=self._selector,
-            lora_config=lora_config,
-            lora_name=lora_name,
-            check_equal=self.args.check_lora_weight_equal,
-            repack_lora_for_ipc=self.args.offload_train,
-        )
-        check_weight_sync_results(ray.get(refs or []), is_lora=True)
-        del long_lived_tensors
-        torch.cuda.ipc_collect()
-        torch.cuda.empty_cache()
-
-
-def _repack_onto_fresh_storage(
-    named_tensors: list[tuple[str, torch.Tensor]],
-) -> dict[str, torch.Tensor]:
-    """Copy CUDA tensors into freshly allocated flat buffers and return views onto them.
-
-    ``torch_memory_saver.region()`` is a ``torch.cuda.use_mem_pool`` context, so anything
-    allocated while building the model -- the LoRA adapter parameters included -- lives in
-    a MemPool the preloaded hook backs with cuMem. cuMem allocations cannot be exported
-    over the legacy CUDA IPC API that ``MultiprocessingSerializer`` uses, so handing their
-    storages straight to the engine fails with "CUDA error: invalid argument" on the first
-    sync. Allocating here, outside any region, gets normal caching-allocator memory whose
-    handles export fine. (This is also why the FlattenedTensorBucket path works: its
-    flattened tensor is likewise allocated at sync time.)
-
-    One buffer per (dtype, device) rather than one clone per tensor: the direct-dict
-    transport relies on the pickler memoizing storages so the engine receives a handful of
-    IPC handles instead of one per adapter tensor.
-    """
-    groups: dict[tuple[torch.dtype, torch.device], list[tuple[str, torch.Tensor]]] = {}
-    for name, tensor in named_tensors:
-        if tensor.is_cuda:
-            groups.setdefault((tensor.dtype, tensor.device), []).append((name, tensor))
-
-    views: dict[str, torch.Tensor] = {}
-    for (dtype, device), items in groups.items():
-        flat = torch.empty(sum(t.numel() for _, t in items), dtype=dtype, device=device)
-        offset = 0
-        for name, tensor in items:
-            view = flat[offset : offset + tensor.numel()].view(tensor.shape)
-            view.copy_(tensor)
-            views[name] = view
-            offset += tensor.numel()
-
-    return {name: views.get(name, tensor) for name, tensor in named_tensors}
-
 
 def _send_to_colocated_engine(
     hf_named_tensors: list[tuple[str, torch.Tensor]],
@@ -258,27 +196,17 @@ def _send_to_colocated_engine(
     ipc_gather_src,
     ipc_gather_group,
     weight_version=None,
-    lora_config: dict | None = None,
-    lora_name: str | None = None,
-    check_equal: bool = False,
     selector: str = "all",
-    repack_lora_for_ipc: bool = False,
 ) -> tuple[list[ObjectRef], Any]:
     # Placeholder ranks (GPU slots reserved but no engine) have no gather group.
     # gather_object is only collective among group members, so we skip entirely.
     if ipc_gather_group is None:
         return [], None
 
-    is_lora = lora_config is not None
     is_gather_src = dist.get_rank() == ipc_gather_src
     long_live_tensors = []
 
-    if is_lora:
-        payload = _repack_onto_fresh_storage(hf_named_tensors) if repack_lora_for_ipc else dict(hf_named_tensors)
-        long_live_tensors.append(payload)
-        converted_named_tensors_by_dtypes = {}
-        serialized_lora = MultiprocessingSerializer.serialize(payload, output_str=True)
-    elif getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
+    if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
         converted_named_tensors_by_dtypes = {"dtype": hf_named_tensors}
     else:
         converted_named_tensors_by_dtypes = {}
@@ -288,7 +216,7 @@ def _send_to_colocated_engine(
                 converted_named_tensors_by_dtypes[dtype] = []
             converted_named_tensors_by_dtypes[dtype].append((name, tensor))
 
-    serialized_tensors: list = [serialized_lora] if is_lora else []
+    serialized_tensors: list = []
     for _dtype, named_tensors in converted_named_tensors_by_dtypes.items():
         flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
         flattened_tensor_data = {
@@ -308,41 +236,14 @@ def _send_to_colocated_engine(
 
     refs = []
     if is_gather_src:
-        if is_lora:
-            try:
-                ray.get(ipc_engine.unload_lora_adapter.remote(lora_name=lora_name))
-            except Exception as _unload_err:
-                logger.debug("lora unload before load skipped: %s", _unload_err)
-
-            expected_checksums = None
-            if check_equal:
-                expected_checksums = {
-                    n: hashlib.sha256(
-                        t.detach().cpu().contiguous().flatten().view(torch.uint8).numpy().tobytes()
-                    ).hexdigest()
-                    for n, t in hf_named_tensors
-                }
-
-            refs.append(
-                ipc_engine.load_lora_adapter_from_tensors.remote(
-                    lora_name=lora_name,
-                    config_dict=lora_config,
-                    serialized_named_tensors=[
-                        per_rank[0] if per_rank else None for per_rank in serialized_named_tensors
-                    ],
-                    expected_checksums=expected_checksums,
-                )
-            )
-
-        else:
-            num_dtypes = len(serialized_named_tensors[0])
-            for i in range(num_dtypes):
-                kwargs = {
-                    "serialized_named_tensors": [tensors[i] for tensors in serialized_named_tensors],
-                    "load_format": "flattened_bucket",
-                    "weight_version": str(weight_version),
-                    "selector": selector,
-                }
-                refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs))
+        num_dtypes = len(serialized_named_tensors[0])
+        for i in range(num_dtypes):
+            kwargs = {
+                "serialized_named_tensors": [tensors[i] for tensors in serialized_named_tensors],
+                "load_format": "flattened_bucket",
+                "weight_version": str(weight_version),
+                "selector": selector,
+            }
+            refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs))
 
     return refs, long_live_tensors

miles/backends/training_utils/weight_update/delta.py

@@ -77,7 +77,6 @@ def connect(
         self._group_name = "miles-disk-delta"
         replica_rank, _ = derive_replica_position(parallel_state, placement)
         self.is_sender = replica_rank == 0
-        self.is_lora_sender = self.is_sender and (placement.gather_pp or parallel_state.pp.rank == 0)
 
     def begin_sync(self, weight_version: int, iter_buckets) -> bool:
         # The first call only captures the baseline snapshot the next sync diffs against.