[HotFix] fix fp8 scale load failed in tp>1 (#2837)

sgl-project · Jan 11, 2025 · f0e15dc · f0e15dc
1 parent f176958
commit f0e15dc
Showing 1 changed file with 2 additions and 7 deletions.
diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py
@@ -437,7 +437,7 @@ def weight_loader_v2(self, param: Parameter, loaded_weight: torch.Tensor):
         if len(loaded_weight.shape) == 0:
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
-        load_column_parallel_weight(param, loaded_weight, self.tp_rank)
+        param.load_column_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(self, input_):
         bias = self.bias if not self.skip_bias_add else None
@@ -1247,12 +1247,7 @@ def weight_loader_v2(self, param: BasevLLMParameter, loaded_weight: torch.Tensor
             assert loaded_weight.numel() == 1
             loaded_weight = loaded_weight.reshape(1)
 
-        load_row_parallel_weight(
-            param,
-            loaded_weight,
-            self.tp_rank,
-            use_presharded_weights=self.use_presharded_weights,
-        )
+        param.load_row_parallel_weight(loaded_weight=loaded_weight)
 
     def forward(self, input_):
         if self.input_is_parallel: