fix(lora): inline adapter tensors into the engine payload

relax/backends/megatron/weight_update/update_weight_from_tensor.py

@@ -25,6 +25,7 @@
     is_lora_adapter_param,
     is_lora_enabled,
     is_lora_merge_mode,
+    serialize_adapter_tensors,
 )
 
 from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer
@@ -440,41 +441,32 @@ def _push_lora_adapter(self, all_params: Mapping[str, torch.Tensor], *, first_sy
 
         config_dict = self._lora_sync.config_dict()
 
-        # SGLang broadcasts this single blob to every TP worker (each slices its own shard), so the
-        # tensors must be host-shared, not CUDA-IPC handles. The default file_descriptor sharing
-        # strategy sends an fd that cannot cross the Ray -> HTTP hops to the server process; switch
-        # to file_system (self-describing /dev/shm filenames in the pickle) around the serialize.
-        from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy
-
+        # The adapter is inlined into the payload rather than shared through host memory:
+        # a /dev/shm reference does not survive the Ray -> HTTP hops to the workers. See
+        # serialize_adapter_tensors for why.
         tensors = {name: t.contiguous() for name, t in full_adapter.items()}
-        prev_strategy = get_sharing_strategy()
-        set_sharing_strategy("file_system")
-        try:
-            serialized = MultiprocessingSerializer.serialize(tensors, output_str=True)
-            t3 = monotonic()
-            if not first_sync:
-                ray.get(self._ipc_engine.unload_lora_adapter.remote(LORA_ADAPTER_NAME))
-            # Keep `tensors` alive across the synchronous load: file_system storages live only while
-            # the producer holds them, and the server maps them during this call.
-            ray.get(
-                self._ipc_engine.load_lora_adapter_from_tensors.remote(
-                    lora_name=LORA_ADAPTER_NAME,
-                    serialized_tensors=serialized,
-                    config_dict=config_dict,
-                    load_format=None,
-                    pinned=False,
-                )
-            )
-            logger.info(
-                "[lora-adapter] tensor push: export=%.2fs gather=%.2fs load=%.2fs (%d tensors, rank=%d)",
-                t_export,
-                t_gather,
-                monotonic() - t3,
-                len(tensors),
-                dist.get_rank(),
+        serialized = serialize_adapter_tensors(tensors)
+        t3 = monotonic()
+        if not first_sync:
+            ray.get(self._ipc_engine.unload_lora_adapter.remote(LORA_ADAPTER_NAME))
+        ray.get(
+            self._ipc_engine.load_lora_adapter_from_tensors.remote(
+                lora_name=LORA_ADAPTER_NAME,
+                serialized_tensors=serialized,
+                config_dict=config_dict,
+                load_format=None,
+                pinned=False,
             )
-        finally:
-            set_sharing_strategy(prev_strategy)
+        )
+        logger.info(
+            "[lora-adapter] tensor push: export=%.2fs gather=%.2fs load=%.2fs (%d tensors, %.1fMB, rank=%d)",
+            t_export,
+            t_gather,
+            monotonic() - t3,
+            len(tensors),
+            len(serialized) / 1e6,
+            dist.get_rank(),
+        )
 
     def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]:
         all_refs = []

relax/utils/megatron_peft_utils.py

@@ -3,6 +3,8 @@
 """Utilities for PEFT (Parameter-Efficient Fine-Tuning) of Megatron in
 Relax."""
 
+import base64
+import pickle
 from typing import Tuple
 
 import torch
@@ -297,6 +299,32 @@ def build_lora_peft(args):
     return create_peft(peft_config)
 
 
+def serialize_adapter_tensors(tensors: dict[str, torch.Tensor]) -> str:
+    """Serialize adapter tensors into a payload that carries its own bytes.
+
+    SGLang broadcasts one blob to every TP worker, so the adapter travels as
+    host memory rather than CUDA IPC handles. Torch's sharing strategies put
+    a reference in the pickle instead of the data, and neither reference
+    survives the Ray -> HTTP hops to the server process. ``file_descriptor``
+    sends an fd the server cannot claim. ``file_system`` sends a ``/dev/shm``
+    path whose storage is reference counted, so the workers that map it first
+    unlink it on the way out, and a straggler rank then opens a file that no
+    longer exists.
+
+    Adapters are small enough to inline: rank 16 on a 30B model is ~24MB. A
+    plain pickle carries the bytes and has no reference to resolve. SGLang
+    reads it unchanged, since ``MultiprocessingSerializer.deserialize``
+    base64-decodes and unpickles.
+
+    Args:
+        tensors: Adapter parameter name -> CPU tensor.
+
+    Returns:
+        A base64 string accepted by SGLang's tensor-load entry points.
+    """
+    return base64.b64encode(pickle.dumps(dict(tensors))).decode("utf-8")
+
+
 __all__ = [
     "LORA_ADAPTER_NAME",
     "count_adapter_parameters",
@@ -308,4 +336,5 @@ def build_lora_peft(args):
     "is_lora_merge_mode",
     "is_lora_adapter_mode",
     "build_lora_peft",
+    "serialize_adapter_tensors",
 ]

tests/utils/test_megatron_peft_utils.py

@@ -14,7 +14,9 @@
 - Base<->adapter param-name prefix round-trip in the fast bridge path
 """
 
+import base64
 import json
+import pickle
 import sys
 import types
 from argparse import Namespace
@@ -29,6 +31,7 @@
     is_lora_adapter_param,
     is_lora_enabled,
     is_lora_merge_mode,
+    serialize_adapter_tensors,
     write_hf_peft_adapter,
 )
 
@@ -242,3 +245,46 @@ def test_base_prefix_strips_grouped_expert_suffix(self):
         base = "decoder.layers.0.mlp.experts.linear_fc1"
         # Grouped experts carry a trailing weight index (weight3) that must be stripped.
         assert _base_param_prefix(f"{base}.to_wrap.weight3") == base
+
+
+class TestSerializeAdapterTensors:
+    """The adapter payload must carry its bytes, not a reference to them.
+
+    Sharing the tensors through host memory instead puts a handle in the
+    pickle, and that handle is only resolvable while the producer holds the
+    storage, which no longer holds once the payload reaches the engine workers.
+    """
+
+    @staticmethod
+    def _adapter():
+        # Large enough that a handle-only payload is unmistakably smaller.
+        return {
+            "layers.0.self_attn.q_proj.lora_A.weight": torch.randn(256, 256),
+            "layers.0.self_attn.q_proj.lora_B.weight": torch.randn(256, 256),
+        }
+
+    def test_round_trips_through_plain_unpickling(self):
+        """SGLang deserializes with base64 + unpickle and nothing else."""
+        tensors = self._adapter()
+
+        restored = pickle.loads(base64.b64decode(serialize_adapter_tensors(tensors), validate=True))
+
+        assert restored.keys() == tensors.keys()
+        for name, tensor in tensors.items():
+            assert torch.equal(restored[name], tensor)
+
+    def test_payload_is_at_least_as_large_as_the_tensors(self):
+        """A payload holding a handle is orders of magnitude smaller."""
+        tensors = self._adapter()
+        nbytes = sum(t.numel() * t.element_size() for t in tensors.values())
+
+        payload = base64.b64decode(serialize_adapter_tensors(tensors), validate=True)
+
+        assert len(payload) >= nbytes
+
+    def test_payload_holds_no_shared_memory_handle(self):
+        """A shared storage serializes as a /torch_<pid> handle, not as
+        bytes."""
+        payload = base64.b64decode(serialize_adapter_tensors(self._adapter()), validate=True)
+
+        assert b"/torch_" not in payload