diff --git a/apps/sft/main.py b/apps/sft/main.py index 38ae27e1e..be3cea6d6 100644 --- a/apps/sft/main.py +++ b/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,15 +207,14 @@ 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 @@ -222,7 +222,8 @@ def forward_backward( # 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,7 +275,7 @@ 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, @@ -282,8 +283,7 @@ def train_step(self, batch) -> None: # 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() diff --git a/src/forge/data/__init__.py b/src/forge/data/__init__.py index 6564817a9..2312a7df6 100644 --- a/src/forge/data/__init__.py +++ b/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", diff --git a/src/forge/data/collate.py b/src/forge/data/collate.py index adc7b8200..cc70ec703 100644 --- a/src/forge/data/collate.py +++ b/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]: diff --git a/src/forge/data/utils.py b/src/forge/data/utils.py index f1ca07773..7281e9d1a 100644 --- a/src/forge/data/utils.py +++ b/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,14 +305,14 @@ 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 @@ -317,17 +320,19 @@ def extract_epoch_from_batch(batch: dict) -> int: 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]}" ) diff --git a/tests/unit_tests/datasets/test_stop_after_one_epoch.py b/tests/unit_tests/datasets/test_stop_after_one_epoch.py index 89ff0790c..955e6c312 100644 --- a/tests/unit_tests/datasets/test_stop_after_one_epoch.py +++ b/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,23 +26,26 @@ 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: @@ -49,33 +53,46 @@ class TestExtractEpochFromBatch: 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)