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

llama_stack/distributions/meta-reference-gpu/config.yaml

@@ -1,140 +0,0 @@
-version: 2
-image_name: meta-reference-gpu
-apis:
-- agents
-- datasetio
-- eval
-- inference
-- safety
-- scoring
-- tool_runtime
-- vector_io
-providers:
-  inference:
-  - provider_id: meta-reference-inference
-    provider_type: inline::meta-reference
-    config:
-      model: ${env.INFERENCE_MODEL}
-      checkpoint_dir: ${env.INFERENCE_CHECKPOINT_DIR:=null}
-      quantization:
-        type: ${env.QUANTIZATION_TYPE:=bf16}
-      model_parallel_size: ${env.MODEL_PARALLEL_SIZE:=0}
-      max_batch_size: ${env.MAX_BATCH_SIZE:=1}
-      max_seq_len: ${env.MAX_SEQ_LEN:=4096}
-  - provider_id: sentence-transformers
-    provider_type: inline::sentence-transformers
-  vector_io:
-  - provider_id: faiss
-    provider_type: inline::faiss
-    config:
-      persistence:
-        namespace: vector_io::faiss
-        backend: kv_default
-  safety:
-  - provider_id: llama-guard
-    provider_type: inline::llama-guard
-    config:
-      excluded_categories: []
-  agents:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      persistence:
-        agent_state:
-          namespace: agents
-          backend: kv_default
-        responses:
-          table_name: responses
-          backend: sql_default
-          max_write_queue_size: 10000
-          num_writers: 4
-  eval:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      kvstore:
-        namespace: eval
-        backend: kv_default
-  datasetio:
-  - provider_id: huggingface
-    provider_type: remote::huggingface
-    config:
-      kvstore:
-        namespace: datasetio::huggingface
-        backend: kv_default
-  - provider_id: localfs
-    provider_type: inline::localfs
-    config:
-      kvstore:
-        namespace: datasetio::localfs
-        backend: kv_default
-  scoring:
-  - provider_id: basic
-    provider_type: inline::basic
-  - provider_id: llm-as-judge
-    provider_type: inline::llm-as-judge
-  - provider_id: braintrust
-    provider_type: inline::braintrust
-    config:
-      openai_api_key: ${env.OPENAI_API_KEY:=}
-  tool_runtime:
-  - provider_id: brave-search
-    provider_type: remote::brave-search
-    config:
-      api_key: ${env.BRAVE_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: tavily-search
-    provider_type: remote::tavily-search
-    config:
-      api_key: ${env.TAVILY_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: rag-runtime
-    provider_type: inline::rag-runtime
-  - provider_id: model-context-protocol
-    provider_type: remote::model-context-protocol
-storage:
-  backends:
-    kv_default:
-      type: kv_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db
-    sql_default:
-      type: sql_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db
-  stores:
-    metadata:
-      namespace: registry
-      backend: kv_default
-    inference:
-      table_name: inference_store
-      backend: sql_default
-      max_write_queue_size: 10000
-      num_writers: 4
-    conversations:
-      table_name: openai_conversations
-      backend: sql_default
-    prompts:
-      namespace: prompts
-      backend: kv_default
-registered_resources:
-  models:
-  - metadata: {}
-    model_id: ${env.INFERENCE_MODEL}
-    provider_id: meta-reference-inference
-    model_type: llm
-  - metadata:
-      embedding_dimension: 768
-    model_id: nomic-embed-text-v1.5
-    provider_id: sentence-transformers
-    model_type: embedding
-  shields: []
-  vector_dbs: []
-  datasets: []
-  scoring_fns: []
-  benchmarks: []
-  tool_groups:
-  - toolgroup_id: builtin::websearch
-    provider_id: tavily-search
-  - toolgroup_id: builtin::rag
-    provider_id: rag-runtime
-server:
-  port: 8321

llama_stack/distributions/meta-reference-gpu/doc_template.md

@@ -1,119 +0,0 @@
----
-orphan: true
----
-# Meta Reference GPU Distribution
-
-```{toctree}
-:maxdepth: 2
-:hidden:
-
-self
-```
-
-The `llamastack/distribution-{{ name }}` distribution consists of the following provider configurations:
-
-{{ providers_table }}
-
-Note that you need access to nvidia GPUs to run this distribution. This distribution is not compatible with CPU-only machines or machines with AMD GPUs.
-
-{% if run_config_env_vars %}
-### Environment Variables
-
-The following environment variables can be configured:
-
-{% for var, (default_value, description) in run_config_env_vars.items() %}
-- `{{ var }}`: {{ description }} (default: `{{ default_value }}`)
-{% endfor %}
-{% endif %}
-
-
-## Prerequisite: Downloading Models
-
-Please check that you have llama model checkpoints downloaded in `~/.llama` before proceeding. See [installation guide](../../references/llama_cli_reference/download_models.md) here to download the models using the Hugging Face CLI.
-```
-
-## Running the Distribution
-
-You can do this via venv or Docker which has a pre-built image.
-
-### Via Docker
-
-This method allows you to get started quickly without having to build the distribution code.
-
-```bash
-LLAMA_STACK_PORT=8321
-docker run \
-  -it \
-  --pull always \
-  --gpu all \
-  -p $LLAMA_STACK_PORT:$LLAMA_STACK_PORT \
-  -v ~/.llama:/root/.llama \
-  -e INFERENCE_MODEL=meta-llama/Llama-3.2-3B-Instruct \
-  llamastack/distribution-{{ name }} \
-  --port $LLAMA_STACK_PORT
-```
-
-If you are using Llama Stack Safety / Shield APIs, use:
-
-```bash
-docker run \
-  -it \
-  --pull always \
-  --gpu all \
-  -p $LLAMA_STACK_PORT:$LLAMA_STACK_PORT \
-  -v ~/.llama:/root/.llama \
-  -e INFERENCE_MODEL=meta-llama/Llama-3.2-3B-Instruct \
-  -e SAFETY_MODEL=meta-llama/Llama-Guard-3-1B \
-  llamastack/distribution-{{ name }} \
-  --port $LLAMA_STACK_PORT
-```
-
-### Via Docker with Custom Run Configuration
-
-You can also run the Docker container with a custom run configuration file by mounting it into the container:
-
-```bash
-# Set the path to your custom config.yaml file
-CUSTOM_RUN_CONFIG=/path/to/your/custom-config.yaml
-LLAMA_STACK_PORT=8321
-
-docker run \
-  -it \
-  --pull always \
-  --gpu all \
-  -p $LLAMA_STACK_PORT:$LLAMA_STACK_PORT \
-  -v ~/.llama:/root/.llama \
-  -v $CUSTOM_RUN_CONFIG:/app/custom-config.yaml \
-  -e RUN_CONFIG_PATH=/app/custom-config.yaml \
-  llamastack/distribution-{{ name }} \
-  --port $LLAMA_STACK_PORT
-```
-
-**Note**: The run configuration must be mounted into the container before it can be used. The `-v` flag mounts your local file into the container, and the `RUN_CONFIG_PATH` environment variable tells the entrypoint script which configuration to use.
-
-{% if run_configs %}
-Available run configurations for this distribution:
-{% for config in run_configs %}
-- `{{ config }}`
-{% endfor %}
-{% endif %}
-
-### Via venv
-
-Make sure you have the Llama Stack CLI available.
-
-```bash
-llama stack list-deps meta-reference-gpu | xargs -L1 uv pip install
-INFERENCE_MODEL=meta-llama/Llama-3.2-3B-Instruct \
-llama stack run distributions/{{ name }}/config.yaml \
-  --port 8321
-```
-
-If you are using Llama Stack Safety / Shield APIs, use:
-
-```bash
-INFERENCE_MODEL=meta-llama/Llama-3.2-3B-Instruct \
-SAFETY_MODEL=meta-llama/Llama-Guard-3-1B \
-llama stack run distributions/{{ name }}/run-with-safety.yaml \
-  --port 8321
-```

llama_stack/distributions/meta-reference-gpu/meta_reference.py

@@ -1,163 +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.
-
-from pathlib import Path
-
-from llama_stack.core.datatypes import (
-    BuildProvider,
-    ModelInput,
-    Provider,
-    ShieldInput,
-    ToolGroupInput,
-)
-from llama_stack.distributions.template import DistributionTemplate, RunConfigSettings
-from llama_stack.providers.inline.inference.meta_reference import (
-    MetaReferenceInferenceConfig,
-)
-from llama_stack.providers.inline.inference.sentence_transformers import (
-    SentenceTransformersInferenceConfig,
-)
-from llama_stack.providers.inline.vector_io.faiss.config import FaissVectorIOConfig
-from llama_stack_api import ModelType
-
-
-def get_distribution_template() -> DistributionTemplate:
-    providers = {
-        "inference": [BuildProvider(provider_type="inline::meta-reference")],
-        "vector_io": [
-            BuildProvider(provider_type="inline::faiss"),
-            BuildProvider(provider_type="remote::chromadb"),
-            BuildProvider(provider_type="remote::pgvector"),
-        ],
-        "safety": [BuildProvider(provider_type="inline::llama-guard")],
-        "agents": [BuildProvider(provider_type="inline::meta-reference")],
-        "eval": [BuildProvider(provider_type="inline::meta-reference")],
-        "datasetio": [
-            BuildProvider(provider_type="remote::huggingface"),
-            BuildProvider(provider_type="inline::localfs"),
-        ],
-        "scoring": [
-            BuildProvider(provider_type="inline::basic"),
-            BuildProvider(provider_type="inline::llm-as-judge"),
-            BuildProvider(provider_type="inline::braintrust"),
-        ],
-        "tool_runtime": [
-            BuildProvider(provider_type="remote::brave-search"),
-            BuildProvider(provider_type="remote::tavily-search"),
-            BuildProvider(provider_type="inline::rag-runtime"),
-            BuildProvider(provider_type="remote::model-context-protocol"),
-        ],
-    }
-    name = "meta-reference-gpu"
-    inference_provider = Provider(
-        provider_id="meta-reference-inference",
-        provider_type="inline::meta-reference",
-        config=MetaReferenceInferenceConfig.sample_run_config(
-            model="${env.INFERENCE_MODEL}",
-            checkpoint_dir="${env.INFERENCE_CHECKPOINT_DIR:=null}",
-        ),
-    )
-    embedding_provider = Provider(
-        provider_id="sentence-transformers",
-        provider_type="inline::sentence-transformers",
-        config=SentenceTransformersInferenceConfig.sample_run_config(),
-    )
-    vector_io_provider = Provider(
-        provider_id="faiss",
-        provider_type="inline::faiss",
-        config=FaissVectorIOConfig.sample_run_config(f"~/.llama/distributions/{name}"),
-    )
-
-    inference_model = ModelInput(
-        model_id="${env.INFERENCE_MODEL}",
-        provider_id="meta-reference-inference",
-    )
-    embedding_model = ModelInput(
-        model_id="nomic-embed-text-v1.5",
-        provider_id="sentence-transformers",
-        model_type=ModelType.embedding,
-        metadata={
-            "embedding_dimension": 768,
-        },
-    )
-    safety_model = ModelInput(
-        model_id="${env.SAFETY_MODEL}",
-        provider_id="meta-reference-safety",
-    )
-    default_tool_groups = [
-        ToolGroupInput(
-            toolgroup_id="builtin::websearch",
-            provider_id="tavily-search",
-        ),
-        ToolGroupInput(
-            toolgroup_id="builtin::rag",
-            provider_id="rag-runtime",
-        ),
-    ]
-
-    return DistributionTemplate(
-        name=name,
-        distro_type="self_hosted",
-        description="Use Meta Reference for running LLM inference",
-        template_path=Path(__file__).parent / "doc_template.md",
-        providers=providers,
-        run_configs={
-            "config.yaml": RunConfigSettings(
-                provider_overrides={
-                    "inference": [inference_provider, embedding_provider],
-                    "vector_io": [vector_io_provider],
-                },
-                default_models=[inference_model, embedding_model],
-                default_tool_groups=default_tool_groups,
-            ),
-            "run-with-safety.yaml": RunConfigSettings(
-                provider_overrides={
-                    "inference": [
-                        inference_provider,
-                        embedding_provider,
-                        Provider(
-                            provider_id="meta-reference-safety",
-                            provider_type="inline::meta-reference",
-                            config=MetaReferenceInferenceConfig.sample_run_config(
-                                model="${env.SAFETY_MODEL}",
-                                checkpoint_dir="${env.SAFETY_CHECKPOINT_DIR:=null}",
-                            ),
-                        ),
-                    ],
-                    "vector_io": [vector_io_provider],
-                },
-                default_models=[
-                    inference_model,
-                    safety_model,
-                    embedding_model,
-                ],
-                default_shields=[ShieldInput(shield_id="${env.SAFETY_MODEL}")],
-                default_tool_groups=default_tool_groups,
-            ),
-        },
-        run_config_env_vars={
-            "LLAMA_STACK_PORT": (
-                "8321",
-                "Port for the Llama Stack distribution server",
-            ),
-            "INFERENCE_MODEL": (
-                "meta-llama/Llama-3.2-3B-Instruct",
-                "Inference model loaded into the Meta Reference server",
-            ),
-            "INFERENCE_CHECKPOINT_DIR": (
-                "null",
-                "Directory containing the Meta Reference model checkpoint",
-            ),
-            "SAFETY_MODEL": (
-                "meta-llama/Llama-Guard-3-1B",
-                "Name of the safety (Llama-Guard) model to use",
-            ),
-            "SAFETY_CHECKPOINT_DIR": (
-                "null",
-                "Directory containing the Llama-Guard model checkpoint",
-            ),
-        },
-    )

llama_stack/distributions/meta-reference-gpu/run-with-safety.yaml

@@ -1,155 +0,0 @@
-version: 2
-image_name: meta-reference-gpu
-apis:
-- agents
-- datasetio
-- eval
-- inference
-- safety
-- scoring
-- tool_runtime
-- vector_io
-providers:
-  inference:
-  - provider_id: meta-reference-inference
-    provider_type: inline::meta-reference
-    config:
-      model: ${env.INFERENCE_MODEL}
-      checkpoint_dir: ${env.INFERENCE_CHECKPOINT_DIR:=null}
-      quantization:
-        type: ${env.QUANTIZATION_TYPE:=bf16}
-      model_parallel_size: ${env.MODEL_PARALLEL_SIZE:=0}
-      max_batch_size: ${env.MAX_BATCH_SIZE:=1}
-      max_seq_len: ${env.MAX_SEQ_LEN:=4096}
-  - provider_id: sentence-transformers
-    provider_type: inline::sentence-transformers
-  - provider_id: meta-reference-safety
-    provider_type: inline::meta-reference
-    config:
-      model: ${env.SAFETY_MODEL}
-      checkpoint_dir: ${env.SAFETY_CHECKPOINT_DIR:=null}
-      quantization:
-        type: ${env.QUANTIZATION_TYPE:=bf16}
-      model_parallel_size: ${env.MODEL_PARALLEL_SIZE:=0}
-      max_batch_size: ${env.MAX_BATCH_SIZE:=1}
-      max_seq_len: ${env.MAX_SEQ_LEN:=4096}
-  vector_io:
-  - provider_id: faiss
-    provider_type: inline::faiss
-    config:
-      persistence:
-        namespace: vector_io::faiss
-        backend: kv_default
-  safety:
-  - provider_id: llama-guard
-    provider_type: inline::llama-guard
-    config:
-      excluded_categories: []
-  agents:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      persistence:
-        agent_state:
-          namespace: agents
-          backend: kv_default
-        responses:
-          table_name: responses
-          backend: sql_default
-          max_write_queue_size: 10000
-          num_writers: 4
-  eval:
-  - provider_id: meta-reference
-    provider_type: inline::meta-reference
-    config:
-      kvstore:
-        namespace: eval
-        backend: kv_default
-  datasetio:
-  - provider_id: huggingface
-    provider_type: remote::huggingface
-    config:
-      kvstore:
-        namespace: datasetio::huggingface
-        backend: kv_default
-  - provider_id: localfs
-    provider_type: inline::localfs
-    config:
-      kvstore:
-        namespace: datasetio::localfs
-        backend: kv_default
-  scoring:
-  - provider_id: basic
-    provider_type: inline::basic
-  - provider_id: llm-as-judge
-    provider_type: inline::llm-as-judge
-  - provider_id: braintrust
-    provider_type: inline::braintrust
-    config:
-      openai_api_key: ${env.OPENAI_API_KEY:=}
-  tool_runtime:
-  - provider_id: brave-search
-    provider_type: remote::brave-search
-    config:
-      api_key: ${env.BRAVE_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: tavily-search
-    provider_type: remote::tavily-search
-    config:
-      api_key: ${env.TAVILY_SEARCH_API_KEY:=}
-      max_results: 3
-  - provider_id: rag-runtime
-    provider_type: inline::rag-runtime
-  - provider_id: model-context-protocol
-    provider_type: remote::model-context-protocol
-storage:
-  backends:
-    kv_default:
-      type: kv_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/kvstore.db
-    sql_default:
-      type: sql_sqlite
-      db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/meta-reference-gpu}/sql_store.db
-  stores:
-    metadata:
-      namespace: registry
-      backend: kv_default
-    inference:
-      table_name: inference_store
-      backend: sql_default
-      max_write_queue_size: 10000
-      num_writers: 4
-    conversations:
-      table_name: openai_conversations
-      backend: sql_default
-    prompts:
-      namespace: prompts
-      backend: kv_default
-registered_resources:
-  models:
-  - metadata: {}
-    model_id: ${env.INFERENCE_MODEL}
-    provider_id: meta-reference-inference
-    model_type: llm
-  - metadata: {}
-    model_id: ${env.SAFETY_MODEL}
-    provider_id: meta-reference-safety
-    model_type: llm
-  - metadata:
-      embedding_dimension: 768
-    model_id: nomic-embed-text-v1.5
-    provider_id: sentence-transformers
-    model_type: embedding
-  shields:
-  - shield_id: ${env.SAFETY_MODEL}
-  vector_dbs: []
-  datasets: []
-  scoring_fns: []
-  benchmarks: []
-  tool_groups:
-  - toolgroup_id: builtin::websearch
-    provider_id: tavily-search
-  - toolgroup_id: builtin::rag
-    provider_id: rag-runtime
-server:
-  port: 8321

llama_stack/models/llama/hadamard_utils.py

@@ -1,88 +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
-import re
-
-import torch
-from torch import nn
-
-
-def hadamard_transform(x: torch.Tensor) -> torch.Tensor:
-    """Hadamard transform.
-
-    This function performs the Hadamard transform on the input tensor 'x'.
-    The Hadamard transform is a linear transformation that multiplies the input
-    tensor by the Hadamard matrix of dimension n x n, where n is the size of
-    the last dimension of the input tensor.
-    """
-    *_, n = x.shape
-    m = int(math.log2(n))
-    assert n == 1 << m, "n must be a power of 2"
-    x = x[..., None]
-    inv_sqrt2 = 0.5**0.5
-    for _ in range(m):
-        top = x[..., ::2, :] + x[..., 1::2, :]
-        bot = x[..., ::2, :] - x[..., 1::2, :]
-        x = torch.cat((top, bot), dim=-1)
-        x *= inv_sqrt2
-    res = x.squeeze(-2)
-    return res
-
-
-class HadamardModule(torch.nn.Module):
-    """A module that applies the Hadamard transform to the input tensor.
-
-    Args:
-        group_size: The size of the groups that the input tensor will be divided into
-            before applying the Hadamard transform.
-    """
-
-    def __init__(self, group_size: int) -> None:
-        super().__init__()
-        self.group_size = group_size
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        reshape_back = False
-        orig_shape = x.shape
-        if self.group_size != x.shape[-1]:
-            reshape_back = True
-            x = x.reshape(-1, x.shape[-1] // self.group_size, self.group_size)
-        x = hadamard_transform(x)
-        if reshape_back:
-            x = x.reshape(orig_shape)
-        return x
-
-
-def add_hadamard_transform_for_spinquant(model: torch.nn.Module, prefix: str = "") -> None:
-    """
-    Adds a Hadamard transform to the last linear layer of each feedforward network (FFN) in the model.
-    This function recursively traverses the model's children and looks for layers that match the pattern
-    "layers.<digit>.feed_forward.w2", where <digit> is one or more digits. When such a layer is found,
-    it is replaced with a new sequential module that consists of a HadamardModule followed by the original
-    layer. The HadamardModule applies the Hadamard transform to the input tensor.
-
-    See `SpinQuant <https://arxiv.org/abs/2405.16406>_` paper for more details.
-
-    Args:
-        model: An instance of 'torch.nn.Module' (e.g., Transformer model).
-        prefix: A string prefix to add to the full name of each child module.
-
-    Returns:
-        None
-    """
-
-    pattern_last_linear_ffn = r"layers.\d+.feed_forward.w2"
-    for module_name, module in model.named_children():
-        child_full_name = prefix + "." + module_name
-        if re.search(pattern_last_linear_ffn, child_full_name):
-            # Module matching this pattern should be nn.Linear with in_features
-            assert isinstance(module, nn.Linear), f"Expected nn.Linear, got {type(module)}"
-            new_module = nn.Sequential(HadamardModule(group_size=module.in_features), module)
-            del module
-            setattr(model, module_name, new_module)
-        else:
-            add_hadamard_transform_for_spinquant(module, (prefix + "." if prefix else prefix) + module_name)