[RL] Allow passing tensors of different dtypes for FlattenedTensorBucket

Author: zhuzilin

python/sglang/srt/weight_sync/tensor_bucket.py

@@ -48,7 +48,7 @@ def __init__(
             flattened_tensors: List[torch.Tensor] = [None] * len(named_tensors)
 
             for i, (name, tensor) in enumerate(named_tensors):
-                flattened = tensor.flatten()
+                flattened = tensor.flatten().view(torch.uint8)
                 flattened_tensors[i] = flattened
 
                 # Store metadata
@@ -93,14 +93,12 @@ def reconstruct_tensors(self) -> List[Tuple[str, torch.Tensor]]:
         reconstructed = [None] * len(self.metadata)
 
         for i, meta in enumerate(self.metadata):
-            tensor = self.flattened_tensor[meta.start_idx : meta.end_idx].reshape(
-                meta.shape
+            tensor = (
+                self.flattened_tensor[meta.start_idx : meta.end_idx]
+                .view(meta.dtype)
+                .reshape(meta.shape)
             )
 
-            # batch dtype conversion (if needed)
-            if tensor.dtype != meta.dtype:
-                tensor = tensor.to(meta.dtype)
-
             reconstructed[i] = (meta.name, tensor)
 
         return reconstructed