diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index bfb8009b2..9b426cc26 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -588,7 +588,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear): needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) if loaded_shard_id is None: - # Loaded weight is already fused on disk (qkv/mlp). + # Loaded weight is already fused-in-checkpoint (qkv/mlp). if output_dim is None: if needs_scalar_to_array: param_data, loaded_weight = adjust_scalar_to_fused_array( @@ -731,8 +731,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear): ): """ Handle special case for models where MLP layers are already - fused on disk. In this case, we have no shard id. This function - determmines the shard id by splitting these layers and then calls + fused-in-checkpoint. In this case, we have no shard id. This function + determines the shard id by splitting these layers and then calls the weight loader using the shard id. An example of a model with these fused layers: @@ -832,12 +832,28 @@ class MergedColumnParallelLinear(ColumnParallelLinear): ): if loaded_shard_id is None or isinstance(loaded_shard_id, tuple): if isinstance(param, PerTensorScaleParameter): - param.load_merged_column_weight( - loaded_weight=loaded_weight, - shard_id=0, - tp_rank=self.tp_rank, - tp_size=self.tp_size, - ) + if loaded_weight.numel() != 1: + raise ValueError( + "Expected scalar scale for fused-in-checkpoint " + "merged-column checkpoint load, got shape " + f"{tuple(loaded_weight.shape)}" + ) + if loaded_shard_id is None: + # The checkpoint tensor is already fused-in-checkpoint, so a + # scalar scale applies to the entire merged matrix. Fill + # every logical slot so later reductions only see valid + # scale values. + shard_ids = range(param.data.shape[0]) + else: + shard_ids = loaded_shard_id + + for shard_id in shard_ids: + param.load_merged_column_weight( + loaded_weight=loaded_weight, + shard_id=shard_id, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ) return elif isinstance(param, BlockQuantScaleParameter): self._load_merged_block_scale(param, loaded_weight) @@ -1007,8 +1023,8 @@ class QKVParallelLinear(ColumnParallelLinear): ): """ Handle special case for models where QKV layers are already - fused on disk. In this case, we have no shard id. This function - determmines the shard id by splitting these layers and then calls + fused-in-checkpoint. In this case, we have no shard id. This function + determines the shard id by splitting these layers and then calls the weight loader using the shard id. An example of a model with these fused layers: @@ -1084,7 +1100,19 @@ class QKVParallelLinear(ColumnParallelLinear): ): if loaded_shard_id is None: # special case for certain models if isinstance(param, PerTensorScaleParameter): - param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0) + # The checkpoint tensor is already fused-in-checkpoint, so a scalar + # scale applies to the entire QKV matrix. Fill every logical + # slot so later reductions only see valid scale values. + if loaded_weight.numel() != 1: + raise ValueError( + "Expected scalar scale for fused-in-checkpoint QKV " + "checkpoint load when loaded_shard_id is None, got " + f"shape {tuple(loaded_weight.shape)}" + ) + for shard_id in param.qkv_idxs: + param.load_qkv_weight( + loaded_weight=loaded_weight, shard_id=shard_id + ) return elif type(param) in (RowvLLMParameter, BasevLLMParameter): param.load_qkv_weight(loaded_weight=loaded_weight) @@ -1156,7 +1184,7 @@ class QKVParallelLinear(ColumnParallelLinear): needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False) if loaded_shard_id is None: - # Loaded weight is already fused on disk (qkv/mlp). + # Loaded weight is already fused-in-checkpoint (qkv/mlp). if output_dim is None: if needs_scalar_to_array: param_data, loaded_weight = adjust_scalar_to_fused_array( diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py new file mode 100644 index 000000000..7773928c7 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py @@ -0,0 +1,80 @@ +import unittest + +import torch + +from sglang.srt.layers.linear import MergedColumnParallelLinear, QKVParallelLinear +from sglang.srt.layers.parameter import PerTensorScaleParameter +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestModelOptNvfp4(CustomTestCase): + def _make_layer(self): + return MergedColumnParallelLinear( + input_size=16, + output_sizes=[16, 16], + bias=False, + tp_rank=0, + tp_size=1, + ) + + def _make_qkv_layer(self): + return QKVParallelLinear( + hidden_size=16, + head_size=8, + total_num_heads=2, + total_num_kv_heads=2, + bias=False, + tp_rank=0, + tp_size=1, + ) + + def test_fused_scalar_scale_load_fills_all_logical_slots(self): + layer = self._make_layer() + scale = PerTensorScaleParameter( + data=torch.empty(2, dtype=torch.float32), + weight_loader=layer.weight_loader_v2, + ) + + layer.weight_loader_v2(scale, torch.tensor(0.25, dtype=torch.float32)) + + torch.testing.assert_close(scale, torch.tensor([0.25, 0.25])) + + def test_fused_scalar_scale_load_rejects_non_scalar(self): + layer = self._make_layer() + scale = PerTensorScaleParameter( + data=torch.empty(2, dtype=torch.float32), + weight_loader=layer.weight_loader_v2, + ) + + with self.assertRaisesRegex(ValueError, "Expected scalar scale"): + layer.weight_loader_v2(scale, torch.tensor([0.25, 0.5])) + + def test_fused_qkv_scalar_scale_load_fills_all_logical_slots(self): + layer = self._make_qkv_layer() + scale = PerTensorScaleParameter( + data=torch.empty(3, dtype=torch.float32), + weight_loader=layer.weight_loader_v2, + ) + + layer.weight_loader_v2(scale, torch.tensor(0.125, dtype=torch.float32)) + + torch.testing.assert_close(scale, torch.tensor([0.125, 0.125, 0.125])) + + def test_explicit_shard_scale_loads_stay_independent(self): + layer = self._make_layer() + scale = PerTensorScaleParameter( + data=torch.empty(2, dtype=torch.float32), + weight_loader=layer.weight_loader_v2, + ) + + layer.weight_loader_v2(scale, torch.tensor(0.25, dtype=torch.float32), 0) + layer.weight_loader_v2(scale, torch.tensor(0.5, dtype=torch.float32), 1) + + torch.testing.assert_close(scale, torch.tensor([0.25, 0.5])) + + +if __name__ == "__main__": + unittest.main()