make inference server load checkpoints for fp8 inference

- introduce quantization related args for inference config
- also kill GeneratorArgs

toolchain/inference/quantization/loader.py

@@ -0,0 +1,106 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed in accordance with the terms of the Llama 3 Community License Agreement.
+
+import os
+
+from typing import Optional
+
+import torch
+from models.llama3_1.api.model import Transformer, TransformerBlock
+
+from toolchain.inference.api.config import (
+    CheckpointQuantizationFormat,
+    InlineImplConfig,
+)
+from toolchain.inference.api.datatypes import (
+    QuantizationType,
+)
+
+from termcolor import cprint
+
+def is_fbgemm_available() -> bool:
+    try:
+        import fbgemm_gpu.experimental.gen_ai  # noqa: F401
+        return True
+    except (ImportError, ModuleNotFoundError):
+        return False
+
+
+def convert_to_quantized_model(
+    model: Transformer,
+    config: InlineImplConfig,
+    fp8_activation_scale_ub: Optional[float] = 1200.0,
+) -> Transformer:
+    if config.quantization.type == QuantizationType.bf16.value:
+        return model
+
+    elif config.quantization.type != QuantizationType.fp8.value:
+        raise ValueError("Only FP8 quantization is supported")
+
+    from .fp8_impls import Fp8ScaledWeights, load_fp8, quantize_fp8
+
+
+
+    checkpoint = config.checkpoint_config.checkpoint
+    # Move weights to GPU with quantization
+    if checkpoint.quantization_format == CheckpointQuantizationFormat.fp8_mixed.value:
+        cprint("Loading fp8 scales...", "yellow")
+        fp8_scales_path = os.path.join(
+            checkpoint.checkpoint_dir, f"fp8_scales_{get_model_parallel_rank()}.pt"
+        )
+        assert os.path.isfile(
+            fp8_scales_path
+        ), f"fp8_scales_path not found for rank {get_model_parallel_rank()}"
+        fp8_scales = torch.load(fp8_scales_path, weights_only=True)
+
+        for block in model.layers:
+            if isinstance(block, TransformerBlock):
+                if block.layer_id == 0 or block.layer_id == (model.n_layers - 1):
+                    continue
+                block.feed_forward.w1.weight = load_fp8(
+                    block.feed_forward.w1.weight,
+                    fp8_scales[
+                        f"{block.layer_id}_feed_forward.w1_{get_model_parallel_rank()}"
+                    ],
+                    fp8_activation_scale_ub,
+                )
+                block.feed_forward.w3.weight = load_fp8(
+                    block.feed_forward.w3.weight,
+                    fp8_scales[
+                        f"{block.layer_id}_feed_forward.w3_{get_model_parallel_rank()}"
+                    ],
+                    fp8_activation_scale_ub,
+                )
+                block.feed_forward.w2.weight = load_fp8(
+                    block.feed_forward.w2.weight,
+                    fp8_scales[
+                        f"{block.layer_id}_feed_forward.w2_{get_model_parallel_rank()}"
+                    ],
+                    fp8_activation_scale_ub,
+                )
+    else:
+        cprint("Quantizing fp8 weights from bf16...", "yellow")
+        for block in model.layers:
+            if isinstance(block, TransformerBlock):
+                if block.layer_id == 0 or block.layer_id == (model.n_layers - 1):
+                    continue
+                block.feed_forward.w1.weight = quantize_fp8(
+                    block.feed_forward.w1.weight,
+                    fp8_activation_scale_ub,
+                    output_device=torch.device("cuda"),
+                )
+                block.feed_forward.w3.weight = quantize_fp8(
+                    block.feed_forward.w3.weight,
+                    fp8_activation_scale_ub,
+                    output_device=torch.device("cuda"),
+                )
+                block.feed_forward.w2.weight = quantize_fp8(
+                    block.feed_forward.w2.weight,
+                    fp8_activation_scale_ub,
+                    output_device=torch.device("cuda"),
+                )
+
+    for _, parameter in model.named_parameters():
+        if not isinstance(parameter, Fp8ScaledWeights):
+            parameter.data = parameter.to(device="cuda")
+    return model