basic validator implementation (#1362)

Update PR Summary:
Implements a validator that can be easily plugged into the training loop
and configured from the job specific config file.
Changes:
- Created validation section in job_config with enabled, dataset, freq,
and steps fields
- Created a builder function for validator in train_spec
- Created a separate builder function for validation dataset in
hf_dataset.py
- Created validator class
- Validator class initializes a build_validation_hf_loader but leaves
this dataloader function unexposed to the train_spec
  - Validator class supports ddp, fsdp, cp, and tp (but not pp yet)
- Integrated validation call into training loop
- Creates an integration test to test parallelization

Updated tests training the same base model weights from a seed
checkpoint:

| FSDP=2 | FSDP=2,TP=4 |
| --- | --- |
| <img width="978" alt="Screenshot 2025-07-09 at 4 33 53 PM"
src="https://github.com/user-attachments/assets/a1fa9fa7-df2f-4302-aa4a-d556a5699ba9"
/> | <img width="978" alt="Screenshot 2025-07-09 at 4 33 53 PM"
src="https://github.com/user-attachments/assets/a1fa9fa7-df2f-4302-aa4a-d556a5699ba9"
/> |

| FSDP=2,CP=4 | FSDP=2,TP=2,CP=2 |
| --- | --- |
| <img width="972" alt="Screenshot 2025-07-09 at 4 39 35 PM"
src="https://github.com/user-attachments/assets/56d62841-5841-4969-85b1-803705892465"
/> | <img width="970" alt="Screenshot 2025-07-09 at 4 28 57 PM"
src="https://github.com/user-attachments/assets/f7d33fa8-ca2c-48f1-931c-8d4c017a47ce"
/> |

Author: wesleytruong

tests/integration_tests.py

@@ -509,6 +509,20 @@ def build_test_list():
             "gradient_accumulation",
             ngpu=2,
         ),
+        OverrideDefinitions(
+            [
+                [
+                    "--validation.enabled",
+                    "--validation.dataset c4_test",
+                    "--parallelism.data_parallel_replicate_degree=2",
+                    "--parallelism.tensor_parallel_degree=2",
+                    "--parallelism.context_parallel_degree=2",
+                ],
+            ],
+            "Validation test with fsdp, tp, cp",
+            "validation_fsdp_tp_cp",
+            ngpu=8,
+        ),
     ]
     return integration_tests_flavors
 

torchtitan/components/validate.py

@@ -0,0 +1,163 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Generator
+
+import torch
+import torch.nn as nn
+from torch.distributed.fsdp import FSDPModule
+from torchtitan.components.dataloader import BaseDataLoader
+from torchtitan.components.loss import LossFunction
+from torchtitan.components.tokenizer import Tokenizer
+from torchtitan.config_manager import JobConfig
+from torchtitan.datasets.hf_datasets import build_hf_validation_dataloader
+from torchtitan.distributed import ParallelDims, utils as dist_utils
+from torchtitan.tools import utils
+from torchtitan.tools.logging import logger
+
+
+class BaseValidator:
+    def __init__(self, job_config: JobConfig):
+        self.job_config = job_config
+
+    def validate(self, model_parts: list[nn.Module]) -> dict[str, float]:
+        raise NotImplementedError("validate method not implemented")
+
+    def should_validate(self, step: int) -> bool:
+        return step % self.job_config.validation.freq == 0
+
+
+class Validator(BaseValidator):
+    """
+    Simple validator focused on correctness and integration.
+
+    Args:
+        job_config: Job configuration
+        validation_dataloader: The validation dataloader
+        loss_fn: Loss function to use for validation
+        model: The model to validate (single model, no parallelism)
+    """
+
+    validation_dataloader: BaseDataLoader
+
+    def __init__(
+        self,
+        job_config: JobConfig,
+        dp_world_size: int,
+        dp_rank: int,
+        tokenizer: Tokenizer,
+        parallel_dims: ParallelDims,
+        world_mesh: torch.distributed.DeviceMesh,
+        loss_fn: LossFunction,
+        validation_context: Generator[None, None, None],
+        maybe_enable_amp: Generator[None, None, None],
+    ):
+        self.job_config = job_config
+        self.parallel_dims = parallel_dims
+        self.world_mesh = world_mesh
+        self.loss_fn = loss_fn
+        self.validation_dataloader = build_hf_validation_dataloader(
+            job_config=job_config,
+            dp_world_size=dp_world_size,
+            dp_rank=dp_rank,
+            tokenizer=tokenizer,
+        )
+        self.validation_context = validation_context
+        self.maybe_enable_amp = maybe_enable_amp
+
+    @torch.no_grad()
+    def validate(
+        self,
+        model_parts: list[nn.Module],
+    ) -> dict[str, float]:
+        # Set model to eval mode
+        # TODO: currently does not support pipeline parallelism
+        model = model_parts[0]
+        model.eval()
+
+        accumulated_losses = []
+        device_type = utils.device_type
+        num_steps = 0
+
+        for input_dict, labels in self.validation_dataloader:
+            if (
+                self.job_config.validation.steps != -1
+                and num_steps >= self.job_config.validation.steps
+            ):
+                break
+
+            for k, v in input_dict.items():
+                input_dict[k] = v.to(device_type)
+            inputs = input_dict["input"]
+            labels = labels.to(device_type)
+
+            optional_context_parallel_ctx = (
+                dist_utils.create_context_parallel_ctx(
+                    cp_mesh=self.world_mesh["cp"],
+                    cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts],
+                    cp_seq_dims=[1, 1] + [0 for _ in model_parts],
+                    cp_no_restore_buffers={inputs, labels},
+                    cp_rotate_method=self.job_config.parallelism.context_parallel_rotate_method,
+                )
+                if self.parallel_dims.cp_enabled
+                else None
+            )
+
+            with self.validation_context(optional_context_parallel_ctx):
+                assert len(model_parts) == 1
+                with self.maybe_enable_amp:
+                    predictions = model(inputs)
+                    loss = self.loss_fn(predictions, labels)
+
+            accumulated_losses.append(loss.detach())
+
+            num_steps += 1
+
+        # Compute average loss
+        loss = torch.sum(torch.stack(accumulated_losses))
+        loss /= num_steps
+        if self.parallel_dims.dp_cp_enabled:
+            global_avg_loss = dist_utils.dist_mean(loss, self.world_mesh["dp_cp"])
+        else:
+            global_avg_loss = loss
+
+        logger.info(
+            f"Validation completed. Average loss: {global_avg_loss:.4f} over {num_steps} batches"
+        )
+
+        # Reshard after run forward pass
+        # This is to ensure the model weights are sharded the same way for checkpoint saving.
+        for module in model.modules():
+            if isinstance(module, FSDPModule):
+                module.reshard()
+
+        # Set model back to train mode
+        model.train()
+
+
+def build_validator(
+    job_config: JobConfig,
+    dp_world_size: int,
+    dp_rank: int,
+    tokenizer: Tokenizer,
+    parallel_dims: ParallelDims,
+    world_mesh: torch.distributed.DeviceMesh,
+    loss_fn: LossFunction,
+    validation_context: Generator[None, None, None],
+    maybe_enable_amp: Generator[None, None, None],
+) -> BaseValidator:
+    """Build a simple validator focused on correctness."""
+    return Validator(
+        job_config=job_config,
+        dp_world_size=dp_world_size,
+        dp_rank=dp_rank,
+        tokenizer=tokenizer,
+        parallel_dims=parallel_dims,
+        world_mesh=world_mesh,
+        loss_fn=loss_fn,
+        validation_context=validation_context,
+        maybe_enable_amp=maybe_enable_amp,
+    )

torchtitan/config_manager.py

@@ -665,6 +665,35 @@ class Experimental:
     """
 
 
+@dataclass
+class Validation:
+    enabled: bool = False
+    """Enable validation to default run validation after each training loop"""
+
+    dataset: str = "c4_validation"
+    """Dataset to use for validation"""
+
+    dataset_path: str | None = None
+    """Path to dataset to use for validation"""
+
+    local_batch_size: int = 8
+    """Batch size for validation"""
+
+    seq_len: int = 2048
+    """Sequence length for validation"""
+
+    freq: int = 10
+    """Frequency of validation"""
+
+    steps: int = -1
+    """Number of steps to take in the validation set, -1 means consuming all the data in the validation dataset"""
+
+    def __post_init__(self):
+        assert (
+            self.steps > 0 or self.steps == -1
+        ), "validation steps must be positive or -1"
+
+
 @dataclass
 class JobConfig:
     """
@@ -689,6 +718,7 @@ class JobConfig:
     memory_estimation: MemoryEstimation = field(default_factory=MemoryEstimation)
     fault_tolerance: FaultTolerance = field(default_factory=FaultTolerance)
     experimental: Experimental = field(default_factory=Experimental)
+    validation: Validation = field(default_factory=Validation)
 
     def to_dict(self) -> dict[str, Any]:
         return asdict(self)

torchtitan/datasets/hf_datasets.py

@@ -5,6 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 from dataclasses import dataclass
+
+from functools import partial
 from typing import Any, Callable
 
 import torch
@@ -20,9 +22,9 @@
 from torchtitan.tools.logging import logger
 
 
-def _load_c4_dataset(dataset_path: str):
+def _load_c4_dataset(dataset_path: str, split: str):
     """Load C4 dataset with default configuration."""
-    return load_dataset(dataset_path, name="en", split="train", streaming=True)
+    return load_dataset(dataset_path, name="en", split=split, streaming=True)
 
 
 def _process_c4_text(sample: dict[str, Any]) -> str:
@@ -41,14 +43,19 @@ class DatasetConfig:
 DATASETS = {
     "c4": DatasetConfig(
         path="allenai/c4",
-        loader=_load_c4_dataset,
+        loader=partial(_load_c4_dataset, split="train"),
         text_processor=_process_c4_text,
     ),
     "c4_test": DatasetConfig(
         path="tests/assets/c4_test",
         loader=lambda path: load_dataset(path, split="train"),
         text_processor=_process_c4_text,
     ),
+    "c4_validation": DatasetConfig(
+        path="allenai/c4",
+        loader=partial(_load_c4_dataset, split="validation"),
+        text_processor=_process_c4_text,
+    ),
 }
 
 
@@ -193,3 +200,33 @@ def build_hf_dataloader(
         dp_world_size=dp_world_size,
         batch_size=batch_size,
     )
+
+
+def build_hf_validation_dataloader(
+    dp_world_size: int,
+    dp_rank: int,
+    tokenizer: Tokenizer,
+    job_config: JobConfig,
+) -> ParallelAwareDataloader:
+    """Build a validation data loader for HuggingFace datasets."""
+    dataset_name = job_config.validation.dataset
+    dataset_path = job_config.validation.dataset_path
+    batch_size = job_config.validation.local_batch_size
+    seq_len = job_config.validation.seq_len
+
+    hf_ds = HuggingFaceDataset(
+        dataset_name=dataset_name,
+        dataset_path=dataset_path,
+        tokenizer=tokenizer,
+        seq_len=seq_len,
+        dp_rank=dp_rank,
+        dp_world_size=dp_world_size,
+        infinite=False,
+    )
+
+    return ParallelAwareDataloader(
+        dataset=hf_ds,
+        dp_rank=dp_rank,
+        dp_world_size=dp_world_size,
+        batch_size=batch_size,
+    )

torchtitan/models/llama3/__init__.py

@@ -9,6 +9,7 @@
 from torchtitan.components.loss import build_cross_entropy_loss
 from torchtitan.components.lr_scheduler import build_lr_schedulers
 from torchtitan.components.optimizer import build_optimizers
+from torchtitan.components.validate import build_validator
 from torchtitan.datasets.hf_datasets import build_hf_dataloader
 from torchtitan.datasets.tokenizer.tiktoken import build_tiktoken_tokenizer
 from torchtitan.protocols.train_spec import register_train_spec, TrainSpec
@@ -81,5 +82,6 @@
         build_dataloader_fn=build_hf_dataloader,
         build_tokenizer_fn=build_tiktoken_tokenizer,
         build_loss_fn=build_cross_entropy_loss,
+        build_validator_fn=build_validator,
     )
 )

torchtitan/models/llama3/train_configs/debug_model.toml

@@ -71,3 +71,9 @@ selective_ac_option = '2'  # 'int' = ac every positive int layer or 'op', ac bas
 enable_fsdp_float8_all_gather = false
 precompute_float8_dynamic_scale_for_fsdp = false
 filter_fqns = ["output"]
+
+[validation]
+enabled = false
+dataset = "c4_validation"
+freq = 5
+steps = 10

torchtitan/protocols/train_spec.py

@@ -23,6 +23,7 @@
 from torchtitan.components.metrics import MetricsProcessor
 from torchtitan.components.optimizer import OptimizersContainer
 from torchtitan.components.tokenizer import Tokenizer
+from torchtitan.components.validate import BaseValidator
 from torchtitan.config_manager import JobConfig
 from torchtitan.distributed import ParallelDims
 
@@ -80,6 +81,7 @@ def init_weights(self, buffer_device: torch.device | None = None) -> None:
     [OptimizersContainer, JobConfig], LRSchedulersContainer
 ]
 LossFunctionBuilder: TypeAlias = Callable[..., LossFunction]
+ValidatorBuilder: TypeAlias = Callable[..., BaseValidator]
 
 
 @dataclass
@@ -94,6 +96,7 @@ class TrainSpec:
     build_dataloader_fn: DataLoaderBuilder
     build_tokenizer_fn: TokenizerBuilder | None
     build_loss_fn: LossFunctionBuilder
+    build_validator_fn: ValidatorBuilder | None = None
     build_metrics_processor_fn: MetricsProcessorBuilder | None = None