llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/inline/post_training/huggingface/utils.py

@@ -0,0 +1,269 @@
+# 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
+import signal
+import sys
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import psutil
+import torch
+from datasets import Dataset
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from llama_stack.apis.datasetio import DatasetIO
+from llama_stack.apis.post_training import Checkpoint, TrainingConfig
+from llama_stack.log import get_logger
+
+from .config import HuggingFacePostTrainingConfig
+
+logger = get_logger(name=__name__, category="post_training")
+
+
+def setup_environment():
+    """Setup common environment variables for training."""
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+    os.environ["MKL_THREADING_LAYER"] = "GNU"
+    os.environ["MKL_SERVICE_FORCE_INTEL"] = "0"
+    os.environ["MKL_NUM_THREADS"] = "1"
+
+
+def bytes_to_gb(to_convert: int) -> str:
+    """Converts memory stats to GB and formats to 2 decimal places.
+    Args:
+        to_convert: Memory value in bytes
+    Returns:
+        str: Memory value in GB formatted to 2 decimal places
+    """
+    return f"{(to_convert / (1024**3)):.2f}"
+
+
+def get_memory_stats(device: torch.device) -> dict[str, Any]:
+    """Get memory statistics for the given device."""
+    stats = {
+        "system_memory": {
+            "total": bytes_to_gb(psutil.virtual_memory().total),
+            "available": bytes_to_gb(psutil.virtual_memory().available),
+            "used": bytes_to_gb(psutil.virtual_memory().used),
+            "percent": psutil.virtual_memory().percent,
+        }
+    }
+
+    if device.type == "cuda":
+        stats["device_memory"] = {
+            "allocated": bytes_to_gb(torch.cuda.memory_allocated(device)),
+            "reserved": bytes_to_gb(torch.cuda.memory_reserved(device)),
+            "max_allocated": bytes_to_gb(torch.cuda.max_memory_allocated(device)),
+        }
+    elif device.type == "mps":
+        # MPS doesn't provide direct memory stats, but we can track system memory
+        stats["device_memory"] = {
+            "note": "MPS memory stats not directly available",
+            "system_memory_used": bytes_to_gb(psutil.virtual_memory().used),
+        }
+    elif device.type == "cpu":
+        # For CPU, we track process memory usage
+        process = psutil.Process()
+        stats["device_memory"] = {
+            "process_rss": bytes_to_gb(process.memory_info().rss),
+            "process_vms": bytes_to_gb(process.memory_info().vms),
+            "process_percent": process.memory_percent(),
+        }
+
+    return stats
+
+
+def setup_torch_device(device_str: str) -> torch.device:
+    """Initialize and validate a PyTorch device.
+    This function handles device initialization and validation for different device types:
+    - CUDA: Validates CUDA availability and handles device selection
+    - MPS: Validates MPS availability for Apple Silicon
+    - CPU: Basic validation
+    - HPU: Raises error as it's not supported
+    Args:
+        device_str: String specifying the device ('cuda', 'cpu', 'mps')
+    Returns:
+        torch.device: The initialized and validated device
+    Raises:
+        RuntimeError: If device initialization fails or device is not supported
+    """
+    try:
+        device = torch.device(device_str)
+    except RuntimeError as e:
+        raise RuntimeError(f"Error getting Torch Device {str(e)}") from e
+
+    # Validate device capabilities
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError(
+                f"{device.type}: Torch has no CUDA/ROCm support or could not detect a compatible device."
+            )
+        if device.index is None:
+            device = torch.device(device.type, torch.cuda.current_device())
+    elif device.type == "mps":
+        if not torch.backends.mps.is_available():
+            raise RuntimeError(f"{device.type}: Torch has no MPS support or could not detect a compatible device.")
+    elif device.type == "hpu":
+        raise RuntimeError(f"{device.type}: training does not support Intel Gaudi.")
+
+    return device
+
+
+async def load_rows_from_dataset(datasetio_api: DatasetIO, dataset_id: str) -> list[dict[str, Any]]:
+    """Load dataset from llama stack dataset provider"""
+    try:
+        all_rows = await datasetio_api.iterrows(
+            dataset_id=dataset_id,
+            limit=-1,
+        )
+        if not isinstance(all_rows.data, list):
+            raise RuntimeError("Expected dataset data to be a list")
+        return all_rows.data
+    except Exception as e:
+        raise RuntimeError(f"Failed to load dataset: {str(e)}") from e
+
+
+def load_model(
+    model: str,
+    device: torch.device,
+    provider_config: HuggingFacePostTrainingConfig,
+) -> AutoModelForCausalLM:
+    """Load and initialize the model for training.
+    Args:
+        model: The model identifier to load
+        device: The device to load the model onto
+        provider_config: Provider-specific configuration
+    Returns:
+        The loaded and initialized model
+    Raises:
+        RuntimeError: If model loading fails
+    """
+    logger.info("Loading the base model")
+    try:
+        model_config = AutoConfig.from_pretrained(model, **provider_config.model_specific_config)
+        model_obj = AutoModelForCausalLM.from_pretrained(
+            model,
+            torch_dtype="auto" if device.type != "cpu" else "float32",
+            quantization_config=None,
+            config=model_config,
+            **provider_config.model_specific_config,
+        )
+        # Always move model to specified device
+        model_obj = model_obj.to(device)
+        logger.info(f"Model loaded and moved to device: {model_obj.device}")
+        return model_obj
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model: {str(e)}") from e
+
+
+def split_dataset(ds: Dataset) -> tuple[Dataset, Dataset]:
+    """Split dataset into train and validation sets.
+    Args:
+        ds: Dataset to split
+    Returns:
+        tuple: (train_dataset, eval_dataset)
+    """
+    logger.info("Splitting dataset into train and validation sets")
+    train_val_split = ds.train_test_split(test_size=0.1, seed=42)
+    train_dataset = train_val_split["train"]
+    eval_dataset = train_val_split["test"]
+    logger.info(f"Split dataset into {len(train_dataset)} training and {len(eval_dataset)} validation examples")
+    return train_dataset, eval_dataset
+
+
+def setup_signal_handlers():
+    """Setup signal handlers for graceful shutdown."""
+
+    def signal_handler(signum, frame):
+        logger.info(f"Received signal {signum}, initiating graceful shutdown")
+        sys.exit(0)
+
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+
+def calculate_training_steps(steps_per_epoch: int, config: TrainingConfig) -> dict[str, int]:
+    """Calculate training steps and logging configuration.
+    Args:
+        steps_per_epoch: Number of training steps per epoch
+        config: Training configuration
+    Returns:
+        dict: Dictionary with calculated step values
+    """
+    total_steps = steps_per_epoch * config.n_epochs
+    max_steps = min(config.max_steps_per_epoch, total_steps)
+    logging_steps = max(1, steps_per_epoch // 50)  # Log 50 times per epoch
+
+    logger.info("Training configuration:")
+    logger.info(f"- Steps per epoch: {steps_per_epoch}")
+    logger.info(f"- Total steps: {total_steps}")
+    logger.info(f"- Max steps: {max_steps}")
+    logger.info(f"- Logging steps: {logging_steps}")
+
+    return {"total_steps": total_steps, "max_steps": max_steps, "logging_steps": logging_steps}
+
+
+def get_save_strategy(output_dir_path: Path | None) -> tuple[str, str]:
+    """Get save and evaluation strategy based on output directory.
+    Args:
+        output_dir_path: Optional path to save the model
+    Returns:
+        tuple: (save_strategy, eval_strategy)
+    """
+    if output_dir_path:
+        logger.info(f"Will save checkpoints to {output_dir_path}")
+        return "epoch", "epoch"
+    return "no", "no"
+
+
+def create_checkpoints(
+    output_dir_path: Path, job_uuid: str, model: str, config: TrainingConfig, final_model_name: str
+) -> list[Checkpoint]:
+    """Create checkpoint objects from training output.
+    Args:
+        output_dir_path: Path to the training output directory
+        job_uuid: Unique identifier for the training job
+        model: Model identifier
+        config: Training configuration
+        final_model_name: Name of the final model directory ("merged_model" for SFT, "dpo_model" for DPO)
+    Returns:
+        List of Checkpoint objects
+    """
+    checkpoints = []
+
+    # Add checkpoint directories
+    checkpoint_dirs = sorted(
+        [d for d in output_dir_path.glob("checkpoint-*") if d.is_dir()],
+        key=lambda x: int(x.name.split("-")[1]),
+    )
+
+    for epoch_number, checkpoint_dir in enumerate(checkpoint_dirs, start=1):
+        created_time = datetime.fromtimestamp(os.path.getctime(checkpoint_dir), tz=UTC)
+        checkpoint = Checkpoint(
+            identifier=checkpoint_dir.name,
+            created_at=created_time,
+            epoch=epoch_number,
+            post_training_job_id=job_uuid,
+            path=str(checkpoint_dir),
+        )
+        checkpoints.append(checkpoint)
+
+    # Add final model
+    final_model_path = output_dir_path / final_model_name
+    if final_model_path.exists():
+        training_type = "sft" if final_model_name == "merged_model" else "dpo"
+        checkpoint = Checkpoint(
+            identifier=f"{model}-{training_type}-{config.n_epochs}",
+            created_at=datetime.now(UTC),
+            epoch=config.n_epochs,
+            post_training_job_id=job_uuid,
+            path=str(final_model_path),
+        )
+        checkpoints.append(checkpoint)
+
+    return checkpoints

llama_stack/providers/inline/post_training/torchtune/__init__.py

@@ -0,0 +1,27 @@
+# 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 Any
+
+from llama_stack.core.datatypes import Api
+
+from .config import TorchtunePostTrainingConfig
+
+# post_training api and the torchtune provider is still experimental and under heavy development
+
+
+async def get_provider_impl(
+    config: TorchtunePostTrainingConfig,
+    deps: dict[Api, Any],
+):
+    from .post_training import TorchtunePostTrainingImpl
+
+    impl = TorchtunePostTrainingImpl(
+        config,
+        deps[Api.datasetio],
+        deps[Api.datasets],
+    )
+    return impl

llama_stack/providers/inline/post_training/torchtune/common/__init__.py

@@ -0,0 +1,5 @@
+# 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/providers/inline/post_training/torchtune/common/checkpointer.py

@@ -0,0 +1,240 @@
+# 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 json
+import os
+import shutil
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import save_file
+from torchtune import training
+from torchtune.models import convert_weights
+from torchtune.training.checkpointing._utils import (
+    ADAPTER_CONFIG_FNAME,
+    ADAPTER_MODEL_FNAME,
+    REPO_ID_FNAME,
+    SUFFIXES_TO_NOT_COPY,
+    ModelType,
+    copy_files,
+    safe_torch_load,
+)
+from torchtune.utils._logging import get_logger
+
+logger = get_logger("DEBUG")
+
+
+class TorchtuneCheckpointer:
+    def __init__(
+        self,
+        model_id: str,
+        training_algorithm: str,
+        checkpoint_dir: str,
+        checkpoint_files: list[str],
+        output_dir: str,
+        model_type: str,
+    ):
+        # Fail fast if ``checkpoint_files`` is invalid
+        # TODO: support loading more than one file
+        if len(checkpoint_files) != 1:
+            raise ValueError(
+                "Currently we only support reading from a single torchtune checkpoint file. "
+                f"Got {len(checkpoint_files)} files instead."
+            )
+        self._checkpoint_file = checkpoint_files[0]
+        self._model_id = model_id
+        self._training_algorithm = training_algorithm
+        self._checkpoint_dir = Path(checkpoint_dir)
+        self._model_type = ModelType[model_type]
+        self._output_dir = output_dir
+        # get ckpt paths
+        self._checkpoint_path = Path.joinpath(self._checkpoint_dir, self._checkpoint_file)
+
+    def load_checkpoint(self) -> dict[str, Any]:
+        """
+        Load Meta checkpoint from file. Currently only loading from a single file is supported.
+        """
+        state_dict: dict[str, Any] = {}
+        model_state_dict = safe_torch_load(self._checkpoint_path)
+        if self._model_type == ModelType.LLAMA3_VISION:
+            from torchtune.models.llama3_2_vision._convert_weights import (
+                llama3_vision_meta_to_tune,
+            )
+
+            state_dict[training.MODEL_KEY] = llama3_vision_meta_to_tune(model_state_dict)
+        else:
+            state_dict[training.MODEL_KEY] = convert_weights.meta_to_tune(model_state_dict)
+
+        # llama3_2 has tied weights, so we need to remove the output.weight key
+        if self._model_type == ModelType.LLAMA3_2:
+            logger.info(
+                "Identified model_type = Llama3_2. Ignoring output.weight in"
+                " checkpoint in favor of the tok_embedding.weight"
+                " tied weights."
+            )
+            state_dict[training.MODEL_KEY].pop("output.weight")
+
+        return state_dict
+
+    def save_checkpoint(
+        self,
+        state_dict: dict[str, Any],
+        epoch: int,
+        adapter_only: bool = False,
+        checkpoint_format: str | None = None,
+    ) -> str:
+        model_file_path = Path(self._output_dir) / f"{self._model_id}-{self._training_algorithm}-{epoch}"
+        if checkpoint_format == "meta" or checkpoint_format is None:
+            self._save_meta_format_checkpoint(model_file_path, state_dict, adapter_only)
+        elif checkpoint_format == "huggingface":
+            # Note: for saving hugging face format checkpoints, we only suppport saving adapter weights now
+            self._save_hf_format_checkpoint(model_file_path, state_dict)
+        else:
+            raise ValueError(f"Unsupported checkpoint format: {format}")
+        return str(model_file_path)
+
+    def _save_meta_format_checkpoint(
+        self,
+        model_file_path: Path,
+        state_dict: dict[str, Any],
+        adapter_only: bool = False,
+    ) -> None:
+        model_file_path.mkdir(parents=True, exist_ok=True)
+
+        # copy the related files for inference
+        source_path = Path.joinpath(self._checkpoint_dir, "params.json")
+        if source_path.exists():
+            shutil.copy(
+                source_path,
+                Path.joinpath(model_file_path, "params.json"),
+            )
+        source_path = Path.joinpath(self._checkpoint_dir, "tokenizer.model")
+        if source_path.exists():
+            shutil.copy(
+                source_path,
+                Path.joinpath(model_file_path, "tokenizer.model"),
+            )
+        source_path = Path.joinpath(self._checkpoint_dir, "orig_params.json")
+        if source_path.exists():
+            shutil.copy(
+                source_path,
+                Path.joinpath(model_file_path, "orig_params.json"),
+            )
+
+        if not adapter_only:
+            model_state_dict = state_dict[training.MODEL_KEY]
+            if self._model_type == ModelType.LLAMA3_VISION:
+                from torchtune.models.llama3_2_vision._convert_weights import (
+                    llama3_vision_tune_to_meta,
+                )
+
+                state_dict[training.MODEL_KEY] = llama3_vision_tune_to_meta(model_state_dict)
+            else:
+                # llama3_2 has tied weights, so we need to add the output.weight key
+                if self._model_type == ModelType.LLAMA3_2 and "output.weight" not in model_state_dict:
+                    model_state_dict["output.weight"] = model_state_dict["tok_embeddings.weight"]
+
+                state_dict[training.MODEL_KEY] = convert_weights.tune_to_meta(model_state_dict)
+
+            model_file_name = Path.joinpath(model_file_path, "consolidated.00.pth")
+
+            torch.save(state_dict[training.MODEL_KEY], model_file_name)
+            logger.info(
+                "Model checkpoint of size "
+                f"{os.path.getsize(model_file_name) / 1000**3:.2f} GB "
+                f"saved to {model_file_name}"
+            )
+
+        if training.ADAPTER_KEY in state_dict:
+            adapter_file_path = model_file_path / "adapter"
+            adapter_file_path.mkdir(parents=True, exist_ok=True)
+            adapter_file_name = Path.joinpath(adapter_file_path, "adapter.pth")
+            torch.save(state_dict[training.ADAPTER_KEY], adapter_file_name)
+            logger.info(
+                "Adapter checkpoint of size "
+                f"{os.path.getsize(adapter_file_name) / 1000**3:.2f} GB "
+                f"saved to {adapter_file_name}"
+            )
+
+        elif adapter_only:
+            raise ValueError(
+                "Adapter checkpoint not found in state_dict. Please ensure that the state_dict contains adapter weights."
+            )
+
+    def _save_hf_format_checkpoint(
+        self,
+        model_file_path: Path,
+        state_dict: dict[str, Any],
+    ) -> None:
+        # the config.json file contains model params needed for state dict conversion
+        config = json.loads(Path.joinpath(self._checkpoint_dir.parent, "config.json").read_text())
+
+        # repo_id is necessary for when saving an adapter config, so its compatible with HF.
+        # This json file is produced and saved in the download step.
+        # contents are {"repo_id": "some_model/some_model_version"}
+        repo_id_path = Path.joinpath(self._checkpoint_dir.parent, REPO_ID_FNAME).with_suffix(".json")
+        self.repo_id = None
+        if repo_id_path.exists():
+            with open(repo_id_path) as json_file:
+                data = json.load(json_file)
+                self.repo_id = data.get("repo_id")
+
+        if training.ADAPTER_KEY in state_dict:
+            # TODO: saving it "as is" is a requirement because, if we only save with
+            # convert_weights.tune_to_peft_adapter_weights, we do NOT have a fn
+            # convert_weights.peft_to_tune. The .pt format is not needed, but
+            # it is an easy way to distinguish the adapters. Ideally we should save only one.
+            output_path = Path.joinpath(model_file_path, ADAPTER_MODEL_FNAME).with_suffix(".pt")
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            torch.save(state_dict[training.ADAPTER_KEY], output_path)
+            logger.info(
+                f"Adapter checkpoint of size {os.path.getsize(output_path) / 1024**3:.2f} GiB saved to {output_path}"
+            )
+
+            state_dict[training.ADAPTER_KEY] = convert_weights.tune_to_peft_adapter_weights(
+                state_dict[training.ADAPTER_KEY],
+                num_heads=config["num_attention_heads"],
+                num_kv_heads=config["num_key_value_heads"],
+                dim=config["hidden_size"],
+                head_dim=config.get("head_dim", None),
+            )
+            output_path = Path.joinpath(model_file_path, "adapter", ADAPTER_MODEL_FNAME)
+            output_path.parent.mkdir(parents=True, exist_ok=True)
+            output_path = output_path.with_suffix(".safetensors")
+            save_file(
+                state_dict[training.ADAPTER_KEY],
+                output_path,
+                metadata={"format": "pt"},
+            )
+            logger.info(
+                f"Adapter checkpoint of size {os.path.getsize(output_path) / 1024**3:.2f} GiB saved to {output_path}"
+            )
+        else:
+            raise ValueError(
+                "Adapter checkpoint not found in state_dict. Please ensure that the state_dict contains adapter weights."
+            )
+
+        if training.ADAPTER_CONFIG in state_dict:
+            state_dict[training.ADAPTER_CONFIG] = convert_weights.tune_to_peft_adapter_config(
+                adapter_config=state_dict[training.ADAPTER_CONFIG],
+                base_model_name_or_path=self.repo_id,
+            )
+
+            output_path = Path.joinpath(model_file_path, "adapter", ADAPTER_CONFIG_FNAME).with_suffix(".json")
+            with open(output_path, "w") as f:
+                json.dump(state_dict[training.ADAPTER_CONFIG], f)
+            logger.info(
+                f"Adapter checkpoint of size {os.path.getsize(output_path) / 1024**3:.2f} GiB saved to {output_path}"
+            )
+
+        # Save all files in ckpt_dir, except model weights and mapping, to output_dir/epoch_{epoch}
+        # So its easy to run inference with the model using this epoch's checkpoint
+        copy_files(
+            self._checkpoint_dir.parent,
+            model_file_path,
+            ignore_suffixes=SUFFIXES_TO_NOT_COPY,
+        )

llama_stack/providers/inline/post_training/torchtune/common/utils.py

@@ -0,0 +1,99 @@
+# 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, IAny, nc. 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 collections.abc import Callable
+
+import torch
+from pydantic import BaseModel
+from torchtune.data._messages import InputOutputToMessages, ShareGPTToMessages
+from torchtune.models.llama3 import llama3_tokenizer
+from torchtune.models.llama3._tokenizer import Llama3Tokenizer
+from torchtune.models.llama3_1 import lora_llama3_1_8b
+from torchtune.models.llama3_2 import lora_llama3_2_3b
+from torchtune.modules.transforms import Transform
+
+from llama_stack.apis.post_training import DatasetFormat
+from llama_stack.models.llama.sku_list import resolve_model
+from llama_stack.models.llama.sku_types import Model
+
+BuildLoraModelCallable = Callable[..., torch.nn.Module]
+BuildTokenizerCallable = Callable[..., Llama3Tokenizer]
+
+
+class ModelConfig(BaseModel):
+    model_definition: BuildLoraModelCallable
+    tokenizer_type: BuildTokenizerCallable
+    checkpoint_type: str
+
+
+MODEL_CONFIGS: dict[str, ModelConfig] = {
+    "Llama3.2-3B-Instruct": ModelConfig(
+        model_definition=lora_llama3_2_3b,
+        tokenizer_type=llama3_tokenizer,
+        checkpoint_type="LLAMA3_2",
+    ),
+    "Llama3.1-8B-Instruct": ModelConfig(
+        model_definition=lora_llama3_1_8b,
+        tokenizer_type=llama3_tokenizer,
+        checkpoint_type="LLAMA3",
+    ),
+}
+
+DATA_FORMATS: dict[str, Transform] = {
+    "instruct": InputOutputToMessages,
+    "dialog": ShareGPTToMessages,
+}
+
+
+def _validate_model_id(model_id: str) -> Model:
+    model = resolve_model(model_id)
+    if model is None or model.core_model_id.value not in MODEL_CONFIGS:
+        raise ValueError(f"Model {model_id} is not supported.")
+    return model
+
+
+async def get_model_definition(
+    model_id: str,
+) -> BuildLoraModelCallable:
+    model = _validate_model_id(model_id)
+    model_config = MODEL_CONFIGS[model.core_model_id.value]
+    if not hasattr(model_config, "model_definition"):
+        raise ValueError(f"Model {model_id} does not have model definition.")
+    return model_config.model_definition
+
+
+async def get_tokenizer_type(
+    model_id: str,
+) -> BuildTokenizerCallable:
+    model = _validate_model_id(model_id)
+    model_config = MODEL_CONFIGS[model.core_model_id.value]
+    if not hasattr(model_config, "tokenizer_type"):
+        raise ValueError(f"Model {model_id} does not have tokenizer_type.")
+    return model_config.tokenizer_type
+
+
+async def get_checkpointer_model_type(
+    model_id: str,
+) -> str:
+    """
+    checkpointer model type is used in checkpointer for some special treatment on some specific model types
+    For example, llama3.2 model tied weights (https://github.com/pytorch/torchtune/blob/main/torchtune/training/checkpointing/_checkpointer.py#L1041)
+    """
+    model = _validate_model_id(model_id)
+    model_config = MODEL_CONFIGS[model.core_model_id.value]
+    if not hasattr(model_config, "checkpoint_type"):
+        raise ValueError(f"Model {model_id} does not have checkpoint_type.")
+    return model_config.checkpoint_type
+
+
+async def get_data_transform(data_format: DatasetFormat) -> Transform:
+    return DATA_FORMATS[data_format.value]

llama_stack/providers/inline/post_training/torchtune/config.py

@@ -0,0 +1,20 @@
+# 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 Any, Literal
+
+from pydantic import BaseModel
+
+
+class TorchtunePostTrainingConfig(BaseModel):
+    torch_seed: int | None = None
+    checkpoint_format: Literal["meta", "huggingface"] | None = "meta"
+
+    @classmethod
+    def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
+        return {
+            "checkpoint_format": "meta",
+        }

llama_stack/providers/inline/post_training/torchtune/datasets/__init__.py

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