diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 96ed322df..302cc9f77 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -684,16 +684,23 @@ class Engine(EngineBase): ) def load_lora_adapter_from_tensors( - self, lora_name: str, tensors: List[Tuple[str, torch.Tensor]], config_dict: Dict + self, + lora_name: str, + tensors, + config_dict: Dict, + load_format: Optional[str] = None, ): - # Load LoRA adapter again - serialized_tensors = MultiprocessingSerializer.serialize( - tensors, output_str=True - ) + if load_format == "flattened_bucket": + serialized_tensors = tensors + else: + serialized_tensors = MultiprocessingSerializer.serialize( + tensors, output_str=True + ) lora_req = LoadLoRAAdapterFromTensorsReqInput( lora_name=lora_name, config_dict=config_dict, serialized_tensors=serialized_tensors, + load_format=load_format, ) return self.loop.run_until_complete( self.tokenizer_manager.load_lora_adapter_from_tensors(lora_req, None) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d543242ff..5f0db8137 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1766,6 +1766,7 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq): pinned: bool = False added_tokens_config: Optional[Dict[str, Any]] = None lora_id: Optional[str] = None + load_format: Optional[str] = None def to_ref(self) -> LoRARef: return LoRARef( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 86b009df4..8408cc905 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -49,6 +49,7 @@ from sglang.srt.utils.hf_transformers_utils import ( get_tokenizer_from_processor, ) from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions +from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket if TYPE_CHECKING: from sglang.srt.managers.cache_controller import LayerDoneCounter @@ -187,7 +188,17 @@ class BaseTpWorker(ABC): ): # The LoRA code handles TP sharding internally using slice_lora_a_weights # and slice_lora_b_weights methods (see lora/layers.py:46-49, mem_pool.py:437-440). - tensors = MultiprocessingSerializer.deserialize(recv_req.serialized_tensors) + if recv_req.load_format == "flattened_bucket": + flattened_data = MultiprocessingSerializer.deserialize( + recv_req.serialized_tensors + ) + bucket = FlattenedTensorBucket( + flattened_tensor=flattened_data["flattened_tensor"], + metadata=flattened_data["metadata"], + ) + tensors = dict(bucket.reconstruct_tensors()) + else: + tensors = MultiprocessingSerializer.deserialize(recv_req.serialized_tensors) result = self.model_runner.load_lora_adapter_from_tensors( recv_req.to_ref(), tensors, diff --git a/test/registered/rl/test_lora_load_from_tensor.py b/test/registered/rl/test_lora_load_from_tensor.py index fe21e9940..fd0113a73 100644 --- a/test/registered/rl/test_lora_load_from_tensor.py +++ b/test/registered/rl/test_lora_load_from_tensor.py @@ -329,6 +329,38 @@ class TestLoRALoadFromTensor(CustomTestCase): print("\n[Test]LoRA logprob comparison test passed!") + def test_lora_e2e_load_from_flattened_bucket(self): + """Test loading LoRA via FlattenedTensorBucket format (RL weight sync path).""" + from sglang.srt.utils import MultiprocessingSerializer + from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket + + named_tensors = list(self.lora_tensors.items()) + bucket = FlattenedTensorBucket(named_tensors=[(n, t) for n, t in named_tensors]) + bucket_dict = { + "flattened_tensor": bucket.get_flattened_tensor(), + "metadata": bucket.get_metadata(), + } + serialized = MultiprocessingSerializer.serialize(bucket_dict, output_str=True) + + result = self.engine.load_lora_adapter_from_tensors( + lora_name="self_cognition_Alice_flattened", + tensors=serialized, + config_dict=self.lora_config_dict, + load_format="flattened_bucket", + ) + self.assertTrue(result.success, f"Failed: {result.error_message}") + + output = self.engine.generate( + prompt=[TEST_PROMPT], + sampling_params={"max_new_tokens": MAX_NEW_TOKENS, "temperature": 0.0}, + lora_path=["self_cognition_Alice_flattened"], + ) + self.assertEqual( + output[0]["text"][: len(EXPECTED_OUTPUT)], + EXPECTED_OUTPUT, + "Output after applying LoRA via flattened bucket does not match expected", + ) + @classmethod def tearDownClass(cls): cls.engine.shutdown()