Always Pre-Split Microbatches for PP

[sanketpurandare]

Stack from ghstack (oldest at bottom):

• -> Always Pre-Split Microbatches for PP #3856

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

[fegin]

Let's move this to a method pp_forward_backward_step and this block becomes:

if parallel_dims.pp_enabled:
     return self.pp_forward_backward_step(...)

non_pp_logic

Reason: the if/else blocks now share nothing except for local alias. Before this PR, the two at least share the inputs, labels, extra_kwargs = self.post_dataloading_process(input_dict, labels). And the pp logic becomes more complicated now and justifies a method for readability.

[fegin]

We cannot use PP with gradient accumulation after this PR, right? If so, do we error out somewhere if users try to use them together?

[sanketpurandare]

Using pp with gradient accumulation doesn't make sense. PP is already doing gradient accumulation. So we can disable gradient accumulation for pp.

[tianyu-l]

I don't think that's true. See previous discussions in #1733

Conclusion is that we should continue to support gradient accumulation + PP.

[tianyu-l]

The main change / challenge in this PR is to properly differentiate PP microbatching vs. gradient accumulation microbatching, from UX / config perspective.

E.g. a simple way is to let user specify "local microbatch size"

• without PP it's the local_batch_size today, which is used together with global_batch_size to determine number of grad accum steps
• with PP, it's the local microbatch size, whose product with num PP microbatches * num grad accum steps would have to equal global_batch_size.

[fegin]

Why do we need pyrefly ignore now? Let's try to avoid using pyrefly ignore. Let's use cast if we really cannot figure out why.

[tianyu-l]

I don't think that's true. See previous discussions in #1733

Conclusion is that we should continue to support gradient accumulation + PP.

[tianyu-l]

This is computed twice in constructor and here

[tianyu-l]

it feels to me this can be improved -- two pieces of code seem almost identical

[tianyu-l]

The main change / challenge in this PR is to properly differentiate PP microbatching vs. gradient accumulation microbatching, from UX / config perspective.

E.g. a simple way is to let user specify "local microbatch size"

• without PP it's the local_batch_size today, which is used together with global_batch_size to determine number of grad accum steps
• with PP, it's the local microbatch size, whose product with num PP microbatches * num grad accum steps would have to equal global_batch_size.