00ebf19cca · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/distributed/utils.py

+18

@@ -64,6 +64,24 @@ def divide(numerator, denominator):

64 return numerator // denominator

65

66

67 +def verify_group_size_divides_partition(

68 + input_size_per_partition: int,

69 + group_size: int,

70 + layer_name: str | None = None,

71 + extra_suggestion: str = "",

72 +) -> None:

73 + """Validate that a TP-sharded layer holds a whole number of quant groups."""

74 + if input_size_per_partition % group_size == 0:

75 + return

76 + location = f" for layer '{layer_name}'" if layer_name else ""

77 + raise ValueError(

78 + f"Weight {input_size_per_partition=}{location} is not divisible by "

79 + f"{group_size=}. This happens when tensor_parallel_size splits the layer input "

80 + "into shards that are not a whole number of quant groups. Consider reducing "

81 + f"tensor_parallel_size{extra_suggestion}."

82 + )

83 +

84 +

85 def is_weak_contiguous(inp: torch.Tensor) -> bool:

86 """Check that *inp* occupies a single contiguous block of memory.

87

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py

+2 -1

@@ -6,6 +6,7 @@ from collections.abc import Callable

6 import torch

7 from compressed_tensors.quantization import ActivationOrdering

8

9 +from vllm.distributed.utils import verify_group_size_divides_partition

10 from vllm.logger import init_logger

11 from vllm.model_executor.kernels.linear import (

12 MPLinearLayerConfig,

@@ -112,7 +113,7 @@ class CompressedTensorsW4A8Fp8(CompressedTensorsScheme):

113 scales_and_zp_size = input_size // group_size

114

115 if partition_scales:

115 - assert input_size_per_partition % group_size == 0

116 + verify_group_size_divides_partition(input_size_per_partition, group_size)

117 scales_and_zp_size = input_size_per_partition // group_size

118

119 weight = PackedvLLMParameter(

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py

+3 -3

@@ -5,6 +5,7 @@ from collections.abc import Callable

5

6 import torch

7

8 +from vllm.distributed.utils import verify_group_size_divides_partition

9 from vllm.logger import init_logger

10 from vllm.model_executor.kernels.linear import (

11 MPLinearLayerConfig,

@@ -79,9 +80,8 @@ class CompressedTensorsW4A8Int(CompressedTensorsScheme):

80 effective_group_size = self.group_size

81

82 # Ensure group_size divides input_size_per_partition

82 - assert input_size_per_partition % effective_group_size == 0, (

83 - f"input_size_per_partition {input_size_per_partition}"

84 - f" not divisible by group_size {effective_group_size}"

83 + verify_group_size_divides_partition(

84 + input_size_per_partition, effective_group_size

85 )

86

87 # Determine scale partitioning

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py

+4 -1

@@ -8,6 +8,7 @@ from fractions import Fraction

8 import torch

9 from compressed_tensors.quantization import ActivationOrdering

10

11 +from vllm.distributed.utils import verify_group_size_divides_partition

12 from vllm.logger import init_logger

13 from vllm.model_executor.kernels.linear import (

14 MarlinLinearKernel,

@@ -151,7 +152,9 @@ class CompressedTensorsWNA16(CompressedTensorsScheme):

152 scales_and_zp_size = input_size // group_size

153

154 if partition_scales:

154 - assert input_size_per_partition % group_size == 0

155 + verify_group_size_divides_partition(

156 + input_size_per_partition, group_size, self.layer_name

157 + )

158 scales_and_zp_size = input_size_per_partition // group_size

159

160 packed_input_dim = math.ceil(input_size_per_partition * self.num_bits / 32)

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa8o8.py

+4 -1

@@ -13,6 +13,7 @@ from collections.abc import Callable

13 import torch

14 from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32

15

16 +from vllm.distributed.utils import verify_group_size_divides_partition

17 from vllm.model_executor.kernels.linear import (

18 MPLinearLayerConfig,

19 choose_mp_linear_kernel,

@@ -171,7 +172,9 @@ class CompressedTensorsWNA8O8Int(CompressedTensorsScheme):

172 scales = (input_size_per_partition if partitioned else input_size) // group_size

173 scale_data = torch.empty(out, scales, dtype=params_dtype)

174 if partitioned:

174 - assert input_size_per_partition % group_size == 0

175 + verify_group_size_divides_partition(

176 + input_size_per_partition, group_size, self.layer_name

177 + )

178 weight_scale = GroupQuantScaleParameter(

179 data=scale_data, output_dim=0, input_dim=1, weight_loader=weight_loader

180 )

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+6 -6

@@ -9,6 +9,7 @@ import torch

9

10 import vllm.envs as envs

11 from vllm import _custom_ops as ops

12 +from vllm.distributed.utils import verify_group_size_divides_partition

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.fused_moe import RoutedExperts

15 from vllm.model_executor.layers.linear import LinearBase

@@ -193,12 +194,11 @@ def verify_marlin_supports_shape(

194 "with --quantization gptq."

195 )

196

196 - if group_size < input_size and input_size_per_partition % group_size != 0:

197 - raise ValueError(

198 - f"Weight input_size_per_partition = {input_size_per_partition}"

199 - f" is not divisible by group_size = {group_size}. "

200 - "Consider reducing tensor_parallel_size or running "

201 - "with --quantization gptq."

197 + if group_size < input_size:

198 + verify_group_size_divides_partition(

199 + input_size_per_partition,

200 + group_size,

201 + extra_suggestion=" or running with --quantization gptq",

202 )

203

204

[Bugfix][Quant] Raise actionable error instead of bare assert for group-size/TP mismatch (#46230) (#46236)