corrected tp and cp integration in validation, removed some repetitive integration tests

Author: wesleytruong

tests/integration_tests.py

@@ -513,45 +513,19 @@ def build_test_list():
             [
                 [
                     "--validation.enabled",
-                    "--validation.dataset c4_validation",
+                    "--validation.dataset c4_test",
                 ],
             ],
             "Validation test no parallelism",
             "validation_no_parallel",
             ngpu=1,
         ),
-        OverrideDefinitions(
-            [
-                [
-                    "--validation.enabled",
-                    "--validation.dataset c4_validation",
-                    "--parallelism.data_parallel_shard_degree=1",
-                    "--parallelism.data_parallel_replicate_degree=4",
-                ]
-            ],
-            "Validation test with DP",
-            "validation_dp",
-            ngpu=4,
-        ),
-        OverrideDefinitions(
-            [
-                [
-                    "--validation.enabled",
-                    "--validation.dataset c4_validation",
-                    "--parallelism.data_parallel_shard_degree=2",
-                    "--parallelism.data_parallel_replicate_degree=2",
-                ]
-            ],
-            "Validation test with FSDP",
-            "validation_fsdp",
-            ngpu=4,
-        ),
         OverrideDefinitions(
             [
                 [
                     "--checkpoint.enable_checkpoint",
                     "--validation.enabled",
-                    "--validation.dataset c4_validation",
+                    "--validation.dataset c4_test",
                     "--parallelism.data_parallel_shard_degree=2",
                     "--parallelism.data_parallel_replicate_degree=2",
                 ]
@@ -560,37 +534,6 @@ def build_test_list():
             "validation_fsdp_checkpoint",
             ngpu=4,
         ),
-        OverrideDefinitions(
-            [
-                [
-                    "--validation.enabled",
-                    "--validation.dataset c4_validation",
-                    "--parallelism.data_parallel_shard_degree=2",
-                    "--parallelism.data_parallel_replicate_degree=1",
-                    "--parallelism.tensor_parallel_degree=2",
-                    "--parallelism.context_parallel_degree=2",
-                ]
-            ],
-            "Validation test with FSDP, TP, CP",
-            "validation_fsdp_tp_cp",
-            ngpu=8,
-        ),
-        OverrideDefinitions(
-            [
-                [
-                    "--checkpoint.enable_checkpoint",
-                    "--validation.enabled",
-                    "--validation.dataset c4_validation",
-                    "--parallelism.data_parallel_shard_degree=2",
-                    "--parallelism.data_parallel_replicate_degree=1",
-                    "--parallelism.tensor_parallel_degree=2",
-                    "--parallelism.context_parallel_degree=2",
-                ]
-            ],
-            "Validation checkpoint test with FSDP, TP, CP",
-            "validation_fsdp_tp_cp_checkpoint",
-            ngpu=8,
-        ),
     ]
     return integration_tests_flavors
 

torchtitan/components/validate.py

@@ -4,10 +4,11 @@
 # 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 torch.distributed.tensor import DTensor
 from torchtitan.components.dataloader import BaseDataLoader
 from torchtitan.components.loss import LossFunction
@@ -52,6 +53,8 @@ def __init__(
         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
@@ -63,6 +66,8 @@ def __init__(
             dp_rank=dp_rank,
             tokenizer=tokenizer,
         )
+        self.validation_context = validation_context
+        self.maybe_enable_amp = maybe_enable_amp
 
     @torch.no_grad()
     def validate(
@@ -76,44 +81,52 @@ def validate(
 
         accumulated_losses = []
         device_type = utils.device_type
-        num_val_steps = 0
+        num_steps = 0
 
         for input_dict, labels in self.validation_dataloader:
             if (
                 self.job_config.validation.steps != -1
-                and num_val_steps >= self.job_config.validation.steps
+                and num_steps >= self.job_config.validation.steps
             ):
                 break
 
             for k, v in input_dict.items():
                 input_dict[k] = v.to(device_type)
-            labels = labels.to(device_type)
-
             inputs = input_dict["input"]
-            predictions = model(inputs)
+            labels = labels.to(device_type)
 
-            if self.parallel_dims.loss_parallel_enabled:
-                if isinstance(predictions, torch.Tensor) and not isinstance(
-                    predictions, DTensor
-                ):
-                    predictions = DTensor.from_local(predictions, self.world_mesh["tp"])
-                if isinstance(labels, torch.Tensor) and not isinstance(labels, DTensor):
-                    labels = DTensor.from_local(labels, self.world_mesh["tp"])
-            loss = self.loss_fn(predictions, labels)
+            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_val_steps += 1
+            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_val_steps} batches"
+            f"Validation completed. Average loss: {global_avg_loss:.4f} over {num_steps} batches"
         )
 
         # Reshard after run forward pass
@@ -125,8 +138,6 @@ def validate(
         # Set model back to train mode
         model.train()
 
-        return {"validation_loss": global_avg_loss}
-
 
 def build_validator(
     job_config: JobConfig,
@@ -136,6 +147,8 @@ def build_validator(
     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(
@@ -146,4 +159,6 @@ def build_validator(
         parallel_dims=parallel_dims,
         world_mesh=world_mesh,
         loss_fn=loss_fn,
+        validation_context=validation_context,
+        maybe_enable_amp=maybe_enable_amp,
     )

torchtitan/train.py

@@ -219,7 +219,7 @@ def __init__(self, job_config: JobConfig):
         if parallel_dims.pp_enabled:
             if not self.train_spec.pipelining_fn:
                 raise RuntimeError(
-                    f"pipeline parallel is enabled but {self.train_spec.name} "
+                    f"Pipeline parallel is enabled but {self.train_spec.name} "
                     f"does not support pipelining"
                 )
 
@@ -336,7 +336,9 @@ def __init__(self, job_config: JobConfig):
                 tokenizer=tokenizer,
                 parallel_dims=parallel_dims,
                 world_mesh=world_mesh,
-                loss_fn=self.loss_fn,
+                loss_fn=self.train_spec.build_loss_fn(job_config),
+                validation_context=self.train_context,
+                maybe_enable_amp=self.maybe_enable_amp,
             )
 
         logger.info(
@@ -525,7 +527,7 @@ def train(self):
                     self.job_config.validation.enabled
                     and self.validator.should_validate(self.step)
                 ):
-                    validation_metrics = self.validator.validate(self.model_parts)
+                    self.validator.validate(self.model_parts)
 
                 self.checkpointer.save(
                     self.step, last_step=(self.step == job_config.training.steps)