Always Pre-Split Microbatches for PP

Summary:
- Have PP training and validation dataloaders emit pipeline microbatches directly.
- Group dataloader-produced microbatches inline into pre-split args/kwargs/targets for torch.distributed.pipelining schedule.step/eval.
- Remove the PP+Varlen guard and switch the GPT-OSS PP integration test back to the default varlen config.
- Update copied PP train paths in TorchFT and Forge to use the pre-split schedule API.

Motivation:
TorchTitan previously built a full local batch and relied on torch.distributed.pipelining to split it. That does not work for batch-dependent non-tensor metadata such as varlen attention metadata. Owning the split at the dataloader/trainer boundary lets TorchTitan generate block masks and varlen metadata per pipeline microbatch before calling the PP schedule.

Test Plan:
- python -m py_compile torchtitan/trainer.py torchtitan/components/validate.py torchtitan/models/common/decoder.py torchtitan/models/gpt_oss/config_registry.py tests/integration_tests/models.py torchtitan/experiments/torchft/trainer.py torchtitan/experiments/forge/engine.py torchtitan/experiments/forge/example_train.py
- pre-commit run --files torchtitan/trainer.py torchtitan/components/validate.py torchtitan/models/common/decoder.py torchtitan/models/gpt_oss/config_registry.py tests/integration_tests/models.py torchtitan/experiments/torchft/trainer.py torchtitan/experiments/forge/engine.py torchtitan/experiments/forge/example_train.py
- python -m tests.integration_tests.run_tests <output_dir> --test_suite models --test_name gpt_oss_pp+fsdp+ep+sacop --ngpu 8
- NGPU=2 LOG_RANK=0,1 ./run_train.sh --dump_folder <output_dir> --module gpt_oss --config gpt_oss_debugmodel --parallelism.pipeline_parallel_degree 2 --parallelism.pipeline_parallel_schedule 1F1B --training.steps 1 --validator.enable --validator.steps 1

ghstack-source-id: e45a042
Pull Request resolved: #3856

Author: sanketpurandare

tests/integration_tests/models.py

@@ -194,7 +194,7 @@ def build_model_tests_list() -> list[OverrideDefinitions]:
         OverrideDefinitions(
             [
                 [
-                    "--module gpt_oss --config gpt_oss_debugmodel_flex",
+                    "--module gpt_oss --config gpt_oss_debugmodel",
                     "--parallelism.data_parallel_shard_degree 4",
                     "--parallelism.pipeline_parallel_degree 2",
                     "--parallelism.pipeline_parallel_schedule Interleaved1F1B",

torchtitan/components/validate.py

@@ -227,28 +227,45 @@ def validate(
         accumulated_losses = []
         device_type = utils.device_type
         num_steps = 0
+        num_microbatches = (
+            self.local_batch_size // self.parallelism.pipeline_parallel_microbatch_size
+            if parallel_dims.pp_enabled
+            else 1
+        )
 
         validation_dataloader = self.dl_config.build(
             dp_world_size=self.dp_world_size,
             dp_rank=self.dp_rank,
             tokenizer=self.tokenizer,
             seq_len=self.seq_len,
-            local_batch_size=self.local_batch_size,
+            local_batch_size=(
+                self.parallelism.pipeline_parallel_microbatch_size
+                if parallel_dims.pp_enabled
+                else self.local_batch_size
+            ),
         )
 
-        for input_dict, labels in validation_dataloader:
+        validation_iterator = iter(validation_dataloader)
+        while True:
             # pyrefly: ignore [missing-attribute, unsupported-operation]
             if self.config.steps != -1 and num_steps >= self.config.steps:
                 break
 
-            self.metrics_processor.ntokens_since_last_log += labels.numel()
-            for k, v in input_dict.items():
-                input_dict[k] = v.to(device_type)
-            labels = labels.to(device_type)
-
-            # Count valid tokens for this batch
-            local_valid_tokens = torch.tensor(0, dtype=torch.int64, device=device_type)
-            local_valid_tokens += (labels != IGNORE_INDEX).sum()
+            try:
+                microbatches = []
+                local_valid_tokens = torch.tensor(
+                    0, dtype=torch.int64, device=device_type
+                )
+                for _ in range(num_microbatches):
+                    input_dict, labels = next(validation_iterator)
+                    self.metrics_processor.ntokens_since_last_log += labels.numel()
+                    for k, v in input_dict.items():
+                        input_dict[k] = v.to(device_type)
+                    labels = labels.to(device_type)
+                    local_valid_tokens += (labels != IGNORE_INDEX).sum()
+                    microbatches.append((input_dict, labels))
+            except StopIteration:
+                break
 
             # All-reduce token count across DP ranks to get global token count
             if parallel_dims.dp_enabled:
@@ -259,32 +276,43 @@ def validate(
             else:
                 global_valid_tokens = float(local_valid_tokens.item())
 
-            # Process data (extract inputs, handle attention masks, CP sharding)
-            inputs, labels, extra_kwargs = self.post_dataloading_process(
-                input_dict, labels, model_parts
-            )
-
             if parallel_dims.pp_enabled:
                 assert self.pp_schedule is not None
                 assert self.pp_has_first_stage is not None
                 assert self.pp_has_last_stage is not None
-                # Pipeline Parallel forward inside eval() call
-                with self.validation_context():
-                    targets, losses = (
-                        (labels, []) if self.pp_has_last_stage else (None, None)
+
+                arg_mbs: list[tuple[torch.Tensor, ...]] = []
+                kwarg_mbs: list[dict[str, Any]] = []
+                target_mbs: list[torch.Tensor] | None = (
+                    [] if self.pp_has_last_stage else None
+                )
+
+                for input_dict, labels in microbatches:
+                    inputs, labels, extra_kwargs = self.post_dataloading_process(
+                        input_dict, labels, model_parts
                     )
+                    if self.pp_has_first_stage:
+                        arg_mbs.append((inputs,))
+                    kwarg_mbs.append(extra_kwargs)
+                    if target_mbs is not None:
+                        target_mbs.append(labels)
+
+                with self.validation_context():
+                    losses = [] if self.pp_has_last_stage else None
                     if self.pp_has_first_stage:
                         self.pp_schedule.eval(
-                            inputs,
-                            **extra_kwargs,
-                            target=targets,
+                            arg_mbs,
+                            kwargs=kwarg_mbs,
+                            target=target_mbs,
                             losses=losses,
+                            pre_split_args_kwargs=True,
                         )
                     else:
                         self.pp_schedule.eval(
-                            **extra_kwargs,
-                            target=targets,
+                            kwargs=kwarg_mbs,
+                            target=target_mbs,
                             losses=losses,
+                            pre_split_args_kwargs=True,
                         )
 
                 # accumulate losses across pipeline microbatches
@@ -296,6 +324,12 @@ def validate(
                 else:
                     loss_sum = torch.tensor([-1.0], device=device_type)
             else:
+                assert len(microbatches) == 1
+                input_dict, labels = microbatches[0]
+                # Process data (extract inputs, handle attention masks, CP sharding)
+                inputs, labels, extra_kwargs = self.post_dataloading_process(
+                    input_dict, labels, model_parts
+                )
                 with self.validation_context():
                     assert len(model_parts) == 1
                     predictions = model_parts[0](inputs, **extra_kwargs)

torchtitan/experiments/forge/example_train.py

@@ -59,7 +59,11 @@ def __init__(self, config: TitanTrainer.Config):
             dp_rank=self.dp_rank,
             tokenizer=self.tokenizer,
             seq_len=config.training.seq_len,
-            local_batch_size=config.training.local_batch_size,
+            local_batch_size=(
+                config.parallelism.pipeline_parallel_microbatch_size
+                if self.parallel_dims.pp_enabled
+                else config.training.local_batch_size
+            ),
         )
 
         model_args = self.model_config
@@ -205,53 +209,24 @@ def forward_backward_step(
     ) -> torch.Tensor:
         model_parts = self.model_parts
         parallel_dims = self.parallel_dims
+        assert not parallel_dims.pp_enabled
 
         inputs, labels, extra_kwargs = self.post_dataloading_process(input_dict, labels)
 
-        if parallel_dims.pp_enabled:
-            # Pipeline Parallel forward / backward inside step() call
-            with self.train_context():
-                targets, losses = (
-                    (labels, []) if self.pp_has_last_stage else (None, None)
-                )
-                if self.pp_has_first_stage:
-                    self.pp_schedule.step(
-                        inputs,
-                        **extra_kwargs,
-                        target=targets,
-                        losses=losses,
-                    )
-                else:
-                    self.pp_schedule.step(
-                        **extra_kwargs,
-                        target=targets,
-                        losses=losses,
-                    )
+        # Non-PP forward / backward
+        with self.train_context():
+            assert len(model_parts) == 1
+            pred = model_parts[0](inputs, **extra_kwargs)
+            # Compute loss sum (reduction='sum')
+            loss_sum, _ = self.loss_fn(pred, labels)
 
-            # accumulate losses across pipeline microbatches
-            # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
-            loss = (
-                # Rescale PP loss to be "local loss sum / global valid tokens"
-                # because each microbatch could have different number of valid tokens
-                (torch.sum(torch.stack(losses)) / global_valid_tokens).to(self.device)
-                if self.pp_has_last_stage
-                else torch.tensor([-1.0], device=self.device)
-            )
-        else:
-            # Non-PP forward / backward
-            with self.train_context():
-                assert len(model_parts) == 1
-                pred = model_parts[0](inputs, **extra_kwargs)
-                # Compute loss sum (reduction='sum')
-                loss_sum, _ = self.loss_fn(pred, labels)
+            # Scale the loss by the inverse of the total weight denominator before backward
+            # This ensures gradients are properly normalized across all microbatches
+            loss = loss_sum / global_valid_tokens
 
-                # Scale the loss by the inverse of the total weight denominator before backward
-                # This ensures gradients are properly normalized across all microbatches
-                loss = loss_sum / global_valid_tokens
-
-                # need to free pred before bwd to avoid peaking memory
-                del pred
-                loss.backward()
+            # need to free pred before bwd to avoid peaking memory
+            del pred
+            loss.backward()
 
         # The returned loss here is local SUM loss / global_valid_tokens
         return loss
@@ -264,15 +239,23 @@ def train_step(
         # Keep these variables local to shorten the code as these are
         # the major variables that are used in the training loop.
         parallel_dims = self.parallel_dims
+        num_microbatches = (
+            self.config.training.local_batch_size
+            // self.config.parallelism.pipeline_parallel_microbatch_size
+            if parallel_dims.pp_enabled
+            else 1
+        )
 
         # Collect all microbatches on CPU and count total valid tokens
-        # Here we assume the inputs/labels are on GPU
-        microbatches = []
+        batches = []
         local_valid_tokens = torch.tensor(0, dtype=torch.int64)
-        for _microbatch in range(self.gradient_accumulation_steps):
-            input_dict, labels = next(data_iterator)
-            local_valid_tokens += (labels != IGNORE_INDEX).sum()
-            microbatches.append((input_dict, labels))
+        for _ in range(self.gradient_accumulation_steps):
+            step_microbatches = []
+            for _ in range(num_microbatches):
+                input_dict, labels = next(data_iterator)
+                local_valid_tokens += (labels != IGNORE_INDEX).sum()
+                step_microbatches.append((input_dict, labels))
+            batches.append(step_microbatches)
 
         # All-reduce to get global token count across DP ranks
         # Move to GPU for distributed communication
@@ -284,20 +267,72 @@ def train_step(
         else:
             global_valid_tokens = float(local_valid_tokens.item())
 
-        # Process each microbatch: move to GPU, forward/backward, then free
+        # Process each batch: move to GPU, forward/backward, then free
         accumulated_losses = []
-        for input_dict, labels in microbatches:
-            # Move tensors to GPU
-            for k, v in input_dict.items():
-                if isinstance(v, torch.Tensor):
-                    input_dict[k] = v.to(self.device)
-            labels = labels.to(self.device)
-
-            loss = self.forward_backward_step(
-                input_dict=input_dict,
-                labels=labels,
-                global_valid_tokens=global_valid_tokens,
-            )
+        for step_microbatches in batches:
+            if parallel_dims.pp_enabled:
+                arg_mbs: list[tuple[torch.Tensor, ...]] = []
+                kwarg_mbs: list[dict[str, Any]] = []
+                target_mbs: list[torch.Tensor] | None = (
+                    [] if self.pp_has_last_stage else None
+                )
+
+                for input_dict, labels in step_microbatches:
+                    for k, v in input_dict.items():
+                        if isinstance(v, torch.Tensor):
+                            input_dict[k] = v.to(self.device)
+                    labels = labels.to(self.device)
+
+                    inputs, labels, extra_kwargs = self.post_dataloading_process(
+                        input_dict, labels
+                    )
+                    if self.pp_has_first_stage:
+                        arg_mbs.append((inputs,))
+                    kwarg_mbs.append(extra_kwargs)
+                    if target_mbs is not None:
+                        target_mbs.append(labels)
+
+                with self.train_context():
+                    losses = [] if self.pp_has_last_stage else None
+                    if self.pp_has_first_stage:
+                        self.pp_schedule.step(
+                            arg_mbs,
+                            kwargs=kwarg_mbs,
+                            target=target_mbs,
+                            losses=losses,
+                            pre_split_args_kwargs=True,
+                        )
+                    else:
+                        self.pp_schedule.step(
+                            kwargs=kwarg_mbs,
+                            target=target_mbs,
+                            losses=losses,
+                            pre_split_args_kwargs=True,
+                        )
+
+                # accumulate losses across pipeline microbatches
+                # TODO: PP+FSDP unexpectedly puts the loss back to the CPU
+                if self.pp_has_last_stage:
+                    assert losses is not None
+                    # Rescale PP loss to be "local loss sum / global valid tokens"
+                    # because each microbatch could have different number of valid tokens
+                    loss = (torch.sum(torch.stack(losses)) / global_valid_tokens).to(
+                        self.device
+                    )
+                else:
+                    loss = torch.tensor([-1.0], device=self.device)
+            else:
+                assert len(step_microbatches) == 1
+                input_dict, labels = step_microbatches[0]
+                for k, v in input_dict.items():
+                    if isinstance(v, torch.Tensor):
+                        input_dict[k] = v.to(self.device)
+                labels = labels.to(self.device)
+                loss = self.forward_backward_step(
+                    input_dict=input_dict,
+                    labels=labels,
+                    global_valid_tokens=global_valid_tokens,
+                )
             accumulated_losses.append(loss.detach())
 
         grad_norm = dist_utils.clip_grad_norm_(