feat: Enable DPO training with HuggingFace inline provider (#2825)

What does this PR do?

This PR adds support for Direct Preference Optimization (DPO) training
via the existing HuggingFace inline provider. It introduces a new DPO
training recipe, config schema updates, dataset integration, and
end-to-end testing to support preference-based fine-tuning with TRL.

Test Plan

Added integration test:

tests/integration/post_training/test_post_training.py::TestPostTraining::test_preference_optimize

Ran tests on both CPU and CUDA environments

---------

Co-authored-by: Ubuntu <[email protected]>
Co-authored-by: Ashwin Bharambe <[email protected]>

Author: 3 people

docs/source/providers/post_training/inline_huggingface.md

@@ -24,6 +24,10 @@ HuggingFace-based post-training provider for fine-tuning models using the Huggin
 | `weight_decay` | `<class 'float'>` | No | 0.01 |  |
 | `dataloader_num_workers` | `<class 'int'>` | No | 4 |  |
 | `dataloader_pin_memory` | `<class 'bool'>` | No | True |  |
+| `dpo_beta` | `<class 'float'>` | No | 0.1 |  |
+| `use_reference_model` | `<class 'bool'>` | No | True |  |
+| `dpo_loss_type` | `Literal['sigmoid', 'hinge', 'ipo', 'kto_pair'` | No | sigmoid |  |
+| `dpo_output_dir` | `<class 'str'>` | No | ./checkpoints/dpo |  |
 
 ## Sample Configuration
 

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

@@ -67,6 +67,12 @@ class HuggingFacePostTrainingConfig(BaseModel):
     # Can improve data transfer speed to GPU but uses more memory
     dataloader_pin_memory: bool = True
 
+    # DPO-specific parameters
+    dpo_beta: float = 0.1
+    use_reference_model: bool = True
+    dpo_loss_type: Literal["sigmoid", "hinge", "ipo", "kto_pair"] = "sigmoid"
+    dpo_output_dir: str = "./checkpoints/dpo"
+
     @classmethod
     def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]:
         return {"checkpoint_format": "huggingface", "distributed_backend": None, "device": "cpu"}

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

@@ -25,6 +25,9 @@
 from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device import (
     HFFinetuningSingleDevice,
 )
+from llama_stack.providers.inline.post_training.huggingface.recipes.finetune_single_device_dpo import (
+    HFDPOAlignmentSingleDevice,
+)
 from llama_stack.providers.utils.scheduler import JobArtifact, Scheduler
 from llama_stack.providers.utils.scheduler import JobStatus as SchedulerJobStatus
 from llama_stack.schema_utils import webmethod
@@ -36,6 +39,7 @@ class TrainingArtifactType(Enum):
 
 
 _JOB_TYPE_SUPERVISED_FINE_TUNE = "supervised-fine-tune"
+_JOB_TYPE_DPO_TRAINING = "dpo-training"
 
 
 class HuggingFacePostTrainingImpl:
@@ -119,12 +123,37 @@ async def preference_optimize(
         hyperparam_search_config: dict[str, Any],
         logger_config: dict[str, Any],
     ) -> PostTrainingJob:
-        raise NotImplementedError("DPO alignment is not implemented yet")
+        async def handler(on_log_message_cb, on_status_change_cb, on_artifact_collected_cb):
+            on_log_message_cb("Starting HF DPO alignment")
 
-    async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
-        return ListPostTrainingJobsResponse(
-            data=[PostTrainingJob(job_uuid=job.id) for job in self._scheduler.get_jobs()]
-        )
+            recipe = HFDPOAlignmentSingleDevice(
+                job_uuid=job_uuid,
+                datasetio_api=self.datasetio_api,
+                datasets_api=self.datasets_api,
+            )
+
+            resources_allocated, checkpoints = await recipe.train(
+                model=finetuned_model,
+                output_dir=f"{self.config.dpo_output_dir}/{job_uuid}",
+                job_uuid=job_uuid,
+                dpo_config=algorithm_config,
+                config=training_config,
+                provider_config=self.config,
+            )
+
+            on_artifact_collected_cb(self._resources_stats_to_artifact(resources_allocated))
+            if checkpoints:
+                for checkpoint in checkpoints:
+                    artifact = self._checkpoint_to_artifact(checkpoint)
+                    on_artifact_collected_cb(artifact)
+            else:
+                on_log_message_cb("Warning: No checkpoints were saved during DPO training")
+
+            on_status_change_cb(SchedulerJobStatus.completed)
+            on_log_message_cb("HF DPO alignment completed")
+
+        job_uuid = self._scheduler.schedule(_JOB_TYPE_DPO_TRAINING, job_uuid, handler)
+        return PostTrainingJob(job_uuid=job_uuid)
 
     @staticmethod
     def _get_artifacts_metadata_by_type(job, artifact_type):
@@ -174,3 +203,9 @@ async def cancel_training_job(self, job_uuid: str) -> None:
     async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse | None:
         job = self._scheduler.get_job(job_uuid)
         return PostTrainingJobArtifactsResponse(job_uuid=job_uuid, checkpoints=self._get_checkpoints(job))
+
+    @webmethod(route="/post-training/jobs", method="GET")
+    async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
+        return ListPostTrainingJobsResponse(
+            data=[PostTrainingJob(job_uuid=job.id) for job in self._scheduler.get_jobs()]
+        )