Support Flatten Tensor Update Weights to speed up MOE Update Weights by 20% (#8079)

test/srt/rl/test_update_weights_from_tensor.py

@@ -5,6 +5,7 @@ import unittest
 import torch
 
 import sglang as sgl
+from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
 from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST, CustomTestCase
 
 
@@ -112,6 +113,59 @@ class TestUpdateWeightsFromTensor(CustomTestCase):
 
         engine.shutdown()
 
+    def test_update_weights_from_tensor_load_format_flattened_bucket(self):
+        """Test updating weights using flattened_bucket format"""
+        engine = sgl.Engine(model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST)
+
+        # Create a small set of parameters for testing
+        param_names = [f"model.layers.{i}.mlp.up_proj.weight" for i in range(6, 10)]
+
+        # Check original values
+        _check_param(engine, param_names[0], [0.0087, -0.0214, -0.0004, 0.0039, 0.0110])
+
+        # Create new tensors with different values
+        new_tensors = []
+        for _, name in enumerate(param_names):
+            # Create tensors with different values for each parameter
+            value = 2.0  # Different value for each parameter
+            new_tensor = torch.full((16384, 2048), value, device="cuda")
+            new_tensors.append((name, new_tensor))
+
+        # Create a flattened bucket
+        flattened_bucket = FlattenedTensorBucket(named_tensors=new_tensors)
+
+        # Extract the flattened tensor and metadata in the format expected by model_runner
+        flattened_tensor = flattened_bucket.get_flattened_tensor()
+        metadata = flattened_bucket.get_metadata()
+
+        # Create the dict format expected by _update_weights_from_flattened_bucket
+        bucket_dict = {"flattened_tensor": flattened_tensor, "metadata": metadata}
+
+        # Serialize the bucket data
+        from sglang.srt.utils import MultiprocessingSerializer
+
+        serialized_bucket = MultiprocessingSerializer.serialize(
+            bucket_dict, output_str=True
+        )
+
+        # Create a list where each rank contains the same serialized data
+        # This simulates the distributed environment where each rank has the same data
+        serialized_bucket_list = [serialized_bucket]
+
+        # Update weights using flattened_bucket format
+        time_start = time.perf_counter()
+        engine.update_weights_from_tensor(
+            named_tensors=serialized_bucket_list, load_format="flattened_bucket"
+        )
+        update_time = time.perf_counter() - time_start
+        print(f"Flattened bucket update time: {update_time:.03f}")
+
+        # Verify the weights were updated correctly
+        for i, param_name in enumerate(param_names):
+            _check_param(engine, param_name, [2.0] * 5)
+
+        engine.shutdown()
+
 
 def _check_param(engine, param_name, expect_values):
     actual_values = torch.tensor(engine.get_weights_by_name(param_name))[0, :5]