Refactor SFT forward_backward to use TrainBatch

Replace separate input_dict and labels parameters with unified TrainBatch
dataclass, aligning SFT with GRPO's training interface.

- Import TrainBatch from forge.types
- Update forward_backward signature to accept TrainBatch
- Update train_step and evaluate to create TrainBatch from batch dict

Author: Hossein Kavianihamedani

apps/sft/main.py

@@ -26,6 +26,7 @@
 from forge.data.tokenizer import HuggingFaceModelTokenizer
 from forge.data.utils import StopAfterOneEpoch
 from forge.observability import get_or_create_metric_logger, record_metric, Reduce
+from forge.types import TrainBatch
 from forge.util.config import parse
 from monarch.actor import current_rank, current_size, endpoint
 from omegaconf import DictConfig, OmegaConf
@@ -213,16 +214,16 @@ def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader:
 
     def forward_backward(
         self,
-        input_dict: dict[str, torch.Tensor],
-        labels: torch.Tensor,
+        batch: TrainBatch,
         skip_backward: bool = False,
     ) -> torch.Tensor:
         model_parts = self.model_parts
         parallel_dims = self.parallel_dims
 
         # apply context parallelism if cp is enabled
         # ensure CP handles the separate freqs_cis buffer for each pp stage
-        inputs = input_dict["tokens"]
+        inputs = batch.model_inputs["tokens"]
+        labels = batch.loss_inputs["labels"]
         optional_context_parallel_ctx = (
             dist_utils.create_context_parallel_ctx(
                 cp_mesh=parallel_dims.world_mesh["cp"],
@@ -283,7 +284,11 @@ def train_step(self, batch) -> None:
         # ) as grad_acc:
         parallel_dims = self.parallel_dims
         labels = batch.pop("labels")
-        loss = self.forward_backward(batch, labels)
+        train_batch = TrainBatch(
+            model_inputs=batch,
+            loss_inputs={"labels": labels},
+        )
+        loss = self.forward_backward(train_batch)
 
         grad_norm = dist_utils.clip_grad_norm_(
             [p for m in self.model_parts for p in m.parameters()],
@@ -373,7 +378,11 @@ async def evaluate(self) -> None:
 
                     # Process batch
                     labels = batch.pop("labels")
-                    loss = self.forward_backward(batch, labels, skip_backward=True)
+                    train_batch = TrainBatch(
+                        model_inputs=batch,
+                        loss_inputs={"labels": labels},
+                    )
+                    loss = self.forward_backward(train_batch, skip_backward=True)
                     total_loss += loss
                     num_steps += 1