forked from phoenix-oss/llama-stack-mirror
09b793c4d6
I only tested with "on-the-fly" bf16 -> fp8 conversion, not the "load
from fp8" codepath.
YAML I tested with:
```
providers:
- provider_id: quantized
provider_type: meta-reference-quantized
config:
model: Llama3.1-8B-Instruct
quantization:
type: fp8
```
101 lines
3.7 KiB
Python
101 lines
3.7 KiB
Python
# 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.datatypes import CheckpointQuantizationFormat
|
|
from llama_models.llama3.reference_impl.model import Transformer, TransformerBlock
|
|
|
|
from llama_models.sku_list import resolve_model
|
|
from termcolor import cprint
|
|
from torch import Tensor
|
|
|
|
from llama_stack.apis.inference import QuantizationType
|
|
|
|
from llama_stack.providers.impls.meta_reference.inference.config import (
|
|
MetaReferenceQuantizedInferenceConfig,
|
|
)
|
|
|
|
|
|
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: MetaReferenceQuantizedInferenceConfig,
|
|
checkpoint_dir: str,
|
|
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
|
|
|
|
llama_model = resolve_model(config.model)
|
|
assert llama_model is not None, f"Model {config.model} not found"
|
|
|
|
# Move weights to GPU with quantization
|
|
if llama_model.quantization_format == CheckpointQuantizationFormat.fp8_mixed.value:
|
|
cprint("Loading fp8 scales...", "yellow")
|
|
fp8_scales_path = os.path.join(
|
|
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
|