From 1ceb779b9ae094a70c7437b9e1bca60dd3436d3f Mon Sep 17 00:00:00 2001 From: SakaiXue6666 <2441789115@qq.com> Date: Wed, 12 Aug 2026 23:21:15 +0800 Subject: [PATCH] fix(lora): inline adapter tensors into the engine payload Adapter mode serialized the gathered adapter under torch's file_system sharing strategy, so the payload carried a /dev/shm path rather than the tensor bytes. That storage is reference counted: the TP workers that map it first drop their reference on return and unlink the file, and a rank that arrives late then opens a file that no longer exists. RuntimeError: unable to open shared memory object in read-write mode: No such file or directory The failure is timing dependent, which is why it surfaces as a single rank dying while its peers load the same adapter successfully. Pickle the tensors instead. An adapter is small enough to inline (~24MB at rank 16 on a 30B model) and the payload then has no lifetime coupling to the producer at all. SGLang reads it unchanged, since MultiprocessingSerializer.deserialize base64-decodes and unpickles. The base-weight path is untouched: those tensors stay on device and serialize to CUDA IPC handles, which do not have this problem. --- .../update_weight_from_tensor.py | 58 ++++++++----------- relax/utils/megatron_peft_utils.py | 29 ++++++++++ tests/utils/test_megatron_peft_utils.py | 46 +++++++++++++++ 3 files changed, 100 insertions(+), 33 deletions(-) diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py index 344ac9436..bcb295ac8 100644 --- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py +++ b/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 = [] diff --git a/relax/utils/megatron_peft_utils.py b/relax/utils/megatron_peft_utils.py index 7fd78a9ab..50e56cdd1 100644 --- a/relax/utils/megatron_peft_utils.py +++ b/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", ] diff --git a/tests/utils/test_megatron_peft_utils.py b/tests/utils/test_megatron_peft_utils.py index f16219e21..648d98775 100644 --- a/tests/utils/test_megatron_peft_utils.py +++ b/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_ handle, not as + bytes.""" + payload = base64.b64decode(serialize_adapter_tensors(self._adapter()), validate=True) + + assert b"/torch_" not in payload