Commit 3264cff
committed
Update on "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-poisoned]1 parent 4a550bf commit 3264cff
0 file changed
0 commit comments