llama-stack 0.4.4__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack/models/llama/llama3/multimodal/utils.py

@@ -1,26 +0,0 @@
-# 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.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# top-level folder for each specific model found within the models/ directory at
-# the top-level of this source tree.
-
-import collections
-
-import torch
-
-
-def get_negative_inf_value(dtype):
-    return torch.finfo(dtype).min
-
-
-def to_2tuple(x):
-    if isinstance(x, collections.abc.Iterable):
-        return x
-    return (x, x)

llama_stack/models/llama/llama3/quantization/__init__.py

@@ -1,5 +0,0 @@
-# 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.

llama_stack/models/llama/llama3/quantization/loader.py

@@ -1,316 +0,0 @@
-# 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.
-
-# type: ignore
-import os
-from typing import Any, cast
-
-import torch
-from fairscale.nn.model_parallel.initialize import get_model_parallel_rank
-from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear
-from fairscale.nn.model_parallel.mappings import reduce_from_model_parallel_region
-from torch import Tensor, nn
-from torchao.quantization.GPTQ import Int8DynActInt4WeightLinear
-
-from ...datatypes import QuantizationMode
-from ...quantize_impls import (
-    Fp8ScaledWeights,
-    ffn_swiglu,
-    load_fp8,
-    quantize_fp8,
-)
-from ..model import Transformer, TransformerBlock
-from ..multimodal.model import CrossAttentionTransformer
-
-
-def swiglu_wrapper(
-    self,
-    x: Tensor,
-):
-    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 | CrossAttentionTransformer,
-    checkpoint_dir: str,
-    quantization_mode: str | None = None,
-    fp8_activation_scale_ub: float | None = 1200.0,
-    device: torch.device | None = None,
-) -> Transformer | CrossAttentionTransformer:
-    if quantization_mode == QuantizationMode.fp8_mixed:
-        return convert_to_fp8_quantized_model(model, checkpoint_dir, fp8_activation_scale_ub, device)
-    elif quantization_mode == QuantizationMode.int4_mixed:
-        return convert_to_int4_quantized_model(model, checkpoint_dir, device)
-    else:
-        raise ValueError(f"Unsupported quantization mode: {quantization_mode}")
-
-
-def convert_to_fp8_quantized_model(
-    model: Transformer,
-    checkpoint_dir: str,
-    fp8_activation_scale_ub: float | None = 1200.0,
-    device: torch.device | None = None,
-) -> Transformer:
-    # Move weights to GPU with quantization
-    fp8_scales_path = os.path.join(checkpoint_dir, f"fp8_scales_{get_model_parallel_rank()}.pt")
-    if os.path.isfile(fp8_scales_path):
-        print("Loading fp8 scales...")
-        fp8_scales = torch.load(fp8_scales_path, weights_only=True)
-
-        for _, block in model.named_modules():
-            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:
-        print("Quantizing fp8 weights from bf16...")
-        for _, block in model.named_modules():
-            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)  # type: ignore
-                for key in ("w1", "w3", "w2"):
-                    param = getattr(block.feed_forward, key)
-                    param.weight = quantize_fp8(
-                        param.weight,
-                        fp8_activation_scale_ub,
-                        output_device=device,
-                    )
-
-    for _, parameter in model.named_parameters():
-        if not isinstance(parameter, Fp8ScaledWeights):
-            parameter.data = parameter.to(device=device)
-    return model
-
-
-class Int8DynActInt4WeightLinearLoRA(Int8DynActInt4WeightLinear):
-    """
-    Int8DynActInt4WeightLinear with LoRA adaptor.
-
-    Args:
-        in_features: Number of input features.
-        out_features: Number of output features.
-        bias: Whether to use bias.
-        device: Device to use.
-        group_size: Group size for quantization.
-        precision: Precision of quantization.
-        scales_precision: Precision of scales.
-        lora_rank: Rank of LoRA adaptor.
-        lora_scale: Scale of LoRA adaptor.
-    """
-
-    def __init__(
-        self,
-        in_features: int,
-        out_features: int,
-        bias=False,
-        device=None,
-        # quantization parameters
-        group_size: int = 256,
-        precision: torch.dtype = torch.float32,
-        scales_precision: torch.dtype = torch.float32,
-        # LoRA parameters
-        lora_rank: int | None = None,
-        lora_scale: float | None = None,
-    ) -> None:
-        super().__init__(
-            in_features,
-            out_features,
-            bias=bias,
-            device=device,
-            groupsize=group_size,
-            precision=precision,
-            scales_precision=scales_precision,
-        )
-        self.lora_scale: float | None = None
-        self.adaptor: nn.Sequential | None = None
-        if lora_rank is not None:
-            assert lora_scale is not None, "Please specify lora scale for LoRA."
-            # Low-rank adaptation. See paper for more details: https://arxiv.org/abs/2106.09685
-            self.adaptor = nn.Sequential()
-            self.adaptor.add_module("A", nn.Linear(in_features, lora_rank, bias=False))
-            self.adaptor.add_module("B", nn.Linear(lora_rank, out_features, bias=False))
-            self.lora_scale = lora_scale
-        self._register_load_state_dict_pre_hook(self.load_hook)
-
-    def load_hook(
-        self,
-        state_dict: dict[str, Any],
-        prefix: str,
-        local_metadata: dict[str, Any],
-        strict: bool,
-        missing_keys: list[str],
-        unexpected_keys: list[str],
-        error_msgs: list[str],
-    ) -> None:
-        """A hook to load the quantized weights from the state dict."""
-        if prefix + "zeros" not in state_dict:
-            # Zero-point may not be saved in the state dict. In this case, we assume it's zero.
-            assert prefix + "scales" in state_dict
-            state_dict[prefix + "zeros"] = torch.zeros_like(state_dict[prefix + "scales"])
-
-    def forward(self, input_: torch.Tensor) -> torch.Tensor:
-        module_out = super().forward(input_)
-        if self.adaptor is not None:
-            adaptor_out = self.adaptor(input_) * self.lora_scale
-            return module_out + adaptor_out
-        return module_out
-
-
-class Int8WeightEmbedding(torch.nn.Embedding):
-    """An embedding layer to load int8 weights.
-
-    Args:
-        num_embeddings: Number of embeddings.
-        embedding_dim: Embedding dimension.
-        padding_idx: Padding index.
-    """
-
-    def __init__(
-        self,
-        num_embeddings: int,
-        embedding_dim: int,
-        padding_idx: int,
-        device=None,
-    ) -> None:
-        super().__init__(num_embeddings, embedding_dim, padding_idx, device=device)
-
-        self._register_load_state_dict_pre_hook(self.load_hook)
-
-    def load_hook(
-        self,
-        state_dict: dict[str, Any],
-        prefix: str,
-        local_metadata: dict[str, Any],
-        strict: bool,
-        missing_keys: list[str],
-        unexpected_keys: list[str],
-        error_msgs: list[str],
-    ) -> None:
-        """A hook to load the quantized embedding weight and scales from the state dict."""
-        weights = state_dict.pop(prefix + "weight")
-        scales = state_dict.pop(prefix + "scales")
-        state_dict[prefix + "weight"] = weights * scales
-
-
-class Int8WeightLinear(torch.nn.Linear):
-    """A linear layer to load int8 weights.
-
-    Args:
-        in_features: Number of input features.
-        out_features: Number of output features.
-        bias: Whether to use bias.
-    """
-
-    def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None) -> None:
-        super().__init__(in_features, out_features, bias, device=device)
-
-        self._register_load_state_dict_pre_hook(self.load_hook)
-
-    def load_hook(
-        self,
-        state_dict: dict[str, Any],
-        prefix: str,
-        local_metadata: dict[str, Any],
-        strict: bool,
-        missing_keys: list[str],
-        unexpected_keys: list[str],
-        error_msgs: list[str],
-    ) -> None:
-        """A hook to load the quantized linear weight and scales from the state dict."""
-        weights = state_dict.pop(prefix + "weight")
-        scales = state_dict.pop(prefix + "scales")
-        state_dict[prefix + "weight"] = weights * scales
-
-
-def _prepare_model_int4_weight_int8_dynamic_activation(
-    model: torch.nn.Module,
-    group_size: int,
-    lora_rank: int | None,
-    lora_scale: float | None,
-):
-    """Prepare the model for int4 weight and int8 dynamic activation quantization.
-
-    Note that the weights of embedding and output layers are quantized to int8.
-    """
-    device = None
-    for module_name, module in model.named_children():
-        if module_name == "output":
-            quantized_module = Int8WeightLinear(
-                in_features=module.in_features,
-                out_features=module.out_features,
-                bias=module.bias,
-                device=device,
-            )
-            del module
-            setattr(model, module_name, quantized_module)
-        elif module_name == "tok_embeddings":
-            quantized_module = Int8WeightEmbedding(
-                num_embeddings=module.num_embeddings,
-                embedding_dim=module.embedding_dim,
-                padding_idx=module.padding_idx,
-                device=device,
-            )
-            del module
-            setattr(model, module_name, quantized_module)
-        elif isinstance(module, ColumnParallelLinear | RowParallelLinear | nn.Linear):
-            quantized_module = Int8DynActInt4WeightLinearLoRA(
-                in_features=module.in_features,
-                out_features=module.out_features,
-                bias=False,
-                group_size=group_size,
-                lora_rank=lora_rank,
-                lora_scale=lora_scale,
-                device=device,
-            )
-            del module
-            setattr(model, module_name, quantized_module)
-        else:
-            _prepare_model_int4_weight_int8_dynamic_activation(module, group_size, lora_rank, lora_scale)
-
-    return model
-
-
-def convert_to_int4_quantized_model(
-    model: Transformer | CrossAttentionTransformer,
-    checkpoint_dir: str,
-    device: torch.device | None = None,
-) -> Transformer | CrossAttentionTransformer:
-    """Convert the model to int4 quantized model."""
-    model_args = model.params
-    assert model_args.quantization_args is not None, "Quantization args must be specified."
-    quantization_args = model_args.quantization_args
-    if quantization_args.scheme is None:
-        raise ValueError("Quantization scheme must be specified in 'quantization_args'.")
-
-    if quantization_args.scheme.value != "int4_weight_int8_dynamic_activation":
-        raise NotImplementedError(
-            "Only int4 quantization with 'int4_weight_int8_dynamic_activation' scheme is supported."
-        )
-
-    group_size = model_args.quantization_args.group_size
-    if group_size is None:
-        raise ValueError("'group_size' cannot be None in 'quantization_args'. Please specify it.")
-
-    if model_args.lora_args is None:
-        # Certain quantized models (e.g., SpinQuant) may not have LoRA.
-        lora_rank = None
-        lora_scale = None
-    else:
-        lora_rank = model_args.lora_args.rank
-        lora_scale = model_args.lora_args.scale
-
-    _prepare_model_int4_weight_int8_dynamic_activation(model, group_size, lora_rank, lora_scale)
-    return cast(Transformer | CrossAttentionTransformer, model.to(device=device))

llama_stack/models/llama/llama3_1/__init__.py

@@ -1,12 +0,0 @@
-# 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.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# top-level folder for each specific model found within the models/ directory at
-# the top-level of this source tree.