[feat] Enhance lora_update_weight_from_tensor for RL training (#19314)

python/sglang/srt/entrypoints/engine.py

@@ -684,9 +684,15 @@ 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
+        if load_format == "flattened_bucket":
+            serialized_tensors = tensors
+        else:
             serialized_tensors = MultiprocessingSerializer.serialize(
                 tensors, output_str=True
             )
@@ -694,6 +700,7 @@ class Engine(EngineBase):
             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)

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(

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,6 +188,16 @@ 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).
+        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(),

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()