Always Pre-Split Microbatches for PP

Summary:
- Have PP training and validation dataloaders emit pipeline microbatches directly.
- Keep PP execution in the existing forward_backward_step/eval flow while passing pre-split args/kwargs/targets to torch.distributed.pipelining.
- 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 with the same structure-preserving pattern.

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: 51a5e12
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

@@ -8,7 +8,7 @@
 import time
 from collections.abc import Iterable
 from datetime import timedelta
-from typing import Any
+from typing import Any, cast
 
 import torch
 from torch.distributed.elastic.multiprocessing.errors import record
@@ -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
@@ -206,26 +210,41 @@ def forward_backward_step(
         model_parts = self.model_parts
         parallel_dims = self.parallel_dims
 
-        inputs, labels, extra_kwargs = self.post_dataloading_process(input_dict, labels)
-
         if parallel_dims.pp_enabled:
+            input_dict_mbs = cast(list[dict[str, torch.Tensor]], input_dict)
+            label_mbs = cast(list[torch.Tensor], labels)
+            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 mb_input_dict, mb_labels in zip(input_dict_mbs, label_mbs, strict=True):
+                inputs, mb_labels, extra_kwargs = self.post_dataloading_process(
+                    mb_input_dict, mb_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(mb_labels)
+
             # Pipeline Parallel forward / backward inside step() call
             with self.train_context():
-                targets, losses = (
-                    (labels, []) if self.pp_has_last_stage else (None, None)
-                )
+                losses = [] if self.pp_has_last_stage else None
                 if self.pp_has_first_stage:
                     self.pp_schedule.step(
-                        inputs,
-                        **extra_kwargs,
-                        target=targets,
+                        arg_mbs,
+                        kwargs=kwarg_mbs,
+                        target=target_mbs,
                         losses=losses,
+                        pre_split_args_kwargs=True,
                     )
                 else:
                     self.pp_schedule.step(
-                        **extra_kwargs,
-                        target=targets,
+                        kwargs=kwarg_mbs,
+                        target=target_mbs,
                         losses=losses,
+                        pre_split_args_kwargs=True,
                     )
 
             # accumulate losses across pipeline microbatches
@@ -238,6 +257,10 @@ def forward_backward_step(
                 else torch.tensor([-1.0], device=self.device)
             )
         else:
+            inputs, labels, extra_kwargs = self.post_dataloading_process(
+                input_dict, labels
+            )
+
             # Non-PP forward / backward
             with self.train_context():
                 assert len(model_parts) == 1
@@ -264,14 +287,28 @@ 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 = []
+        microbatches: list[tuple[Any, Any]] = []
         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()
+        for _ in range(self.gradient_accumulation_steps):
+            if parallel_dims.pp_enabled:
+                input_dict = []
+                labels = []
+                for _ in range(num_microbatches):
+                    mb_input_dict, mb_labels = next(data_iterator)
+                    local_valid_tokens += (mb_labels != IGNORE_INDEX).sum()
+                    input_dict.append(mb_input_dict)
+                    labels.append(mb_labels)
+            else:
+                input_dict, labels = next(data_iterator)
+                local_valid_tokens += (labels != IGNORE_INDEX).sum()
             microbatches.append((input_dict, labels))
 
         # All-reduce to get global token count across DP ranks
@@ -287,15 +324,26 @@ def train_step(
         # Process each microbatch: 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)
+            if parallel_dims.pp_enabled:
+                assert isinstance(input_dict, list)
+                assert isinstance(labels, list)
+                for mb_input_dict in input_dict:
+                    for k, v in mb_input_dict.items():
+                        if isinstance(v, torch.Tensor):
+                            mb_input_dict[k] = v.to(self.device)
+                labels = [mb_labels.to(self.device) for mb_labels in labels]
+            else:
+                assert isinstance(input_dict, dict)
+                assert isinstance(labels, torch.Tensor)
+                # 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,
+                labels=labels,  # pyrefly: ignore [bad-argument-type]
                 global_valid_tokens=global_valid_tokens,
             )
             accumulated_losses.append(loss.detach())

torchtitan/experiments/torchft/trainer.py

@@ -96,7 +96,11 @@ def __init__(self, config: Config):
             dp_rank=batch_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 parallel_dims.pp_enabled
+                else config.training.local_batch_size
+            ),
         )
 
         # build model (using meta init)
@@ -378,13 +382,28 @@ 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
         microbatches = []
         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()
+        for _ in range(self.gradient_accumulation_steps):
+            if parallel_dims.pp_enabled:
+                input_dict = []
+                labels = []
+                for _ in range(num_microbatches):
+                    mb_input_dict, mb_labels = next(data_iterator)
+                    local_valid_tokens += (mb_labels != IGNORE_INDEX).sum()
+                    input_dict.append(mb_input_dict)
+                    labels.append(mb_labels)
+            else:
+                input_dict, labels = next(data_iterator)
+                local_valid_tokens += (labels != IGNORE_INDEX).sum()
             microbatches.append((input_dict, labels))
 
         # All-reduce to get global token count across DP ranks
@@ -400,11 +419,22 @@ def train_step(
         # Process each microbatch: 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)
+            if parallel_dims.pp_enabled:
+                assert isinstance(input_dict, list)
+                assert isinstance(labels, list)
+                for mb_input_dict in input_dict:
+                    for k, v in mb_input_dict.items():
+                        if isinstance(v, torch.Tensor):
+                            mb_input_dict[k] = v.to(self.device)
+                labels = [mb_labels.to(self.device) for mb_labels in labels]
+            else:
+                assert isinstance(input_dict, dict)
+                assert isinstance(labels, torch.Tensor)
+                # 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,

torchtitan/models/common/decoder.py

@@ -142,17 +142,6 @@ def update_from_config(
                     "Weight tying is not supported with Pipeline Parallel."
                 )
 
-            if parallelism.pipeline_parallel_degree > 1 and any(
-                layer.attention is not None
-                and isinstance(layer.attention.inner_attention, VarlenAttention.Config)
-                for layer in self.layers
-            ):
-                raise ValueError(
-                    "Pipeline Parallel is not compatible with VarlenAttention. "
-                    "Use a FlexAttention backend (attn_backend='flex' or "
-                    "'flex_flash') for pipelined models."
-                )
-
             tp = parallelism.tensor_parallel_degree
             attention = self.first_attention
             if tp > 1 and attention is not None:

torchtitan/models/gpt_oss/config_registry.py

@@ -65,8 +65,6 @@ def gpt_oss_debugmodel() -> Trainer.Config:
 
 
 def gpt_oss_debugmodel_flex() -> Trainer.Config:
-    # FlexAttention variant. Pipeline Parallel is incompatible with
-    # VarlenAttention, so PP integration tests use this flex config.
     return _gpt_oss_debugmodel(attn_backend="flex")