temp commit

llama_stack/providers/inline/post_training/meta_reference/post_training.py

@@ -8,6 +8,7 @@ from llama_stack.providers.inline.post_training.meta_reference.config import (
     MetaReferencePostTrainingConfig,
 )
 from llama_stack.apis.post_training import *  # noqa
+
 from llama_stack.providers.inline.post_training.meta_reference.recipes.lora_finetuning_single_device import (
     LoraFinetuningSingleDevice,
 )
@@ -20,17 +21,45 @@ class MetaReferencePostTrainingImpl:
         self.config = config
         self.datasetio_api = datasetio_api
 
+    LoraFinetuningConfig(
+        lora_attn_modules=["q_proj", "v_proj", "output_proj"],
+        apply_lora_to_mlp=True,
+        apply_lora_to_output=False,
+        rank=8,
+        alpha=16,
+    )
+
+    OptimizerConfig(
+        lr=3e-4,
+        lr_min=3e-5,
+        weight_decay=0.1,
+        num_warmup_steps=100,
+    )
+
+    TrainingConfig(
+        dtype="bf16",
+        n_epochs=1,
+        max_steps_per_epoch=10,
+        gradient_accumulation_steps=1,
+        batch_size=1,
+        shuffle=1,
+        enable_activation_checkpointing=False,
+        memory_efficient_fsdp_wrap=False,
+        fsdp_cpu_offload=False,
+    )
+
     def supervised_fine_tune(
         self,
-        job_uuid: str,
-        model: str,
-        dataset_id: str,
-        validation_dataset_id: str,
-        algorithm: FinetuningAlgorithm,
-        algorithm_config: LoraFinetuningConfig,
-        optimizer_config: OptimizerConfig,
-        training_config: TrainingConfig,
-        logger_config: Dict[str, Any],
+        job_uuid: str = "1234",
+        model: str = " meta-llama/Llama-3.2-3B-Instruct",
+        dataset_id: str = "alpaca",
+        validation_dataset_id: str = "alpaca",
+        algorithm: FinetuningAlgorithm = FinetuningAlgorithm.lora,
+        algorithm_config: LoraFinetuningConfig = LoraFinetuningConfig,
+        optimizer_config: OptimizerConfig = OptimizerConfig,
+        training_config: TrainingConfig = 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(
@@ -54,3 +83,36 @@ class MetaReferencePostTrainingImpl:
             raise NotImplementedError()
 
         return PostTrainingJob(job_uuid=job_uuid)
+
+    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: ...
+
+    def get_training_jobs(self) -> List[PostTrainingJob]: ...
+
+    # sends SSE stream of logs
+    @webmethod(route="/post-training/job/logs")
+    def get_training_job_logstream(self, job_uuid: str) -> PostTrainingJobLogStream: ...
+
+    @webmethod(route="/post-training/job/status")
+    def get_training_job_status(
+        self, job_uuid: str
+    ) -> PostTrainingJobStatusResponse: ...
+
+    @webmethod(route="/post-training/job/cancel")
+    def cancel_training_job(self, job_uuid: str) -> None: ...
+
+    @webmethod(route="/post-training/job/artifacts")
+    def get_training_job_artifacts(
+        self, job_uuid: str
+    ) -> PostTrainingJobArtifactsResponse: ...

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

@@ -38,7 +38,7 @@ from torchtune.modules.peft import (
     set_trainable_params,
     validate_missing_and_unexpected_for_lora,
 )
-from torchtune.training.lr_scheduler import get_cosine_schedule_with_warmup
+from torchtune.training.lr_schedulers import get_cosine_schedule_with_warmup
 
 log = logging.getLogger(__name__)
 

llama_stack/providers/registry/post_training.py

@@ -12,6 +12,7 @@ from llama_stack.distribution.datatypes import *  # noqa: F403
 META_REFERENCE_DEPS = [
     "torch",
     "torchtune",
+    "torchao",
     "numpy",
 ]
 
@@ -24,5 +25,8 @@ def available_providers() -> List[ProviderSpec]:
             pip_packages=META_REFERENCE_DEPS,
             module="llama_stack.providers.inline.post_training.meta_reference",
             config_class="llama_stack.providers.inline.post_training.meta_reference.MetaReferencePostTrainingConfig",
+            api_dependencies=[
+                Api.datasetio,
+            ],
         ),
     ]