d0502b4928 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -332,7 +332,10 @@ class Fp8Config(QuantizationConfig):

332 fused_mapping=self.packed_modules_mapping,

333 ):

334 return UnquantizedFusedMoEMethod(layer.moe_config)

335 - moe_quant_method = Fp8MoEMethod(self, layer)

335 + if self.is_checkpoint_fp8_serialized:

336 + moe_quant_method = Fp8MoEMethod(self, layer)

337 + else:

338 + moe_quant_method = Fp8OnlineMoEMethod(self, layer)

339 moe_quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)

340 return moe_quant_method

341 elif isinstance(layer, Attention):

@@ -745,8 +748,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):

748 layer.orig_dtype = params_dtype

749 layer.weight_block_size = None

750

748 - if self.quant_config.is_checkpoint_fp8_serialized:

749 - params_dtype = torch.float8_e4m3fn

751 + assert self.quant_config.is_checkpoint_fp8_serialized

752 + params_dtype = torch.float8_e4m3fn

753 +

754 if self.block_quant:

755 assert self.weight_block_size is not None

756 layer.weight_block_size = self.weight_block_size

@@ -773,41 +777,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

777 f"weight quantization block_k = {block_k}."

778 )

779

776 - # if we are doing online quantization, patch the weight

777 - # loaded to call `process_weights_after_loading` in a streaming fashion

778 - # as soon as the last weight chunk is loaded

779 - if not self.quant_config.is_checkpoint_fp8_serialized:

780 - weight_loader = extra_weight_attrs["weight_loader"]

781 - # create a new holder to prevent modifying behavior of any other

782 - # objects which might depend on the old one

783 - new_extra_weight_attrs = extra_weight_attrs

784 -

785 - def patched_weight_loader(param, loaded_weight, *args, **kwargs):

786 - # load the current weight chunk

787 - res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]

788 -

789 - # add a counter to track how many elements we have updated

790 - if not hasattr(layer, "_loaded_numel"):

791 - layer._loaded_numel = 0

792 - layer._loaded_numel += loaded_weight.numel()

793 -

794 - # if we have loaded all of the elements, call

795 - # process_weights_after_loading

796 - target_loaded_numel = layer.w13_weight.numel() + layer.w2_weight.numel()

797 - if layer._loaded_numel == target_loaded_numel:

798 - self.process_weights_after_loading(layer)

799 -

800 - # Delete the bookkeeping

801 - del layer._loaded_numel

802 - # Prevent the usual `process_weights_after_loading` call

803 - # from doing anything

804 - layer._already_called_process_weights_after_loading = True

805 -

806 - return res

807 -

808 - new_extra_weight_attrs["weight_loader"] = patched_weight_loader

809 - extra_weight_attrs = new_extra_weight_attrs

810 -

780 # WEIGHTS

781 w13_weight = torch.nn.Parameter(

782 torch.empty(

@@ -875,21 +844,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):

844 if self.block_quant

845 else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}

846 )

878 - # If loading fp8 checkpoint, pass the weight loaders.

879 - # If loading an fp16 checkpoint, do not (we will quantize in

880 - # process_weights_after_loading()

881 - if self.quant_config.is_checkpoint_fp8_serialized:

882 - set_weight_attrs(w13_weight_scale, extra_weight_attrs)

883 - set_weight_attrs(w2_weight_scale, extra_weight_attrs)

847 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)

848 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)

849

850 # INPUT_SCALES

851 if self.quant_config.activation_scheme == "static":

887 - if not self.quant_config.is_checkpoint_fp8_serialized:

888 - raise ValueError(

889 - "Found static activation scheme for checkpoint that "

890 - "was not serialized fp8."

891 - )

892 -

852 w13_input_scale = torch.nn.Parameter(

853 torch.ones(num_experts, dtype=torch.float32), requires_grad=False

854 )

@@ -986,45 +945,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

945 layer.w2_weight_scale_inv = Parameter(

946 dg_w2_weight_scale_inv, requires_grad=False

947 )

989 -

990 - # If checkpoint is fp16, quantize in place.

991 - elif not self.quant_config.is_checkpoint_fp8_serialized:

992 - fp8_dtype = current_platform.fp8_dtype()

993 - w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)

994 - w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)

995 -

996 - # Re-initialize w13_scale because we directly quantize

997 - # merged w13 weights and generate a single scaling factor.

998 - replace_parameter(

999 - layer,

1000 - "w13_weight_scale",

1001 - torch.ones(

1002 - layer.local_num_experts,

1003 - dtype=torch.float32,

1004 - device=w13_weight.device,

1005 - ),

1006 - )

1007 - for expert in range(layer.local_num_experts):

1008 - w13_weight[expert, :, :], layer.w13_weight_scale[expert] = (

1009 - ops.scaled_fp8_quant(layer.w13_weight.data[expert, :, :])

1010 - )

1011 - w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (

1012 - ops.scaled_fp8_quant(layer.w2_weight.data[expert, :, :])

1013 - )

1014 - replace_parameter(layer, "w13_weight", w13_weight)

1015 - replace_parameter(layer, "w2_weight", w2_weight)

1016 -

1017 - if self.rocm_aiter_moe_enabled:

1018 - # reshaping weights is required for aiter moe kernel.

1019 - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1020 - layer.w13_weight, layer.w2_weight

1021 - )

1022 -

1023 - replace_parameter(layer, "w13_weight", shuffled_w13)

1024 - replace_parameter(layer, "w2_weight", shuffled_w2)

1025 - # If checkpoint is fp8, we need to handle that the

1026 - # MoE kernels require single activation scale and single weight

1027 - # scale for w13 per expert.

948 else:

949 # Fp8 moe kernels require a single activation scale.

950 # We take the max of all the scales in case they differ.

@@ -1387,6 +1307,151 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1307 return result

1308

1309

1310 +class Fp8OnlineMoEMethod(Fp8MoEMethod):

1311 + """MoE method for online FP8 quantization.

1312 + Supports loading quantized FP16/BF16 model checkpoints with dynamic

1313 + activation scaling. The weight scaling factor will be initialized after

1314 + the model weights are loaded.

1315 +

1316 + Args:

1317 + quant_config: The quantization config.

1318 + """

1319 +

1320 + def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):

1321 + super().__init__(quant_config, layer)

1322 + assert not quant_config.is_checkpoint_fp8_serialized

1323 + assert quant_config.activation_scheme == "dynamic"

1324 + assert quant_config.weight_block_size is None

1325 + assert self.flashinfer_moe_backend is None

1326 +

1327 + def create_weights(

1328 + self,

1329 + layer: Module,

1330 + num_experts: int,

1331 + hidden_size: int,

1332 + intermediate_size_per_partition: int,

1333 + params_dtype: torch.dtype,

1334 + **extra_weight_attrs,

1335 + ):

1336 + layer.intermediate_size_per_partition = intermediate_size_per_partition

1337 + layer.hidden_size = hidden_size

1338 + layer.num_experts = num_experts

1339 + layer.orig_dtype = params_dtype

1340 + layer.weight_block_size = None

1341 +

1342 + # We are doing online quantization, patch the weight loaded

1343 + # to call `process_weights_after_loading` in a streaming fashion

1344 + # as soon as the last weight chunk is loaded.

1345 + weight_loader = extra_weight_attrs["weight_loader"]

1346 + # create a new holder to prevent modifying behavior of any other

1347 + # objects which might depend on the old one

1348 + new_extra_weight_attrs = extra_weight_attrs

1349 +

1350 + def patched_weight_loader(param, loaded_weight, *args, **kwargs):

1351 + # load the current weight chunk

1352 + res = weight_loader(param, loaded_weight, *args, **kwargs) # type: ignore[misc]

1353 +

1354 + # add a counter to track how many elements we have updated

1355 + if not hasattr(layer, "_loaded_numel"):

1356 + layer._loaded_numel = 0

1357 + layer._loaded_numel += loaded_weight.numel()

1358 +

1359 + # if we have loaded all of the elements, call

1360 + # process_weights_after_loading

1361 + target_loaded_numel = layer.w13_weight.numel() + layer.w2_weight.numel()

1362 + if layer._loaded_numel == target_loaded_numel:

1363 + self.process_weights_after_loading(layer)

1364 +

1365 + # Delete the bookkeeping

1366 + del layer._loaded_numel

1367 + # Prevent the usual `process_weights_after_loading` call

1368 + # from doing anything

1369 + layer._already_called_process_weights_after_loading = True

1370 +

1371 + return res

1372 +

1373 + new_extra_weight_attrs["weight_loader"] = patched_weight_loader

1374 + extra_weight_attrs = new_extra_weight_attrs

1375 +

1376 + # WEIGHTS

1377 + w13_weight = torch.nn.Parameter(

1378 + torch.empty(

1379 + num_experts,

1380 + 2 * intermediate_size_per_partition,

1381 + hidden_size,

1382 + dtype=params_dtype,

1383 + ),

1384 + requires_grad=False,

1385 + )

1386 + layer.register_parameter("w13_weight", w13_weight)

1387 + set_weight_attrs(w13_weight, extra_weight_attrs)

1388 +

1389 + w2_weight = torch.nn.Parameter(

1390 + torch.empty(

1391 + num_experts,

1392 + hidden_size,

1393 + intermediate_size_per_partition,

1394 + dtype=params_dtype,

1395 + ),

1396 + requires_grad=False,

1397 + )

1398 + layer.register_parameter("w2_weight", w2_weight)

1399 + set_weight_attrs(w2_weight, extra_weight_attrs)

1400 +

1401 + # WEIGHT_SCALES

1402 + # Allocate 2 scales for w1 and w3 respectively.

1403 + # They will be combined to a single scale after weight loading.

1404 + w13_weight_scale = torch.nn.Parameter(

1405 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False

1406 + )

1407 + w2_weight_scale = torch.nn.Parameter(

1408 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False

1409 + )

1410 + layer.register_parameter("w13_weight_scale", w13_weight_scale)

1411 + layer.register_parameter("w2_weight_scale", w2_weight_scale)

1412 +

1413 + layer.w13_input_scale = None

1414 + layer.w2_input_scale = None

1415 +

1416 + self.rocm_aiter_moe_enabled = False

1417 +

1418 + def process_weights_after_loading(self, layer: Module) -> None:

1419 + if getattr(layer, "_already_called_process_weights_after_loading", False):

1420 + return

1421 +

1422 + # Lazy import to avoid importing triton too early.

1423 + self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()

1424 +

1425 + # If checkpoint is fp16, quantize in place.

1426 + fp8_dtype = current_platform.fp8_dtype()

1427 + w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)

1428 + w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)

1429 +

1430 + for expert in range(layer.local_num_experts):

1431 + w13_weight[expert, :, :], layer.w13_weight_scale[expert] = (

1432 + ops.scaled_fp8_quant(layer.w13_weight.data[expert, :, :])

1433 + )

1434 + w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (

1435 + ops.scaled_fp8_quant(layer.w2_weight.data[expert, :, :])

1436 + )

1437 + replace_parameter(layer, "w13_weight", w13_weight)

1438 + replace_parameter(layer, "w2_weight", w2_weight)

1439 +

1440 + # Reshuffle weights for AITER if needed.

1441 + if self.rocm_aiter_moe_enabled:

1442 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1443 + layer.w13_weight, layer.w2_weight

1444 + )

1445 + replace_parameter(layer, "w13_weight", shuffled_w13)

1446 + replace_parameter(layer, "w2_weight", shuffled_w2)

1447 +

1448 + # Rushuffle weights for MARLIN if needed.

1449 + if self.use_marlin:

1450 + prepare_moe_fp8_layer_for_marlin(

1451 + layer, False, input_dtype=self.marlin_input_dtype

1452 + )

1453 +

1454 +

1455 class Fp8KVCacheMethod(BaseKVCacheMethod):

1456 """

1457 Supports loading kv-cache scaling factors from FP8 checkpoints.

[MoE][Refactor 1/N] Separate Online Quantization (#30627)