Refactor SFT forward_backward to use TrainBatch

apps/sft/main.py

@@ -21,11 +21,12 @@
 import torch
 import torchtitan.experiments.forge.train_spec as forge_train_spec
 from forge.controller import ForgeActor
-from forge.data.collate import collate_padded
+from forge.data.collate import collate_sft
 from forge.data.datasets.sft_dataset import AlpacaToMessages, sft_iterable_dataset
 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
@@ -206,23 +207,23 @@ def setup_data(self, dataset_configs: list[dict]) -> StatefulDataLoader:
         dataloader = StatefulDataLoader(
             dataset=dataset,
             batch_size=self.job_config.training.local_batch_size,
-            collate_fn=collate_padded,
+            collate_fn=collate_sft,
         )
 
         return dataloader
 
     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"],
@@ -274,16 +275,15 @@ def forward_backward(
 
         return loss
 
-    def train_step(self, batch) -> None:
+    def train_step(self, batch: TrainBatch) -> None:
         # TODO
         # with GradientAccumulation(
         #     self.gradient_accumulation_steps,
         #     self.model,
         #     self.data_parallel_size,
         # ) as grad_acc:
         parallel_dims = self.parallel_dims
-        labels = batch.pop("labels")
-        loss = self.forward_backward(batch, labels)
+        loss = self.forward_backward(batch)
 
         grad_norm = dist_utils.clip_grad_norm_(
             [p for m in self.model_parts for p in m.parameters()],
@@ -367,13 +367,15 @@ async def evaluate(self) -> None:
                         break
 
                     # Move tensors to device
-                    for key, value in batch.items():
+                    for key, value in batch.model_inputs.items():
                         if isinstance(value, torch.Tensor):
-                            batch[key] = value.to(self.device)
+                            batch.model_inputs[key] = value.to(self.device)
+                    for key, value in batch.loss_inputs.items():
+                        if isinstance(value, torch.Tensor):
+                            batch.loss_inputs[key] = value.to(self.device)
 
                     # Process batch
-                    labels = batch.pop("labels")
-                    loss = self.forward_backward(batch, labels, skip_backward=True)
+                    loss = self.forward_backward(batch, skip_backward=True)
                     total_loss += loss
                     num_steps += 1
 
@@ -435,16 +437,19 @@ async def train(self) -> None:
         # self.pbar = tqdm(initial=self.current_step, total=self.num_training_steps)
 
         while self.current_step < self.num_training_steps:
-            batch = next(dataloader)
+            batch: TrainBatch = next(dataloader)
 
-            # Pop and record metrics from batch before moving to device
-            self.record_batch_metrics(batch.pop("metrics", []))
+            # Pop and record metrics from batch metadata
+            self.record_batch_metrics(batch.meta.pop("metrics", []))
             record_metric("ForgeSFTRecipe/train/step", self.current_step, Reduce.MEAN)
 
             # Move tensors to the appropriate device
-            for k, v in batch.items():
+            for k, v in batch.model_inputs.items():
+                if isinstance(v, torch.Tensor):
+                    batch.model_inputs[k] = v.to("cuda")  # TODO: hardcoded for now
+            for k, v in batch.loss_inputs.items():
                 if isinstance(v, torch.Tensor):
-                    batch[k] = v.to("cuda")  # TODO: hardcoded for now
+                    batch.loss_inputs[k] = v.to("cuda")
 
             self.train_step(batch)
             # self.profiler.step()

src/forge/data/__init__.py

@@ -3,13 +3,14 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-from .collate import collate_packed, collate_padded
+from .collate import collate_packed, collate_padded, collate_sft
 from .metric_transform import DefaultDatasetMetricTransform, MetricTransform
 from .utils import CROSS_ENTROPY_IGNORE_IDX
 
 __all__ = [
     "collate_packed",
     "collate_padded",
+    "collate_sft",
     "CROSS_ENTROPY_IGNORE_IDX",
     "MetricTransform",
     "DefaultDatasetMetricTransform",

src/forge/data/collate.py

@@ -9,6 +9,50 @@
 import torch
 import torch.nn.functional as F
 from forge.data.utils import CROSS_ENTROPY_IGNORE_IDX
+from forge.types import TrainBatch
+
+
+def collate_sft(batch: list[dict[str, Any]]) -> TrainBatch:
+    """
+    Collate function for SFT that returns a TrainBatch directly.
+
+    This is a wrapper around collate_padded that separates model inputs
+    from loss inputs, following the same pattern as GRPO's collate.
+
+    Args:
+        batch: List of samples, each containing tensor and non-tensor fields.
+            Expected keys: "tokens" (model input), "labels" (loss input),
+            and optionally "metrics" and other metadata.
+
+    Returns:
+        TrainBatch with model_inputs, loss_inputs, and meta fields separated.
+    """
+    collated = collate_padded(batch)
+
+    # Separate model inputs, loss inputs, and metadata
+    model_inputs = {}
+    loss_inputs = {}
+    meta = {}
+
+    for key, value in collated.items():
+        if key == "tokens":
+            model_inputs[key] = value
+        elif key == "labels":
+            loss_inputs[key] = value
+        elif key == "metrics":
+            meta[key] = value
+        else:
+            # Other tensor fields go to model_inputs, non-tensor to meta
+            if isinstance(value, torch.Tensor):
+                model_inputs[key] = value
+            else:
+                meta[key] = value
+
+    return TrainBatch(
+        model_inputs=model_inputs,
+        loss_inputs=loss_inputs,
+        meta=meta,
+    )
 
 
 def collate_padded(batch: list[dict[str, Any]]) -> dict[str, Any]:

src/forge/data/utils.py

@@ -4,12 +4,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
+
 import logging
 from enum import Enum
 from typing import Any, Iterator, Literal, Union
 
 import torch
 import torch.distributed as dist
+from forge.types import TrainBatch
 from torch.nn.attention.flex_attention import BlockMask
 
 logger = logging.getLogger(__name__)
@@ -256,7 +259,7 @@ def __init__(
     def __iter__(self):
         return self
 
-    def __next__(self) -> dict:
+    def __next__(self) -> TrainBatch:
         """Get next batch from current epoch.
 
         Returns:
@@ -302,32 +305,34 @@ def __next__(self) -> dict:
         return current_batch
 
 
-def extract_epoch_from_batch(batch: dict) -> int:
+def extract_epoch_from_batch(batch: TrainBatch) -> int:
     """Extract epoch number from batch metrics. Useful to detect epoch changes during validation.
 
     Assumes batch contains field "metrics" with at least one Metric containing "num_epochs" in its key, as it is done in
     `forge.src.data.datasets.HfIterableDataset`.
 
     Args:
-        batch (dict): Batch dictionary with 'metrics' field
+        batch (TrainBatch): TrainBatch with 'metrics' field in meta
 
     Returns:
         int: Max epoch number from metrics
 
     Raises:
         ValueError: If metrics key is missing or no metric with 'num_epochs' found
     """
-    if "metrics" not in batch:
+    metrics = batch.meta.get("metrics")
+
+    if metrics is None:
         raise ValueError(
-            "Batch missing 'metrics' field. Cannot extract epoch from batch."
+            "Batch missing 'metrics' field in meta. Cannot extract epoch from batch."
         )
 
     # Match metrics where 'num_epochs' appears in the key (handles prefixed keys like 'dataset/name/num_epochs')
-    epochs = [metric.value for metric in batch["metrics"] if "num_epochs" in metric.key]
+    epochs = [metric.value for metric in metrics if "num_epochs" in metric.key]
     if epochs:
         return max(epochs)
 
     raise ValueError(
         f"No 'num_epochs' metric found in batch. Got metrics: "
-        f"{[m.key for m in batch['metrics']]}"
+        f"{[m.key for m in metrics]}"
     )

tests/unit_tests/datasets/test_stop_after_one_epoch.py

@@ -13,6 +13,7 @@
 from forge.data.datasets import HfIterableDataset
 from forge.data.utils import extract_epoch_from_batch, StopAfterOneEpoch
 from forge.observability.metrics import Metric, Reduce
+from forge.types import TrainBatch
 from tests.test_utils import gpu_test
 from torch.testing._internal.common_fsdp import FSDPTest
 from torchdata.stateful_dataloader import StatefulDataLoader
@@ -25,57 +26,73 @@ def create_test_json_file(path: Path, num_samples: int) -> None:
             f.write(f'{{"id": {i}, "tokens": [{i}, {i + 1}]}}\n')
 
 
-def simple_collate(batch):
-    """Simple collate function that mimics collate_packed behavior.
+def simple_collate(batch) -> TrainBatch:
+    """Simple collate function that returns a TrainBatch.
 
-    Stacks tensors, extends metrics list, keeps other fields as lists.
+    Stacks tensors into model_inputs, keeps metrics in meta.
     """
-    collated = {}
+    model_inputs = {}
+    meta = {}
+
     for key in batch[0].keys():
         if isinstance(batch[0][key], torch.Tensor):
-            collated[key] = torch.stack([sample[key] for sample in batch], dim=0)
+            model_inputs[key] = torch.stack([sample[key] for sample in batch], dim=0)
         elif key == "metrics":
             # Extend all metrics into a single list
-            collated[key] = []
+            meta[key] = []
             for sample in batch:
-                collated[key].extend(sample[key])
+                meta[key].extend(sample[key])
         else:
-            collated[key] = [sample[key] for sample in batch]
-    return collated
+            meta[key] = [sample[key] for sample in batch]
+
+    return TrainBatch(model_inputs=model_inputs, loss_inputs={}, meta=meta)
 
 
 class TestExtractEpochFromBatch:
     """Test extract_epoch_from_batch helper function."""
 
     def test_extract_epoch_from_batch_success(self):
         """Test extracting epoch from valid batch with metrics."""
-        batch = {
-            "tokens": torch.tensor([1, 2, 3]),
-            "metrics": [
-                Metric(key="dataset/test/num_epochs", value=2, reduction=Reduce.MAX),
-                Metric(
-                    key="dataset/test/other_metric", value=42, reduction=Reduce.MEAN
-                ),
-            ],
-        }
+        batch = TrainBatch(
+            model_inputs={"tokens": torch.tensor([1, 2, 3])},
+            loss_inputs={},
+            meta={
+                "metrics": [
+                    Metric(
+                        key="dataset/test/num_epochs", value=2, reduction=Reduce.MAX
+                    ),
+                    Metric(
+                        key="dataset/test/other_metric", value=42, reduction=Reduce.MEAN
+                    ),
+                ],
+            },
+        )
         epoch = extract_epoch_from_batch(batch)
         assert epoch == 2
 
     def test_extract_epoch_missing_metrics_field(self):
         """Test error when batch has no 'metrics' field."""
-        batch = {"tokens": torch.tensor([1, 2, 3])}
+        batch = TrainBatch(
+            model_inputs={"tokens": torch.tensor([1, 2, 3])},
+            loss_inputs={},
+            meta={},
+        )
         with pytest.raises(ValueError, match="Batch missing 'metrics' field"):
             extract_epoch_from_batch(batch)
 
     def test_extract_epoch_no_num_epochs_metric(self):
         """Test error when no num_epochs metric found."""
-        batch = {
-            "metrics": [
-                Metric(
-                    key="dataset/test/other_metric", value=42, reduction=Reduce.MEAN
-                ),
-            ]
-        }
+        batch = TrainBatch(
+            model_inputs={},
+            loss_inputs={},
+            meta={
+                "metrics": [
+                    Metric(
+                        key="dataset/test/other_metric", value=42, reduction=Reduce.MEAN
+                    ),
+                ],
+            },
+        )
         with pytest.raises(ValueError, match="No 'num_epochs' metric found"):
             extract_epoch_from_batch(batch)