temp commit

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

@@ -20,17 +20,46 @@ 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(
+        optimizer_type=OptimizerType.adamw,
+        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: Optional[str] = "1234",
+        model: Optional[str] = " meta-llama/Llama-3.2-3B-Instruct",
+        dataset_id: Optional[str] = "alpaca",
+        validation_dataset_id: Optional[str] = "alpaca",
+        algorithm: Optional[FinetuningAlgorithm] = FinetuningAlgorithm.lora,
+        algorithm_config: Optional[LoraFinetuningConfig] = LoraFinetuningConfig,
+        optimizer_config: Optional[OptimizerConfig] = OptimizerConfig,
+        training_config: Optional[TrainingConfig] = TrainingConfig,
+        hyperparam_search_config: Optional[Dict[str, Any]] = {},
+        logger_config: Optional[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__)