temp commit

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

@@ -0,0 +1,25 @@
+# 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 Dict
+
+from llama_stack.distribution.datatypes import Api, ProviderSpec
+
+from .config import TorchtunePostTrainingConfig
+
+
+async def get_provider_impl(
+    config: TorchtunePostTrainingConfig,
+    deps: Dict[Api, ProviderSpec],
+):
+    from .post_training import TorchtunePostTrainingImpl
+
+    impl = TorchtunePostTrainingImpl(
+        config,
+        deps[Api.datasetio],
+    )
+    # await impl.initialize()
+    return impl

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 Optional
+
+from pydantic import BaseModel, Field
+
+
+class TorchtunePostTrainingConfig(BaseModel):
+    model: str = Field(
+        default="Llama3.2-3B-Instruct",
+        description="Model descriptor from `llama model list`",
+    )
+    torch_seed: Optional[int] = None
+    # By default, the implementation will look at ~/.llama/checkpoints/<model> but you
+    # can override by specifying the directory explicitly
+    checkpoint_dir: Optional[str] = None

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

@@ -0,0 +1,66 @@
+# 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 BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any, Dict, List, Mapping
+
+import numpy as np
+
+from torch.utils.data import Dataset
+from torchtune.data._common import CROSS_ENTROPY_IGNORE_IDX
+from torchtune.data._messages import validate_messages
+from torchtune.modules.transforms import Transform
+
+
+class SFTDataset(Dataset):
+    def __init__(
+        self,
+        rows: List[Dict[str, Any]],
+        message_transform: Transform,
+        model_transform: Transform,
+    ) -> None:
+        self._rows = rows
+        self._message_transform = message_transform
+        self._model_transform = model_transform
+
+    def __len__(self):
+        return len(self._rows)
+
+    def __getitem__(self, index: int) -> Dict[str, Any]:
+        sample = self._rows[index]
+        return self._prepare_sample(sample)
+
+    def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, Any]:
+        transformed_sample = self._message_transform(sample)
+        if "messages" in transformed_sample:
+            validate_messages(transformed_sample["messages"])
+
+        tokenized_dict = self._model_transform(transformed_sample)
+
+        if not ("tokens" in tokenized_dict and "mask" in tokenized_dict):
+            keys_str = ", ".join(tokenized_dict.keys())
+            error_message = (
+                "model_transform returned the following keys: "
+                f"{keys_str}. Must return 'tokens' and 'mask' as keys."
+            )
+            raise ValueError(error_message)
+
+        # Wherever mask == True, set to CROSS_ENTROPY_IGNORE_IDX. Otherwise keep as tokens
+        tokenized_dict["labels"] = list(
+            np.where(
+                tokenized_dict["mask"],
+                CROSS_ENTROPY_IGNORE_IDX,
+                tokenized_dict["tokens"],
+            )
+        )
+        assert len(tokenized_dict["tokens"]) == len(tokenized_dict["labels"])
+
+        return tokenized_dict

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

@@ -0,0 +1,101 @@
+# 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 llama_stack.apis.datasetio import DatasetIO
+from llama_stack.providers.inline.post_training.torchtune.config import (
+    TorchtunePostTrainingConfig,
+)
+from llama_stack.apis.post_training import *  # noqa
+from llama_stack.providers.inline.post_training.torchtune.recipes.lora_finetuning_single_device import (
+    LoraFinetuningSingleDevice,
+)
+
+
+class PostTrainingSFTRequest(BaseModel):
+    job_uuid: str
+    model: str
+    algorithm: FinetuningAlgorithm
+    algorithm_config: Optional[Union[LoraFinetuningConfig, QATFinetuningConfig]] = None
+    training_config: TrainingConfig
+
+    # TODO: define these
+    hyperparam_search_config: Dict[str, Any]
+    logger_config: Dict[str, Any]
+
+
+class TorchtunePostTrainingImpl:
+    def __init__(
+        self, config: TorchtunePostTrainingConfig, datasetio_api: DatasetIO
+    ) -> None:
+        self.config = config
+        self.datasetio_api = datasetio_api
+
+    async def supervised_fine_tune(
+        self,
+        job_uuid: str,
+        model: str,
+        algorithm: FinetuningAlgorithm,
+        algorithm_config: LoraFinetuningConfig,
+        training_config: TrainingConfig,
+        hyperparam_search_config: Dict[str, Any],
+        logger_config: Dict[str, Any],
+    ) -> PostTrainingJob:
+
+        # wrapper request to make it easier to pass around (internal only, not exposed to API)
+        request = PostTrainingSFTRequest(
+            job_uuid=job_uuid,
+            model=model,
+            algorithm=algorithm,
+            algorithm_config=algorithm_config,
+            training_config=training_config,
+            hyperparam_search_config=hyperparam_search_config,
+            logger_config=logger_config,
+        )
+        if request.algorithm == FinetuningAlgorithm.lora:
+            recipe = LoraFinetuningSingleDevice(
+                self.config, request, self.datasetio_api
+            )
+            await recipe.setup(self.config)
+            await recipe.train()
+        else:
+            raise NotImplementedError()
+
+        return PostTrainingJob(job_uuid=job_uuid)
+
+    async def preference_optimize(
+        self,
+        job_uuid: str,
+        finetuned_model: URL,
+        dataset_id: str,
+        validation_dataset_id: str,
+        algorithm: RLHFAlgorithm,
+        algorithm_config: DPOAlignmentConfig,
+        optimizer_config: OptimizerConfig,
+        training_config: TrainingConfig,
+        hyperparam_search_config: Dict[str, Any],
+        logger_config: Dict[str, Any],
+    ) -> PostTrainingJob: ...
+
+    # TODO @markchen1015 impelment below APIs
+    async def get_training_jobs(self) -> List[PostTrainingJob]: ...
+
+    # sends SSE stream of logs
+    @webmethod(route="/post-training/job/logs")
+    async def get_training_job_logstream(
+        self, job_uuid: str
+    ) -> PostTrainingJobLogStream: ...
+
+    @webmethod(route="/post-training/job/status")
+    async def get_training_job_status(
+        self, job_uuid: str
+    ) -> PostTrainingJobStatusResponse: ...
+
+    @webmethod(route="/post-training/job/cancel")
+    async def cancel_training_job(self, job_uuid: str) -> None: ...
+
+    @webmethod(route="/post-training/job/artifacts")
+    async def get_training_job_artifacts(
+        self, job_uuid: str
+    ) -> PostTrainingJobArtifactsResponse: ...

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -0,0 +1,500 @@
+# 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 logging
+import os
+import time
+from functools import partial
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from llama_models.sku_list import resolve_model
+from llama_stack.apis.datasetio import DatasetIO
+from torch import nn
+from torchtune import utils as torchtune_utils
+from torchtune.training.metric_logging import DiskLogger
+from llama_stack.apis.post_training import *  # noqa
+from llama_stack.distribution.utils.model_utils import model_local_dir
+
+from llama_stack.providers.inline.post_training.torchtune import utils
+from llama_stack.providers.inline.post_training.torchtune.config import (
+    MetaReferencePostTrainingConfig,
+)
+from llama_stack.providers.inline.post_training.torchtune.datasets.sft import SFTDataset
+from llama_stack.providers.inline.post_training.torchtune.post_training import (
+    PostTrainingSFTRequest,
+)
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader, DistributedSampler
+from torchtune import modules, training
+from torchtune.data import AlpacaToMessages, padded_collate_sft
+
+from torchtune.modules.loss import CEWithChunkedOutputLoss
+from torchtune.modules.peft import (
+    get_adapter_params,
+    get_adapter_state_dict,
+    get_lora_module_names,
+    get_merged_lora_ckpt,
+    load_dora_magnitudes,
+    set_trainable_params,
+    validate_missing_and_unexpected_for_lora,
+)
+from torchtune.training.lr_schedulers import get_cosine_schedule_with_warmup
+
+log = logging.getLogger(__name__)
+
+from torchtune.models.llama3._tokenizer import Llama3Tokenizer
+
+
+class LoraFinetuningSingleDevice:
+    # This recipe only supports GPU training
+
+    # This recipe doesn't include several training efficiency setting within origin torchtune repo, including
+    # - compile
+    # - activation offloading
+
+    # Resume from checkpoint hasn't been supported yet
+    # Validation hasn't been supported yet
+
+    # TODO @markchen1015 figure out the logging for this training recipe
+    # and make it work with telemetry
+    def __init__(
+        self,
+        config: MetaReferencePostTrainingConfig,
+        request: PostTrainingSFTRequest,
+        datasetio_api: DatasetIO,
+    ) -> None:
+        # Assume the training only happens on GPU
+        self.config = config
+        self.request = request
+        self._device = torchtune_utils.get_device(device="cuda")
+        self._dtype = training.get_dtype(
+            request.training_config.dtype, device=self._device
+        )
+        self.model_id = config.model
+
+        def model_checkpoint_dir(model) -> str:
+            checkpoint_dir = Path(model_local_dir(model.descriptor()))
+
+            paths = [
+                Path(checkpoint_dir / f"consolidated.{ext}")
+                for ext in ["pth", "00.pth"]
+            ]
+            if not any(p.exists() for p in paths):
+                checkpoint_dir = checkpoint_dir / "original"
+
+            assert checkpoint_dir.exists(), (
+                f"Could not find checkpoints in: {model_local_dir(model.descriptor())}. "
+                f"Please download model using `llama download --model-id {model.descriptor()}`"
+            )
+            return str(checkpoint_dir)
+
+        if config.checkpoint_dir and config.checkpoint_dir != "null":
+            self.checkpoint_dir = config.checkpoint_dir
+        else:
+            model = resolve_model(self.model_id)
+            self.checkpoint_dir = model_checkpoint_dir(model)
+
+        # TODO @markchen1015 make it work with get_training_job_artifacts
+        self._output_dir = self.checkpoint_dir + "/posting_training/"
+
+        self.seed = training.set_seed(seed=config.torch_seed or 42)
+        self.epochs_run = 0
+        self.total_epochs = request.training_config.n_epochs
+        self._shuffle = request.training_config.data_config.shuffle
+        self._batch_size = request.training_config.data_config.batch_size
+
+        # this is important for debugging purpose
+        self.max_steps_per_epoch = request.training_config.max_steps_per_epoch
+        self.global_step = 0
+
+        self._gradient_accumulation_steps = (
+            request.training_config.gradient_accumulation_steps
+        )
+
+        self._clip_grad_norm = 1.0
+        self._enable_activation_checkpointing = (
+            (request.training_config.efficiency_config.enable_activation_checkpointing)
+            if request.training_config.efficiency_config
+            else False
+        )
+        self._enable_activation_offloading = (
+            (request.training_config.efficiency_config.enable_activation_offloading)
+            if request.training_config.efficiency_config
+            else False
+        )
+
+        self.datasetio_api = datasetio_api
+
+    async def load_checkpoint(self):
+        def get_checkpoint_files(checkpoint_dir: str) -> List[str]:
+            try:
+                # List all files in the given directory
+                files = os.listdir(checkpoint_dir)
+                # Filter files that end with .pth
+                pth_files = [file for file in files if file.endswith(".pth")]
+                return pth_files
+            except FileNotFoundError:
+                return [f"Error: The directory '{checkpoint_dir}' does not exist."]
+
+        self._checkpointer = training.FullModelMetaCheckpointer(
+            checkpoint_dir=self.checkpoint_dir,
+            checkpoint_files=get_checkpoint_files(self.checkpoint_dir),
+            output_dir=self._output_dir,
+            model_type=utils.get_checkpointer_model_type(self.model_id),
+        )
+        checkpoint_dict = self._checkpointer.load_checkpoint()
+        return checkpoint_dict
+
+    async def setup(self, config: MetaReferencePostTrainingConfig) -> None:
+        # temporily log to local disk, will figure out how to interop with telemetry
+        self._metric_logger = DiskLogger(log_dir=self._output_dir)
+
+        checkpoint_dict = await self.load_checkpoint()
+
+        self._model = await self._setup_model(
+            enable_activation_checkpointing=self._enable_activation_checkpointing,
+            enable_activation_offloading=self._enable_activation_offloading,
+            base_model_state_dict=checkpoint_dict[training.MODEL_KEY],
+            lora_weights_state_dict=None,
+        )
+        log.info(f"Model is initialized with precision {self._dtype}.")
+
+        self._tokenizer = await self._setup_tokenizer()
+        log.info("Tokenizer is initialized from file.")
+
+        self._optimizer = await self._setup_optimizer(
+            optimizer_config=self.request.training_config.optimizer_config
+        )
+        log.info("Optimizer is initialized.")
+
+        self._loss_fn = CEWithChunkedOutputLoss()
+        self._model.set_num_output_chunks(self._loss_fn.num_output_chunks)
+        log.info("Loss is initialized.")
+
+        self._sampler, self._dataloader = await self._setup_data(
+            tokenizer=self._tokenizer,
+            shuffle=self._shuffle,
+            batch_size=self._batch_size,
+        )
+        log.info("Dataset and Sampler are initialized.")
+
+        # Number of training steps in each epoch depends on the number of batches produced
+        # by the dataloader and the max_steps_per_epoch param set by the user and is used
+        # for logging and tracking training state. This should be computed after the dataloader
+        # has been setup
+        self._steps_per_epoch = (
+            len(self._dataloader) // self._gradient_accumulation_steps
+        )
+        if (
+            self.max_steps_per_epoch is not None
+            and self.max_steps_per_epoch < self._steps_per_epoch
+        ):
+            self._steps_per_epoch = self.max_steps_per_epoch
+            self.global_step = self.epochs_run * self._steps_per_epoch
+
+        # Learning rate scheduler can only be set up after number of steps
+        # has been computed
+        self._lr_scheduler = await self._setup_lr_scheduler(
+            num_warmup_steps=self.request.training_config.optimizer_config.num_warmup_steps,
+            num_training_steps=self.total_epochs * self._steps_per_epoch,
+            last_epoch=self.global_step - 1,
+        )
+        log.info("Learning rate scheduler is initialized.")
+
+        # Used to ignore labels for loss computation
+        self.ignore_labels_cache = torch.full(
+            (self._batch_size, 1), self._loss_fn.ignore_index, device=self._device
+        )
+
+    async def _setup_model(
+        self,
+        enable_activation_checkpointing: bool,
+        enable_activation_offloading: bool,
+        base_model_state_dict: Dict[str, Any],
+        lora_weights_state_dict: Optional[Dict[str, Any]] = None,
+    ) -> nn.Module:
+        self._lora_rank = self.request.algorithm_config.rank
+        self._lora_alpha = self.request.algorithm_config.alpha
+        self._lora_attn_modules = list(self.request.algorithm_config.lora_attn_modules)
+        self._apply_lora_to_mlp = self.request.algorithm_config.apply_lora_to_mlp
+        self._apply_lora_to_output = self.request.algorithm_config.apply_lora_to_output
+        self._use_dora = self.request.algorithm_config.use_dora or False
+
+        with training.set_default_dtype(self._dtype), self._device:
+            model_type = utils.get_model_type(self.model_id)
+            model = model_type(
+                lora_attn_modules=self._lora_attn_modules,
+                apply_lora_to_mlp=self._apply_lora_to_mlp,
+                apply_lora_to_output=self._apply_lora_to_output,
+                lora_rank=self._lora_rank,
+                lora_alpha=self._lora_alpha,
+                quantize_base=False,
+                use_dora=self._use_dora,
+            )
+
+        self.adapter_params = get_adapter_params(model)
+        self._is_dora = any(["magnitude" in k for k in self.adapter_params.keys()])
+
+        set_trainable_params(model, self.adapter_params)
+
+        if enable_activation_checkpointing:
+            training.set_activation_checkpointing(
+                model, auto_wrap_policy={modules.TransformerSelfAttentionLayer}
+            )
+
+        base_missing, base_unexpected = model.load_state_dict(
+            base_model_state_dict, strict=False
+        )
+
+        # This is for any adapters that need to be initialized after base weights
+        # have been loaded (e.g. DoRA).
+        if self._is_dora:
+            for m in model.modules():
+                if hasattr(m, "initialize_dora_magnitude"):
+                    m.initialize_dora_magnitude()
+            load_dora_magnitudes(model)
+        if lora_weights_state_dict:
+            lora_missing, lora_unexpected = model.load_state_dict(
+                lora_weights_state_dict, strict=False
+            )
+        else:
+            lora_missing, lora_unexpected = None, None
+        validate_missing_and_unexpected_for_lora(
+            lora_attn_modules=self._lora_attn_modules,
+            apply_lora_to_mlp=self._apply_lora_to_mlp,
+            apply_lora_to_output=self._apply_lora_to_output,
+            base_missing=base_missing,
+            base_unexpected=base_unexpected,
+            lora_missing=lora_missing,
+            lora_unexpected=lora_unexpected,
+        )
+
+        # Validate model adapter params were loaded in with the expected dtype
+        training.validate_expected_param_dtype(
+            self.adapter_params.items(), dtype=self._dtype
+        )
+
+        # activation offloading
+        self.activations_handling_ctx = training.get_act_offloading_ctx_manager(
+            model, enable_activation_offloading
+        )
+
+        memory_stats = training.get_memory_stats(device=self._device)
+        training.log_memory_stats(memory_stats)
+
+        return model
+
+    async def _setup_tokenizer(
+        self,
+    ) -> Llama3Tokenizer:
+        tokenizer_path = self.checkpoint_dir + "/tokenizer.model"
+        tokenizer_type = utils.get_tokenizer_type(self.model_id)
+        return tokenizer_type(path=tokenizer_path)
+
+    async def _setup_optimizer(self, optimizer_config: OptimizerConfig) -> Optimizer:
+        optimizer = torch.optim.AdamW(
+            params=self._model.parameters(),
+            lr=optimizer_config.lr,
+            betas=(0.9, 0.95),
+            eps=1e-8,
+            weight_decay=0.1,
+        )
+        return optimizer
+
+    async def _setup_data(
+        self, tokenizer: Llama3Tokenizer, shuffle: bool, batch_size: int
+    ) -> Tuple[DistributedSampler, DataLoader]:
+        async def fetch_rows():
+            return await self.datasetio_api.get_rows_paginated(
+                dataset_id=self.request.training_config.data_config.dataset_id,
+                rows_in_page=-1,
+            )
+
+        all_rows = await fetch_rows()
+        rows = all_rows.rows
+
+        # Curretly only support instruct dataset
+        # TODO @markchen1015 make the message_transform swappable and support more dataset types
+        ds = SFTDataset(
+            rows,
+            message_transform=AlpacaToMessages(train_on_input=False),
+            model_transform=tokenizer,
+        )
+
+        sampler = DistributedSampler(
+            ds,
+            num_replicas=1,
+            rank=0,
+            shuffle=shuffle,
+            seed=0,
+        )
+        dataloader = DataLoader(
+            dataset=ds,
+            sampler=sampler,
+            batch_size=batch_size,
+            # dropping last avoids shape issues with compile + flex attention
+            drop_last=True,
+            collate_fn=(
+                partial(
+                    padded_collate_sft,
+                    padding_idx=self._tokenizer.pad_id,
+                    ignore_idx=self._loss_fn.ignore_index,
+                )
+            ),
+        )
+
+        return sampler, dataloader
+
+    async def _setup_lr_scheduler(
+        self,
+        num_warmup_steps: int,
+        num_training_steps: int,
+        last_epoch: int,
+    ) -> Optimizer:
+        lr_scheduler = get_cosine_schedule_with_warmup(
+            self._optimizer,
+            num_warmup_steps=num_warmup_steps,
+            num_training_steps=num_training_steps,
+            last_epoch=last_epoch,
+        )
+        return lr_scheduler
+
+    async def save_checkpoint(self, epoch: int) -> None:
+        ckpt_dict = {}
+
+        adapter_state_dict = get_adapter_state_dict(self._model.state_dict())
+        ckpt_dict.update({training.ADAPTER_KEY: adapter_state_dict})
+
+        # Construct the full state dict with LoRA weights merged into base LLM weights
+        # Move to CPU to avoid a copy on GPU
+        state_dict = {k: v.cpu() for k, v in self._model.state_dict().items()}
+
+        merged_state_dict = get_merged_lora_ckpt(
+            state_dict,
+            rank=self._lora_rank,
+            alpha=self._lora_alpha,
+        )
+
+        ckpt_dict.update({training.MODEL_KEY: merged_state_dict})
+
+        adapter_config = {
+            "r": self._lora_rank,
+            "lora_alpha": self._lora_alpha,
+            "target_modules": get_lora_module_names(
+                self._lora_attn_modules,
+                self._apply_lora_to_mlp,
+                self._apply_lora_to_output,
+            ),
+            "peft_type": "LORA",
+        }
+        ckpt_dict.update({training.ADAPTER_CONFIG: adapter_config})
+
+        self._checkpointer.save_checkpoint(
+            ckpt_dict,
+            epoch=epoch,
+        )
+
+    async def _loss_step(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
+        # Shape [b, s], needed for the loss not the model
+        labels = batch.pop("labels")
+        # run model
+        with self.activations_handling_ctx:
+            logits = self._model(**batch)
+
+        # Shift labels to compute loss
+        # equivalent to doing labels[..., 1:] and logits[..., :-1, :]
+        # But this way we dont need to slice the logits. We just add an ignore index to labels.
+        labels = torch.hstack(
+            (labels[..., 1:], self.ignore_labels_cache[: labels.shape[0]])
+        )
+        if not isinstance(logits, list):
+            labels = labels.reshape(-1)
+            logits = logits.reshape(-1, logits.size(-1))
+
+        loss = self._loss_fn(logits, labels)
+
+        # free logits otherwise it peaks backward memory
+        del logits
+
+        return loss
+
+    async def train(self) -> None:
+        """
+        The core training loop.
+        """
+        # Initialize tokens count and running loss (for grad accumulation)
+        # t0 = time.perf_counter()
+        t0 = time.perf_counter()
+        running_loss = 0
+        num_tokens = 0
+
+        # self.epochs_run should be non-zero when we're resuming from a checkpoint
+        for curr_epoch in range(self.epochs_run, self.total_epochs):
+            # Update the sampler to ensure data is correctly shuffled across epochs
+            # in case shuffle is True
+            self._sampler.set_epoch(curr_epoch)
+
+            for idx, batch in enumerate(self._dataloader):
+                if (
+                    self.max_steps_per_epoch is not None
+                    and (idx // self._gradient_accumulation_steps)
+                    == self.max_steps_per_epoch
+                ):
+                    break
+
+                torchtune_utils.batch_to_device(batch, self._device)
+
+                # Calculate the number of unmasked tokens in the current batch
+                # and increment the total number of tokens seen in the step
+                current_num_tokens = (
+                    batch["labels"] != self._loss_fn.ignore_index
+                ).sum()
+                num_tokens += current_num_tokens
+
+                # Loss is normalized by default so we multiply by the number of tokens
+                # This way we can normalize by the total number of tokens if we're accumulating gradients
+                current_loss = await self._loss_step(batch) * current_num_tokens
+                running_loss += current_loss
+                current_loss.backward()
+
+                # Step with optimizer
+                if (idx + 1) % self._gradient_accumulation_steps == 0:
+                    training.scale_grads(self._model, 1 / num_tokens)
+                    grad_norm = torch.nn.utils.clip_grad_norm_(
+                        self._model.parameters(),
+                        max_norm=float(self._clip_grad_norm),
+                    )
+                    self._optimizer.step()
+                    self._optimizer.zero_grad(set_to_none=True)
+                    self._lr_scheduler.step()
+                    # Update the number of steps when the weights are updated
+                    self.global_step += 1
+
+                    loss_to_log = running_loss.item() / num_tokens
+                    time_per_step = time.perf_counter() - t0
+                    log_dict = {
+                        "loss": loss_to_log,
+                        "lr": self._optimizer.param_groups[0]["lr"],
+                        "tokens_per_second_per_gpu": num_tokens / time_per_step,
+                    }
+                    log_dict.update(training.get_memory_stats(device=self._device))
+                    if self._clip_grad_norm is not None:
+                        log_dict.update({"grad_norm": grad_norm})
+                    self._metric_logger.log_dict(
+                        log_dict,
+                        step=self.global_step,
+                    )
+
+                    # Reset running stats for the next step
+                    running_loss = 0
+                    num_tokens = 0
+                    t0 = time.perf_counter()
+
+            self.epochs_run += 1
+            log.info("Starting checkpoint save...")
+            await self.save_checkpoint(epoch=curr_epoch)

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

@@ -0,0 +1,58 @@
+# 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 typing import Any, Callable, Dict
+
+import torch
+from llama_models.sku_list import resolve_model
+
+from torchtune.models.llama3 import llama3_tokenizer, lora_llama3_8b
+from torchtune.models.llama3._tokenizer import Llama3Tokenizer
+from torchtune.models.llama3_2 import lora_llama3_2_3b
+
+LORA_MODEL_TYPES: Dict[str, Any] = {
+    "Llama3.2-3B-Instruct": lora_llama3_2_3b,
+    "Llama-3-8B-Instruct": lora_llama3_8b,
+}
+
+TOKENIZER_TYPES: Dict[str, Any] = {
+    "Llama3.2-3B-Instruct": llama3_tokenizer,
+    "Llama-3-8B-Instruct": llama3_tokenizer,
+}
+
+CHECKPOINT_MODEL_TYPES: Dict[str, str] = {
+    "Llama3.2-3B-Instruct": "LLAMA3_2",
+}
+
+BuildLoraModelCallable = Callable[..., torch.nn.Module]
+BuildTokenizerCallable = Callable[..., Llama3Tokenizer]
+
+
+def get_model_type(
+    model_id: str,
+) -> BuildLoraModelCallable:
+    model = resolve_model(model_id)
+    return LORA_MODEL_TYPES[model.core_model_id.value]
+
+
+def get_tokenizer_type(
+    model_id: str,
+) -> BuildTokenizerCallable:
+    model = resolve_model(model_id)
+    return TOKENIZER_TYPES[model.core_model_id.value]
+
+
+def get_checkpointer_model_type(
+    model_id: str,
+) -> str:
+    model = resolve_model(model_id)
+    return CHECKPOINT_MODEL_TYPES[model.core_model_id.value]