Validator integration with current metrics processor for logging (#1395)

Integrated the validator together with metrics processor for better
metrics logging.

Key changes:
- Metrics processor is passed to validator within training loop
- Validator can reuse metrics processor's built-in functionalities such
as memory profiling, throughput tracking, and tensorboard/wandb logging

This is how the new logging looks from terminal:
<img width="959" height="374" alt="Screenshot 2025-07-14 at 3 22 56 PM"
src="https://github.com/user-attachments/assets/b16a9e00-3ab2-46ed-a42a-0c92d13697cb"
/>

Author: wesleytruong

torchtitan/components/metrics.py

@@ -403,7 +403,7 @@ def log(
             f"{color.red}step: {step:2}  "
             f"{color.green}loss: {global_avg_loss:7.4f}  "
             f"{color.orange}grad_norm: {grad_norm:7.4f}  "
-            f"{color.yellow}memory: {device_mem_stats.max_reserved_gib:5.2f}GiB"
+            f"{color.turquoise}memory: {device_mem_stats.max_reserved_gib:5.2f}GiB"
             f"({device_mem_stats.max_reserved_pct:.2f}%)  "
             f"{color.blue}tps: {round(tps):,}  "
             f"{color.cyan}tflops: {tflops:,.2f}  "
@@ -415,6 +415,39 @@ def log(
         self.time_last_log = time.perf_counter()
         self.device_memory_monitor.reset_peak_stats()
 
+    def log_validation(self, loss: float, step: int):
+        time_delta = time.perf_counter() - self.time_last_log
+
+        device_mem_stats = self.device_memory_monitor.get_peak_stats()
+
+        # tokens per second per device, abbreviated as tps
+        tps = self.ntokens_since_last_log / (
+            time_delta * self.parallel_dims.non_data_parallel_size
+        )
+
+        metrics = {
+            "validation_metrics/loss": loss,
+            "validation_metrics/throughput(tps)": tps,
+            "validation_metrics/memory/max_active(GiB)": device_mem_stats.max_active_gib,
+            "validation_metrics/memory/max_active(%)": device_mem_stats.max_active_pct,
+            "validation_metrics/memory/max_reserved(GiB)": device_mem_stats.max_reserved_gib,
+            "validation_metrics/memory/max_reserved(%)": device_mem_stats.max_reserved_pct,
+        }
+        self.logger.log(metrics, step)
+
+        color = self.color
+        logger.info(
+            f"{color.yellow}validate step: {step:2}  "
+            f"{color.green}loss: {loss:7.4f}  "
+            f"{color.turquoise}memory: {device_mem_stats.max_reserved_gib:5.2f}GiB"
+            f"({device_mem_stats.max_reserved_pct:.2f}%)  "
+            f"{color.blue}tps: {round(tps):,}{color.reset}"
+        )
+
+        self.ntokens_since_last_log = 0
+        self.time_last_log = time.perf_counter()
+        self.device_memory_monitor.reset_peak_stats()
+
     def close(self):
         self.logger.close()
 

torchtitan/components/validate.py

@@ -11,12 +11,12 @@
 from torch.distributed.fsdp import FSDPModule
 from torchtitan.components.dataloader import BaseDataLoader
 from torchtitan.components.loss import LossFunction
+from torchtitan.components.metrics import MetricsProcessor
 from torchtitan.components.tokenizer import BaseTokenizer
 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:
@@ -53,6 +53,7 @@ def __init__(
         loss_fn: LossFunction,
         validation_context: Generator[None, None, None],
         maybe_enable_amp: Generator[None, None, None],
+        metrics_processor: MetricsProcessor,
     ):
         self.job_config = job_config
         self.parallel_dims = parallel_dims
@@ -65,11 +66,13 @@ def __init__(
         )
         self.validation_context = validation_context
         self.maybe_enable_amp = maybe_enable_amp
+        self.metrics_processor = metrics_processor
 
     @torch.no_grad()
     def validate(
         self,
         model_parts: list[nn.Module],
+        step: int,
     ) -> dict[str, float]:
         # Set model to eval mode
         # TODO: currently does not support pipeline parallelism
@@ -89,6 +92,7 @@ def validate(
             ):
                 break
 
+            self.metrics_processor.ntokens_since_last_log += labels.numel()
             for k, v in input_dict.items():
                 input_dict[k] = v.to(device_type)
             inputs = input_dict["input"]
@@ -124,11 +128,9 @@ def validate(
                 loss, parallel_dims.world_mesh["dp_cp"]
             )
         else:
-            global_avg_loss = loss
+            global_avg_loss = loss.item()
 
-        logger.info(
-            f"Validation completed. Average loss: {global_avg_loss:.4f} over {num_steps} batches"
-        )
+        self.metrics_processor.log_validation(loss=global_avg_loss, step=step)
 
         # Reshard after run forward pass
         # This is to ensure the model weights are sharded the same way for checkpoint saving.
@@ -149,6 +151,7 @@ def build_validator(
     loss_fn: LossFunction,
     validation_context: Generator[None, None, None],
     maybe_enable_amp: Generator[None, None, None],
+    metrics_processor: MetricsProcessor | None = None,
 ) -> BaseValidator:
     """Build a simple validator focused on correctness."""
     return Validator(
@@ -160,4 +163,5 @@ def build_validator(
         loss_fn=loss_fn,
         validation_context=validation_context,
         maybe_enable_amp=maybe_enable_amp,
+        metrics_processor=metrics_processor,
     )

torchtitan/tools/utils.py

@@ -134,6 +134,7 @@ class Color:
     white = "\033[37m"
     reset = "\033[39m"
     orange = "\033[38;2;180;60;0m"
+    turquoise = "\033[38;2;54;234;195m"
 
 
 @dataclass(frozen=True)

torchtitan/train.py

@@ -336,6 +336,7 @@ def __init__(self, job_config: JobConfig):
                 loss_fn=self.train_spec.build_loss_fn(job_config),
                 validation_context=self.train_context,
                 maybe_enable_amp=self.maybe_enable_amp,
+                metrics_processor=self.metrics_processor,
             )
 
         logger.info(
@@ -530,7 +531,7 @@ def train(self):
                     self.job_config.validation.enabled
                     and self.validator.should_validate(self.step)
                 ):
-                    self.validator.validate(self.model_parts)
+                    self.validator.validate(self.model_parts, self.step)
 
                 self.checkpointer.save(
                     self.step, last_step=(self.step == job_config.training.steps)