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

@@ -18,6 +18,7 @@
     resolve_placement,
 )
 from miles.backends.training_utils.weight_update.hf_weight_iterator.atomic_groups import get_hf_atomic_update_groups
+from miles.utils.lora import is_lora_weight_name
 
 logger = logging.getLogger(__name__)
 
@@ -34,13 +35,30 @@ def __init__(self, *args, **kwargs):
     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)
+        # TODO: the PP-local branch is unreachable until actor.py lifts its bridge-only guard
+        # for distributed LoRA; add an e2e for native-LoRA disaggregate when it does
+        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

@@ -566,11 +566,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 = False):
+        """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:
@@ -696,13 +696,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/hf_weight_iterator/__init__.py

@@ -1,9 +1,10 @@
 """Backend-neutral API for streaming training-side weights as HF-named tensors."""
 
 import dataclasses
+import itertools
 from abc import ABC, abstractmethod
 from argparse import Namespace
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator, Mapping, Sequence
 from typing import ClassVar
 
 import torch
@@ -13,7 +14,6 @@
     assemble_atomic_update_groups,
     pack_units_by_size,
 )
-from miles.utils.lora import is_lora_weight_name
 
 
 @dataclasses.dataclass(frozen=True)
@@ -71,21 +71,28 @@ def __init__(
         self.model_name = model_name
         self.quantization_config = quantization_config
 
-    def iter_hf_base_weights(
+    def iter_hf_weights(
         self,
         weights: Mapping[str, torch.Tensor] | None,
         *,
+        include_base: bool = True,
+        adapters: Sequence[tuple[str, object]] = (),
         materialize: bool = True,
     ) -> Iterator[list[tuple[str, torch.Tensor]]]:
-        """Base model weights as size-bounded buckets of HF-named GPU tensors;
+        """Model weights as size-bounded buckets of HF-named GPU tensors;
         atomic update groups are never split across buckets.
 
         ``weights``: backend-native named weights to read; None reads the live
-        model parameters. ``materialize=False`` joins every collective but
-        yields nothing.
+        model parameters. ``adapters``: ``(lora_name, adapter_or_None)`` pairs
+        whose tensors join the stream under ``{lora_name}:{hf_key}`` names.
+        ``materialize=False`` joins every collective but yields nothing.
         """
-        hf_param_units = self._iter_hf_param_units(weights, materialize=materialize)
-        atomic_update_groups = self._hf_atomic_update_groups() if materialize else []
+        hf_param_units = self._iter_hf_param_units(weights, materialize=materialize) if include_base else iter(())
+        for lora_name, adapter in adapters:
+            hf_param_units = itertools.chain(
+                hf_param_units, self._iter_hf_adapter_units(lora_name, adapter, materialize=materialize)
+            )
+        atomic_update_groups = self._hf_atomic_update_groups() if include_base and materialize else []
         hf_param_units = assemble_atomic_update_groups(hf_param_units, atomic_update_groups)
         yield from pack_units_by_size(hf_param_units, self.args.update_weight_buffer_size)
 
@@ -105,27 +112,11 @@ def _hf_atomic_update_groups(self) -> list[AtomicUpdateGroup]:
         """Backend hook: HF-namespace atomic groups for this model. Default none."""
         return []
 
-    def get_hf_lora_weights(self, adapter=None) -> list[tuple[str, torch.Tensor]]:
-        """The complete adapter in HF PEFT naming (lora_A/lora_B), fully
-        gathered, as one list. ``adapter`` selects a multi-LoRA slot; None is
-        the single-LoRA adapter. Collective.
-        """
-        named_tensors = self._export_lora_named_tensors(adapter)
-        if not named_tensors:
-            raise RuntimeError(
-                f"LoRA weight sync failed: the weight iterator produced zero chunks"
-                f"{f' for adapter {adapter!r}' if adapter is not None else ''}. "
-                "No adapter weights were sent to the rollout engine. 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: chunk contains no LoRA weights "
-                "(no lora_A/lora_B names found). Check weight iterator configuration."
-            )
-        return named_tensors
-
     @abstractmethod
-    def _export_lora_named_tensors(self, adapter) -> list[tuple[str, torch.Tensor]]:
-        """Backend hook: the complete adapter as HF-named tensors, fully
-        gathered along every parallel dim."""
+    def _iter_hf_adapter_units(
+        self, lora_name: str, adapter, *, materialize: bool
+    ) -> Iterator[list[tuple[str, torch.Tensor]]]:
+        """Backend hook: this rank's slice of the adapter per ``self.placement``,
+        one unit per parameter, names ``{lora_name}:{hf_key}``, rank-trimmed.
+        Collectives must run lockstep on every rank; ``materialize=False`` joins
+        them but yields nothing."""

miles/backends/training_utils/weight_update/protocol.py

@@ -15,9 +15,10 @@
 class WeightTransferProtocol(ABC):
     """Moves HF-named weight buckets from training ranks to rollout engines.
 
-    ``connect`` makes every pairing decision once: it sets ``is_sender``,
-    ``is_lora_sender``, and whatever send channels the protocol needs. The
-    updater then drives ``send_bucket`` on sender ranks only.
+    ``connect`` makes every pairing decision once: it sets ``is_sender`` and
+    whatever send channels the protocol needs. The updater then drives
+    ``send_bucket`` on sender ranks only; streamed adapter tensors are ordinary
+    bucket entries (``{lora_name}:{hf_key}`` names).
     """
 
     required_placement: ClassVar[WeightUpdatePlacement] = WeightUpdatePlacement(gather_pp=False)
@@ -32,7 +33,6 @@ def __init__(self, args: Namespace) -> None:
         self.is_sender: bool | None = None
         self.group_name = "miles"
         self.update_weight_metrics: dict[str, float] = {}
-        self.is_lora_sender = False
 
     @abstractmethod
     def connect(
@@ -56,17 +56,14 @@ def begin_sync(
         return True
 
     @abstractmethod
-    def send_bucket(self, bucket: list[tuple[str, torch.Tensor]], weight_version: int) -> None: ...
+    def send_bucket(self, bucket: list[tuple[str, torch.Tensor]]) -> None: ...
 
     def after_base_weights(self) -> None:  # noqa: B027 — optional hook
         """Hook after the base-weight stream completes (e.g. await in-flight writes)."""
 
     def finalize(self, weight_version: int) -> None:  # noqa: B027 — optional hook
         """Hook after all sends (e.g. publish + engine reload)."""
 
-    def send_adapter(self, named_tensors, *, lora_name: str, lora_config: dict, upsert: bool) -> None:
-        raise NotImplementedError(f"{type(self).__name__} does not support LoRA weight sync")
-
     def is_fresh(self) -> bool:
         return self.rollout_engines is not None and not self._connection_stale
 

miles/backends/training_utils/weight_update/protocols/broadcast.py

@@ -12,7 +12,6 @@
 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
 from miles.backends.training_utils.weight_update.utils import get_data_replica_rank_and_size
 from miles.utils.distributed_utils import init_process_group
 
@@ -28,7 +27,6 @@ class UpdateWeightFromDistributed(WeightTransferProtocol):
     def __init__(self, args: Namespace) -> None:
         super().__init__(args)
         self._model_update_groups = None
-        self._lora_loaded = False
 
     def connect(
         self,
@@ -53,7 +51,6 @@ def connect(
         replica_rank, _ = get_data_replica_rank_and_size(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(
@@ -63,15 +60,14 @@ def connect(
                 self.args, self.group_name, rollout_engines
             )
 
-    def send_bucket(self, bucket: list[tuple[str, torch.Tensor]], weight_version: int) -> None:
+    def send_bucket(self, bucket: list[tuple[str, torch.Tensor]]) -> None:
         """Serialize NCCL broadcasts and always release the rollout lock."""
         while not ray.get(self.rollout_engine_lock.acquire.remote()):
             time.sleep(0.1)
         try:
             refs = update_weights_from_distributed(
                 self.group_name,
                 self._model_update_groups,
-                weight_version,
                 self.rollout_engines,
                 bucket,
                 selector=self._selector,
@@ -83,53 +79,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,
@@ -192,7 +141,6 @@ def disconnect_rollout_engines_from_distributed(args, group_name, model_update_g
 def update_weights_from_distributed(
     group_name: str,
     group: dist.ProcessGroup,
-    weight_version: int,
     rollout_engines: Sequence[ActorHandle],
     converted_named_tensors: Sequence[tuple[str, torch.Tensor]],
     selector: str = "all",
@@ -207,7 +155,6 @@ def update_weights_from_distributed(
             shapes=[param.shape for _, param in converted_named_tensors],
             selector=selector,
             group_name=group_name,
-            weight_version=str(weight_version),
         )
         for engine in rollout_engines
     ]