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

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

@@ -1,226 +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.
-
-import os
-from collections.abc import Callable
-
-import torch
-from fairscale.nn.model_parallel.initialize import get_model_parallel_rank
-from torch import Tensor, nn
-from torch.nn import functional as F
-
-from llama_stack.log import get_logger
-
-from ...datatypes import QuantizationMode
-from ..model import Transformer, TransformerBlock
-from ..moe import MoE
-
-log = get_logger(name=__name__, category="models::llama")
-
-
-def swiglu_wrapper_no_reduce(
-    self,
-    x: Tensor,
-):
-    from ...quantize_impls import ffn_swiglu
-
-    return ffn_swiglu(x, self.w1.weight, self.w3.weight, self.w2.weight)
-
-
-def experts_batched_swiglu_wrapper(
-    self,
-    x: Tensor,  # (e, g, D)
-    w1: Tensor,  # (e, D, F)
-    w3: Tensor,  # (e, D, F)
-    w2: Tensor,  # (e, F, D)
-) -> torch.Tensor:
-    from ...quantize_impls import bmm_nt
-
-    middle_out_egF = F.silu(bmm_nt(x, w1)) * bmm_nt(x, w3)  # noqa: N806
-    return bmm_nt(middle_out_egF, w2)
-
-
-def convert_to_quantized_model(
-    model: Transformer,
-    checkpoint_dir: str,
-    quantization_mode: str | None = None,
-    fp8_activation_scale_ub: float | None = 1200.0,
-    use_rich_progress: bool = True,
-) -> Transformer:
-    from ...quantize_impls import (
-        Fp8ScaledWeights,
-        Int4ScaledWeights,
-        load_fp8,
-        load_int4,
-        quantize_fp8,
-        quantize_int4,
-    )
-
-    rank = get_model_parallel_rank()
-
-    def should_quantize_block(block: nn.Module) -> bool:
-        if not isinstance(block, TransformerBlock):
-            return False
-
-        is_moe = isinstance(block.feed_forward, MoE)
-        if quantization_mode == QuantizationMode.fp8_mixed:
-            # skip quantization on first and last layers
-            return is_moe and not (block.layer_id == 0 or block.layer_id == (model.n_layers - 1))
-
-        return is_moe
-
-    use_rich_progress = use_rich_progress and rank == 0
-    progress, log_status, update_status = logging_callbacks(use_rich_progress, rank, model, should_quantize_block)
-    if quantization_mode == QuantizationMode.int4_mixed:
-        int4_scales_path = os.path.join(checkpoint_dir, f"int4_scales_{rank}.pt")
-        if os.path.isfile(int4_scales_path):
-            log_status(f"Rank {rank}: Loading int4 scales")
-            int4_scales = torch.load(int4_scales_path, weights_only=True)
-
-            def apply_quantization(key, weight):
-                scale = int4_scales[key]
-                return load_int4(
-                    weight,
-                    scale,
-                    output_device=torch.device("cuda"),
-                )
-
-        else:
-            log_status(f"Rank {rank}: Quantizing int4 weights from bf16")
-
-            def apply_quantization(_, weight):
-                return quantize_int4(weight, output_device=torch.device("cuda"))
-
-    else:
-        fp8_scales_path = os.path.join(checkpoint_dir, f"fp8_scales_{rank}.pt")
-        if os.path.isfile(fp8_scales_path):
-            log_status(f"Rank {rank}: Loading fp8 scales")
-            fp8_scales = torch.load(fp8_scales_path, weights_only=True)
-
-            def apply_quantization(key, weight):
-                scale = fp8_scales[key]
-                return load_fp8(
-                    weight,
-                    scale,
-                    fp8_activation_scale_ub,
-                    output_device=torch.device("cuda"),
-                )
-
-        else:
-            log_status(f"Rank {rank}: Quantizing fp8 weights from bf16")
-
-            def apply_quantization(_, weight):
-                return quantize_fp8(weight, fp8_activation_scale_ub, output_device=torch.device("cuda"))
-
-    processed_blocks = 0
-    try:
-        if use_rich_progress:
-            progress.start()
-
-        for _, block in model.named_modules():
-            if not should_quantize_block(block):
-                continue
-
-            update_status(f"Rank {rank} - Layer {block.layer_id}")
-
-            # Quantize only routed experts, not shared
-            prefix = f"layers.{block.layer_id}.feed_forward"
-            moe = block.feed_forward
-            moe.experts.batched_swiglu = experts_batched_swiglu_wrapper.__get__(moe.experts)
-
-            for key in ("w1", "w3", "w2"):
-                param = getattr(moe.experts, key)
-                update_status(f"Rank {rank} - Layer {block.layer_id} - MoE {key}")
-                setattr(
-                    moe.experts,
-                    key,
-                    apply_quantization(
-                        f"{prefix}.experts.{key}",
-                        param.transpose(1, 2).contiguous(),
-                    ),
-                )
-
-            if quantization_mode == QuantizationMode.int4_mixed:
-                # Quantize shared experts
-                moe.shared_expert.forward = swiglu_wrapper_no_reduce.__get__(moe.shared_expert)
-                for key in ("w1", "w3", "w2"):
-                    param = getattr(moe.shared_expert, key)
-                    update_status(f"Rank {rank} - Layer {block.layer_id} - MoE shared expert {key}")
-                    param.weight = apply_quantization(f"{prefix}.shared_expert.{key}", param.weight)
-
-            processed_blocks += 1
-            update_status(message=None, completed=processed_blocks)
-
-        update_status(f"Rank {rank} - Moving parameters to CUDA")
-
-        param_count = 0
-        for _, parameter in model.named_parameters():
-            if not isinstance(parameter, Fp8ScaledWeights) and not isinstance(parameter, Int4ScaledWeights):
-                parameter.data = parameter.to(device="cuda")
-                param_count += 1
-
-        update_status(f"Rank {rank} - Completed - moved {param_count} parameters to CUDA")
-    finally:
-        if use_rich_progress:
-            progress.stop()
-
-    return model
-
-
-# fp8/int4 loading can be very slow so we add progress bars to make life slightly better
-def logging_callbacks(
-    use_rich_progress: bool,
-    rank: int,
-    model: Transformer,
-    should_quantize_block: Callable[[nn.Module], bool],
-):
-    console = None
-    if use_rich_progress:
-        from rich.console import Console
-
-        console = Console(highlight=False)
-
-    def log_status(message: str) -> None:
-        if use_rich_progress:
-            console.print(message)
-        elif rank == 0:  # Only log from rank 0 for non-rich logging
-            log.info(message)
-
-    total_blocks = sum(1 for _, block in model.named_modules() if should_quantize_block(block))
-    progress = None
-    if use_rich_progress:
-        from rich.progress import (
-            BarColumn,
-            Progress,
-            SpinnerColumn,
-            TextColumn,
-            TimeElapsedColumn,
-            TimeRemainingColumn,
-        )
-
-        progress = Progress(
-            SpinnerColumn(),
-            BarColumn(complete_style="green", finished_style="bright_green"),
-            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
-            TimeElapsedColumn(),
-            TextColumn("ETA:"),
-            TimeRemainingColumn(),
-            TextColumn("[bold]{task.fields[status]}"),
-            console=console,
-            expand=True,
-        )
-        task_id = progress.add_task("[blue]Converting layers...", total=total_blocks, status="Starting")
-
-    def update_status(message: str | None, completed: int | None = None) -> None:
-        if use_rich_progress:
-            if message is not None:
-                progress.update(task_id, status=message)
-            if completed is not None:
-                progress.update(task_id, completed=completed)
-        elif rank == 0 and completed and completed % 10 == 0:
-            log.info(f"Rank {rank}: {completed}/{total_blocks} blocks completed")
-
-    return progress, log_status, update_status

llama_stack/models/llama/llama4/vision/__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/llama4/vision/embedding.py

@@ -1,210 +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.
-
-import math
-from collections.abc import Callable
-from typing import Any
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from fairscale.nn.model_parallel.layers import ColumnParallelLinear, RowParallelLinear
-
-from ..args import VisionArgs
-from .encoder import VisionEncoder
-
-
-class PixelShuffle(nn.Module):
-    def __init__(self, ps_ratio):
-        super().__init__()
-        self.ps_ratio = ps_ratio
-
-    def forward(self, x):
-        # x: [B, N, C], N = number of patches
-        assert self.ps_ratio is not None, "ps_ratio is required for pixel shuffle"
-        assert x.dim() == 3, "pixel shuffle requires encoded patches [B, N, C]"
-        hh = ww = int(math.sqrt(x.shape[1]))
-        x = x.reshape(x.shape[0], hh, ww, -1)
-        x = pixel_shuffle_op(x, ps_ratio=self.ps_ratio)
-        pixel_shuffle_patches = x.reshape(x.shape[0], -1, x.shape[-1])
-        return pixel_shuffle_patches
-
-
-def pixel_shuffle_op(input_x, ps_ratio):
-    n, w, h, c = input_x.size()
-    input_x = input_x.view(n, w, int(h * ps_ratio), int(c / ps_ratio))
-    input_x = input_x.permute(0, 2, 1, 3).contiguous()
-    input_x = input_x.view(
-        n,
-        int(h * ps_ratio),
-        int(w * ps_ratio),
-        int(c / (ps_ratio * ps_ratio)),
-    )
-    input_x = input_x.permute(0, 2, 1, 3).contiguous()
-    return input_x
-
-
-class SimpleMLP(torch.nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        hidden_dim: int,
-        bias: bool = True,
-        dropout: float = 0.0,
-        act_layer: Callable = nn.GELU,
-    ):
-        super().__init__()
-        # layers
-        self.c_fc = ColumnParallelLinear(
-            dim,
-            hidden_dim,
-            bias=bias,
-            gather_output=False,
-        )
-        self.c_proj = RowParallelLinear(
-            hidden_dim,
-            hidden_dim,
-            bias=bias,
-            input_is_parallel=True,
-        )
-        self.non_linearity = act_layer()
-        self.dropout = dropout
-
-    def forward(self, x):
-        hidden = self.c_fc(x)
-        hidden = self.non_linearity(hidden)
-        hidden = F.dropout(hidden, p=self.dropout, training=self.training)
-        return self.non_linearity(self.c_proj(hidden))
-
-
-class PixelShuffleMLP(torch.nn.Module):
-    def __init__(
-        self,
-        ps_ratio: float,
-        input_dim: int,
-        output_dim: int = 4096,
-        add_fc: bool = False,
-    ):
-        super().__init__()
-        self.pixel_shuffle = PixelShuffle(ps_ratio)
-        self.mlp = SimpleMLP(
-            int(input_dim // (ps_ratio**2)),
-            output_dim,
-            bias=False,
-            dropout=0.0,
-            act_layer=nn.GELU,
-        )
-        self.fc = nn.Identity()
-        if add_fc:
-            self.fc = ColumnParallelLinear(
-                output_dim,
-                output_dim,
-                bias=False,
-            )
-
-    def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:
-        encoded_patches = self.pixel_shuffle(encoded_patches)
-        return self.fc(self.mlp(encoded_patches))
-
-
-class VisionEmbeddings(torch.nn.Module):
-    def __init__(self, args: VisionArgs):
-        super().__init__()
-        self.args = args
-
-        image_size = args.image_size
-        patch_size = args.patch_size
-        self.vision_encoder = VisionEncoder(
-            image_size=(image_size.height, image_size.width),
-            patch_size=(patch_size.height, patch_size.width),
-            dim=args.dim,
-            layers=args.n_layers,
-            heads=args.n_heads,
-            mlp_ratio=args.mlp_ratio,
-        )
-        self.vision_encoder = self.vision_encoder.to(torch.bfloat16)
-        self.vision_adapter = PixelShuffleMLP(
-            ps_ratio=args.pixel_shuffle_ratio,
-            input_dim=args.dim,
-            output_dim=args.output_dim,
-        )
-
-        self.output_dim = args.output_dim
-        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 = True,
-        missing_keys: list[str] = None,
-        unexpected_keys: list[str] = None,
-        error_msgs: list[str] = None,
-        return_state_dict: bool = False,
-    ) -> None:
-        original_sd = self.state_dict()
-        for k in state_dict:
-            if k.startswith(prefix) and len(state_dict[k].shape) == 1 and state_dict[k].shape[0] == 0:
-                state_dict[k] = state_dict[k].reshape(original_sd[k[len(prefix) :]].shape)
-
-    def _get_empty_sequence(self, h):
-        return torch.zeros(
-            h.shape[0],
-            h.shape[1],
-            self.output_dim,
-            device=h.device,
-            dtype=h.dtype,
-        )
-
-    # x_images is batched; each batch sample contains a list of images. so this is List[List[torch.Tensor]]
-    # each image is a tensor of shape [num_tiles, C, H, W]
-    def forward(
-        self,
-        image_batch: list[list[torch.Tensor]],
-        image_mask: torch.Tensor,
-        h_ref: torch.Tensor,
-    ) -> torch.Tensor:
-        images_flattened = [image for sample in image_batch for image in sample]
-        images_flattened = torch.vstack(images_flattened).unsqueeze(1).to(h_ref.dtype).to(h_ref.device)
-        embedding = self.vision_encoder(images_flattened)
-        projected_embedding = self.vision_adapter(embedding)
-
-        h_image = self._get_empty_sequence(h_ref)
-        return scatter_embeddings(image_batch, image_mask, h_image, projected_embedding)
-
-
-def scatter_embeddings(image_batch, image_mask, h_image, encoded_patches_proj):
-    # If dynamic transform is used and the batch contains 2 images (where image_1 has 2 chunks and image_2 has 3 chunks),
-    # `num_images_per_sequence` now records the number of chunks per image as `[2, 3]`.
-    # `encoded_patches_proj.split` will then split the image chunks into 2 groups: `[image_1_chunks, image_2_chunks]`.
-    num_images_per_sequence = [sum(image.size(0) for image in sample_images) for sample_images in image_batch]
-
-    assert not torch.isnan(encoded_patches_proj).any()
-    assert sum(num_images_per_sequence) == encoded_patches_proj.size(0), (
-        f"{sum(num_images_per_sequence)=} != {encoded_patches_proj.shape=}"
-    )
-
-    encoded_patches_list = encoded_patches_proj.split(num_images_per_sequence, dim=0)
-    for index in range(h_image.size(0)):
-        encoded_patches_per_sample = encoded_patches_list[index]
-        sample_image_mask = image_mask[index]
-
-        if encoded_patches_per_sample.numel() == 0:
-            continue
-        encoded_patches_per_sample = encoded_patches_per_sample.contiguous().view(
-            -1, encoded_patches_per_sample.size(-1)
-        )
-
-        n_tokens_to_fill = sample_image_mask.sum()
-        assert n_tokens_to_fill <= encoded_patches_per_sample.size(0)
-
-        h_image[index].masked_scatter_(
-            sample_image_mask.expand(-1, h_image.size(-1)),
-            encoded_patches_per_sample[:n_tokens_to_fill],
-        )
-
-    return h_image