Flux Validation (#1518)

# This pr implements the validator class for flux following the method
discussed in Stable Diffusion 3 paper.

The paper shows that creating 8 equidistant timesteps and calculating
the average loss on them will result in a highly correlated loss to
external validation methods such as CLIP or FID score.

This pr's implementation rather than creating 8 stratified timesteps per
sample, only applies one of these equidistant timesteps to each sample
in a round-robin fashion. Aggregated over many samples in a validation
set, this should give a similar validation score as the full timestep
method, but will process more validation samples quickly.

### Implementations
- Integrates the image generation evaluation in the validation step,
users can
- Refactors and combines eval job_config with validation
- Adds an `all_timesteps` option to the job_config to choose whether to
use round robin timesteps or full timesteps per sample
- Creates validator class and validation dataloader for flux, validator
dataloader handles generating timesteps for round-robin method of
validation

### Enabling all timesteps
Developers can enable the full timestamp method of validation by setting
`all_timesteps = True` in the flux validation job config. Enabling
all_timesteps may require tweaking some hyperparams
`validation.local_batch_size, validation.steps` to prevent spiking
memory and optimizing throughput. By using a ratio of around 1/4 for
`validation.local_batch_size` to `training.local_batch_size` will not
spike the memory higher than training when `fsdp = 8`.

Below we can see the difference between round robin and all timesteps.
In the comparison the total number of validation samples processed is
the same, but in `all_timesteps=True` configuration we have to lower the
batch size to prevent memory spiking. All timesteps also achieves a
higher throughput (tps) but still processes total samples of validation
set more slowly.

| Round Robin (batch_size=32, steps=1, fsdp=8) | All Timesteps
(batch_size=8, steps=4, fsdp=8) |
| ---- | --- |
| <img width="682" height="303" alt="Screenshot 2025-08-01 at 3 46
42 PM"
src="https://github.com/user-attachments/assets/30328bfe-4c3c-4912-a329-2b94c834b67b"
/> | <img width="719" height="308" alt="Screenshot 2025-08-01 at 3 30
10 PM"
src="https://github.com/user-attachments/assets/c7325d21-8a7b-41d9-a0d2-74052e425083"
/> |

Author: wesleytruong

torchtitan/components/validate.py

@@ -80,7 +80,7 @@ def validate(
         self,
         model_parts: list[nn.Module],
         step: int,
-    ) -> dict[str, float]:
+    ) -> None:
         # Set model to eval mode
         model = model_parts[0]
         model.eval()

torchtitan/experiments/flux/__init__.py

@@ -17,6 +17,7 @@
 from .model.args import FluxModelArgs
 from .model.autoencoder import AutoEncoderParams
 from .model.model import FluxModel
+from .validate import build_flux_validator
 
 __all__ = [
     "FluxModelArgs",
@@ -117,5 +118,6 @@
         build_dataloader_fn=build_flux_dataloader,
         build_tokenizer_fn=None,
         build_loss_fn=build_mse_loss,
+        build_validator_fn=build_flux_validator,
     )
 )

torchtitan/experiments/flux/dataset/flux_dataset.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import itertools
 import math
 from dataclasses import dataclass
 from typing import Any, Callable, Optional
@@ -103,6 +104,38 @@ def _cc12m_wds_data_processor(
         "image": img,
         "clip_tokens": clip_tokens,  # type: List[int]
         "t5_tokens": t5_tokens,  # type: List[int]
+        "prompt": sample["txt"],  # type: str
+    }
+
+
+def _coco_data_processor(
+    sample: dict[str, Any],
+    t5_tokenizer: FluxTokenizer,
+    clip_tokenizer: FluxTokenizer,
+    output_size: int = 256,
+) -> dict[str, Any]:
+    """
+    Preprocess COCO dataset sample image and text for Flux model.
+
+    Args:
+        sample: A sample from dataset
+        t5_encoder: T5 encoder
+        clip_encoder: CLIP encoder
+        output_size: The output image size
+
+    """
+    img = _process_cc12m_image(sample["image"], output_size=output_size)
+    prompt = sample["caption"]
+    if isinstance(prompt, list):
+        prompt = prompt[0]
+    t5_tokens = t5_tokenizer.encode(prompt)
+    clip_tokens = clip_tokenizer.encode(prompt)
+
+    return {
+        "image": img,
+        "clip_tokens": clip_tokens,  # type: List[int]
+        "t5_tokens": t5_tokens,  # type: List[int]
+        "prompt": prompt,  # type: str
     }
 
 
@@ -126,6 +159,11 @@ class TextToImageDatasetConfig:
         ),
         data_processor=_cc12m_wds_data_processor,
     ),
+    "coco-validation": TextToImageDatasetConfig(
+        path="howard-hou/COCO-Text",
+        loader=lambda path: load_dataset(path, split="validation", streaming=True),
+        data_processor=_coco_data_processor,
+    ),
 }
 
 
@@ -242,8 +280,9 @@ def __iter__(self):
 
             # skip low quality image or image with color channel = 1
             if sample_dict["image"] is None:
+                sample = sample.get("__key__", "unknown")
                 logger.warning(
-                    f"Low quality image {sample['__key__']} is skipped in Flux Dataloader."
+                    f"Low quality image {sample} is skipped in Flux Dataloader."
                 )
                 continue
 
@@ -308,3 +347,87 @@ def build_flux_dataloader(
         dp_world_size=dp_world_size,
         batch_size=batch_size,
     )
+
+
+class FluxValidationDataset(FluxDataset):
+    """
+    Adds logic to generate timesteps for flux validation method described in SD3 paper
+
+    Args:
+    generate_timesteps (bool): Generate stratified timesteps in round-robin style for validation
+    """
+
+    def __init__(
+        self,
+        dataset_name: str,
+        dataset_path: Optional[str],
+        t5_tokenizer: BaseTokenizer,
+        clip_tokenizer: BaseTokenizer,
+        job_config: Optional[JobConfig] = None,
+        dp_rank: int = 0,
+        dp_world_size: int = 1,
+        generate_timesteps: bool = True,
+    ) -> None:
+        # Call parent constructor correctly
+        super().__init__(
+            dataset_name=dataset_name,
+            dataset_path=dataset_path,
+            t5_tokenizer=t5_tokenizer,
+            clip_tokenizer=clip_tokenizer,
+            job_config=job_config,
+            dp_rank=dp_rank,
+            dp_world_size=dp_world_size,
+            infinite=False,
+        )
+
+        # Initialize timestep generation for validation
+        self.generate_timesteps = generate_timesteps
+        if self.generate_timesteps:
+            # Generate stratified timesteps as described in SD3 paper
+            val_timesteps = [1 / 8 * (i + 0.5) for i in range(8)]
+            self.timestep_cycle = itertools.cycle(val_timesteps)
+
+    def __iter__(self):
+        # Get parent iterator and add timesteps to each sample
+        parent_iterator = super().__iter__()
+
+        for sample_dict, labels in parent_iterator:
+            # Add timestep to the sample dict if timestep generation is enabled
+            if self.generate_timesteps:
+                sample_dict["timestep"] = next(self.timestep_cycle)
+
+            yield sample_dict, labels
+
+
+def build_flux_validation_dataloader(
+    dp_world_size: int,
+    dp_rank: int,
+    job_config: JobConfig,
+    # This parameter is not used, keep it for compatibility
+    tokenizer: BaseTokenizer | None,
+    generate_timestamps: bool = True,
+) -> ParallelAwareDataloader:
+    """Build a 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
+
+    t5_tokenizer, clip_tokenizer = build_flux_tokenizer(job_config)
+
+    ds = FluxValidationDataset(
+        dataset_name=dataset_name,
+        dataset_path=dataset_path,
+        t5_tokenizer=t5_tokenizer,
+        clip_tokenizer=clip_tokenizer,
+        job_config=job_config,
+        dp_rank=dp_rank,
+        dp_world_size=dp_world_size,
+        generate_timesteps=generate_timestamps,
+    )
+
+    return ParallelAwareDataloader(
+        dataset=ds,
+        dp_rank=dp_rank,
+        dp_world_size=dp_world_size,
+        batch_size=batch_size,
+    )

torchtitan/experiments/flux/job_config.py

@@ -36,7 +36,7 @@ class Encoder:
 
 
 @dataclass
-class Eval:
+class Validation:
     enable_classifier_free_guidance: bool = False
     """Whether to use classifier-free guidance during sampling"""
     classifier_free_guidance_scale: float = 5.0
@@ -45,8 +45,13 @@ class Eval:
     """How many denoising steps to sample when generating an image"""
     eval_freq: int = 100
     """Frequency of evaluation/sampling during training"""
+    save_img_count: int = 1
+    """ How many images to generate and save during validation, starting from
+    the beginning of validation set, -1 means generate on all samples"""
     save_img_folder: str = "img"
     """Directory to save image generated/sampled from the model"""
+    all_timesteps: bool = False
+    """Whether to generate all stratified timesteps per sample or use round robin"""
 
 
 @dataclass
@@ -57,4 +62,4 @@ class JobConfig:
 
     training: Training = field(default_factory=Training)
     encoder: Encoder = field(default_factory=Encoder)
-    eval: Eval = field(default_factory=Eval)
+    validation: Validation = field(default_factory=Validation)

torchtitan/experiments/flux/sampling.py

@@ -93,7 +93,9 @@ def generate_image(
     img_height = 16 * (job_config.training.img_size // 16)
     img_width = 16 * (job_config.training.img_size // 16)
 
-    enable_classifier_free_guidance = job_config.eval.enable_classifier_free_guidance
+    enable_classifier_free_guidance = (
+        job_config.validation.enable_classifier_free_guidance
+    )
 
     # Tokenize the prompt. Unsqueeze to add a batch dimension.
     clip_tokens = clip_tokenizer.encode(prompt).unsqueeze(0)
@@ -132,7 +134,7 @@ def generate_image(
         model=model,
         img_width=img_width,
         img_height=img_height,
-        denoising_steps=job_config.eval.denoising_steps,
+        denoising_steps=job_config.validation.denoising_steps,
         clip_encodings=batch["clip_encodings"],
         t5_encodings=batch["t5_encodings"],
         enable_classifier_free_guidance=enable_classifier_free_guidance,
@@ -142,7 +144,7 @@ def generate_image(
         empty_clip_encodings=(
             empty_batch["clip_encodings"] if enable_classifier_free_guidance else None
         ),
-        classifier_free_guidance_scale=job_config.eval.classifier_free_guidance_scale,
+        classifier_free_guidance_scale=job_config.validation.classifier_free_guidance_scale,
     )
 
     img = autoencoder.decode(img)

torchtitan/experiments/flux/tests/integration_tests.py

@@ -64,6 +64,9 @@ def build_test_list():
             "Checkpoint Integration Test - Save Model Only fp32",
             "last_save_model_only_fp32",
         ),
+        OverrideDefinitions(
+            [["--validation.enabled"]], "Flux Validation Test", "validation"
+        ),
         # Parallelism tests.
         OverrideDefinitions(
             [

torchtitan/experiments/flux/tests/test_generate_image.py

@@ -57,12 +57,12 @@ def test_generate_image(self):
                 "--training.img_size",
                 str(img_width),
                 # eval params
-                "--eval.denoising_steps",
+                "--validation.denoising_steps",
                 str(num_steps),
-                "--eval.enable_classifier_free_guidance",
-                "--eval.classifier_free_guidance_scale",
+                "--validation.enable_classifier_free_guidance",
+                "--validation.classifier_free_guidance_scale",
                 str(classifier_free_guidance_scale),
-                "--eval.save_img_folder",
+                "--validation.save_img_folder",
                 "img",
             ]
         )
@@ -120,7 +120,7 @@ def test_generate_image(self):
         save_image(
             name=f"img_unit_test_{config.training.seed}.jpg",
             output_dir=os.path.join(
-                config.job.dump_folder, config.eval.save_img_folder
+                config.job.dump_folder, config.validation.save_img_folder
             ),
             x=image,
             add_sampling_metadata=True,

torchtitan/experiments/flux/tests/unit_tests/test_flux_dataloader.py

@@ -79,7 +79,9 @@ def test_load_dataset(self):
                 for i in range(0, num_steps):
                     input_data, labels = next(it)
 
-                    assert len(input_data) == 2  # (clip_encodings, t5_encodings)
+                    assert (
+                        len(input_data) == 3
+                    )  # (clip_encodings, t5_encodings, prompt)
                     assert labels.shape == (batch_size, 3, 256, 256)
                     assert input_data["clip_tokens"].shape == (
                         batch_size,

torchtitan/experiments/flux/train.py

@@ -5,21 +5,18 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
-from typing import Iterable, Optional
+from typing import Optional
 
 import torch
-from torch.distributed.fsdp import FSDPModule
 
 from torchtitan.config import ConfigManager, JobConfig, TORCH_DTYPE_MAP
 from torchtitan.distributed import utils as dist_utils
 from torchtitan.tools.logging import init_logger, logger
 from torchtitan.train import Trainer
 
-from .dataset.tokenizer import build_flux_tokenizer
 from .infra.parallelize import parallelize_encoders
 from .model.autoencoder import load_ae
 from .model.hf_embedder import FluxEmbedder
-from .sampling import generate_image, save_image
 from .utils import (
     create_position_encoding_for_latents,
     pack_latents,
@@ -81,6 +78,15 @@ def __init__(self, job_config: JobConfig):
             job_config=job_config,
         )
 
+        if job_config.validation.enabled:
+            self.validator.flux_init(
+                device=self.device,
+                _dtype=self._dtype,
+                autoencoder=self.autoencoder,
+                t5_encoder=self.t5_encoder,
+                clip_encoder=self.clip_encoder,
+            )
+
     def forward_backward_step(
         self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor
     ) -> torch.Tensor:
@@ -147,64 +153,6 @@ def forward_backward_step(
 
         return loss
 
-    def eval_step(self, prompt: str = "A photo of a cat"):
-        """
-        Evaluate the Flux model.
-        1) generate and save images every few steps. Currently, we run the eval and on the same
-        prompts across all DP ranks. We will change this behavior to run on validation set prompts.
-        Due to random noise generation, results could be different across DP ranks cause we assign
-        different random seeds to each DP rank.
-        2) [TODO] Calculate loss with fixed t value on validation set.
-        """
-
-        t5_tokenizer, clip_tokenizer = build_flux_tokenizer(self.job_config)
-
-        image = generate_image(
-            device=self.device,
-            dtype=self._dtype,
-            job_config=self.job_config,
-            model=self.model_parts[0],
-            prompt=prompt,  # TODO(jianiw): change this to a prompt from validation set
-            autoencoder=self.autoencoder,
-            t5_tokenizer=t5_tokenizer,
-            clip_tokenizer=clip_tokenizer,
-            t5_encoder=self.t5_encoder,
-            clip_encoder=self.clip_encoder,
-        )
-
-        save_image(
-            name=f"image_rank{str(torch.distributed.get_rank())}_{self.step}.png",
-            output_dir=os.path.join(
-                self.job_config.job.dump_folder, self.job_config.eval.save_img_folder
-            ),
-            x=image,
-            add_sampling_metadata=True,
-            prompt=prompt,
-        )
-
-        # Reshard after run forward pass in eval_step.
-        # This is to ensure the model weights are sharded the same way for checkpoint saving.
-        for module in self.model_parts[0].modules():
-            if isinstance(module, FSDPModule):
-                module.reshard()
-
-    def train_step(
-        self, data_iterator: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]]
-    ):
-        super().train_step(data_iterator)
-
-        # Evaluate the model during training
-        if (
-            self.step % self.job_config.eval.eval_freq == 0
-            or self.step == self.job_config.training.steps
-        ):
-            model = self.model_parts[0]
-            model.eval()
-            # We need to set reshard_after_forward before last forward pass.
-            # So the model wieghts are sharded the same way for checkpoint saving.
-            self.eval_step()
-            model.train()
-
 
 if __name__ == "__main__":
     init_logger()