mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-10-04 04:04:14 +00:00
7519b73fcc
The starter distribution added post-training which added torch dependencies which pulls in all the nvidia CUDA libraries. This made our starter container very big. We have worked hard to keep the starter container small so it serves its purpose as a starter. This PR tries to get it back to its size by forking off duplicate "-gpu" providers for post-training. These forked providers are then used for a new `starter-gpu` distribution which can pull in all dependencies.
87 lines
3.4 KiB
Python
87 lines
3.4 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.
|
|
|
|
|
|
from typing import cast
|
|
|
|
from llama_stack.providers.datatypes import AdapterSpec, Api, InlineProviderSpec, ProviderSpec, remote_provider_spec
|
|
|
|
# We provide two versions of these providers so that distributions can package the appropriate version of torch.
|
|
# The CPU version is used for distributions that don't have GPU support -- they result in smaller container images.
|
|
torchtune_def = dict(
|
|
api=Api.post_training,
|
|
pip_packages=["torchtune==0.5.0", "torchao==0.8.0", "numpy"],
|
|
module="llama_stack.providers.inline.post_training.torchtune",
|
|
config_class="llama_stack.providers.inline.post_training.torchtune.TorchtunePostTrainingConfig",
|
|
api_dependencies=[
|
|
Api.datasetio,
|
|
Api.datasets,
|
|
],
|
|
description="TorchTune-based post-training provider for fine-tuning and optimizing models using Meta's TorchTune framework.",
|
|
)
|
|
|
|
huggingface_def = dict(
|
|
api=Api.post_training,
|
|
pip_packages=["trl", "transformers", "peft", "datasets"],
|
|
module="llama_stack.providers.inline.post_training.huggingface",
|
|
config_class="llama_stack.providers.inline.post_training.huggingface.HuggingFacePostTrainingConfig",
|
|
api_dependencies=[
|
|
Api.datasetio,
|
|
Api.datasets,
|
|
],
|
|
description="HuggingFace-based post-training provider for fine-tuning models using the HuggingFace ecosystem.",
|
|
)
|
|
|
|
|
|
def available_providers() -> list[ProviderSpec]:
|
|
return [
|
|
InlineProviderSpec(
|
|
**{
|
|
**torchtune_def,
|
|
"provider_type": "inline::torchtune-cpu",
|
|
"pip_packages": (
|
|
cast(list[str], torchtune_def["pip_packages"])
|
|
+ ["torch torchtune==0.5.0 torchao==0.8.0 --index-url https://download.pytorch.org/whl/cpu"]
|
|
),
|
|
},
|
|
),
|
|
InlineProviderSpec(
|
|
**{
|
|
**huggingface_def,
|
|
"provider_type": "inline::huggingface-cpu",
|
|
"pip_packages": (
|
|
cast(list[str], huggingface_def["pip_packages"])
|
|
+ ["torch --index-url https://download.pytorch.org/whl/cpu"]
|
|
),
|
|
},
|
|
),
|
|
InlineProviderSpec(
|
|
**{
|
|
**torchtune_def,
|
|
"provider_type": "inline::torchtune-gpu",
|
|
"pip_packages": (
|
|
cast(list[str], torchtune_def["pip_packages"]) + ["torch torchtune==0.5.0 torchao==0.8.0"]
|
|
),
|
|
},
|
|
),
|
|
InlineProviderSpec(
|
|
**{
|
|
**huggingface_def,
|
|
"provider_type": "inline::huggingface-gpu",
|
|
"pip_packages": (cast(list[str], huggingface_def["pip_packages"]) + ["torch"]),
|
|
},
|
|
),
|
|
remote_provider_spec(
|
|
api=Api.post_training,
|
|
adapter=AdapterSpec(
|
|
adapter_type="nvidia",
|
|
pip_packages=["requests", "aiohttp"],
|
|
module="llama_stack.providers.remote.post_training.nvidia",
|
|
config_class="llama_stack.providers.remote.post_training.nvidia.NvidiaPostTrainingConfig",
|
|
description="NVIDIA's post-training provider for fine-tuning models on NVIDIA's platform.",
|
|
),
|
|
),
|
|
]
|