Skip to content

Commit 3264cff

Browse files
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

File tree

    0 commit comments

    Comments
     (0)