Initial commit

llama_toolchain/inference/quantization/loader.py

@@ -0,0 +1,105 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# 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 fairscale.nn.model_parallel.mappings import reduce_from_model_parallel_region
+from llama_models.llama3_1.api.model import Transformer, TransformerBlock
+
+from llama_toolchain.inference.api.config import (
+    CheckpointQuantizationFormat,
+    InlineImplConfig,
+)
+from llama_toolchain.inference.api.datatypes import QuantizationType
+
+from termcolor import cprint
+from torch import Tensor
+
+
+def is_fbgemm_available() -> bool:
+    try:
+        import fbgemm_gpu.experimental.gen_ai  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def swiglu_wrapper(
+    self,
+    x: Tensor,
+):
+    from .fp8_impls import ffn_swiglu
+
+    out = ffn_swiglu(x, self.w1.weight, self.w3.weight, self.w2.weight)
+    return reduce_from_model_parallel_region(out)
+
+
+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.forward = swiglu_wrapper.__get__(block.feed_forward)
+                for key in ("w1", "w3", "w2"):
+                    param = getattr(block.feed_forward, key)
+                    param.weight = load_fp8(
+                        param.weight,
+                        fp8_scales[
+                            f"{block.layer_id}_feed_forward.{key}_{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.forward = swiglu_wrapper.__get__(block.feed_forward)
+                for key in ("w1", "w3", "w2"):
+                    param = getattr(block.feed_forward, key)
+                    param.weight = quantize_fp8(
+                        param.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