Pipeline parallelism continued

[BlueCrescent]

What does this PR do?

• Issue: Dicts, as used in modalities as model/loss inputs and outputs, are not supported in the torch pp code.
  • Reason: Direct GPU to GPU communication
  • See: Support dict as input/output for pipeline parallelism pytorch/pytorch#159711
  • Workaround: Overloaded corresponding forward() and call() methods to allow direct tensor input and output.
• Issue: Loss accumulation (for logging) needs to consider pp world size instead of total world size
  • Solution:
    • For the averaged loss we only increment the batch count on the last stages of a PP schedule.
    • For the current batch loss we added the pp world size, computed from the device mesh, to the Trainer class.
• Issue: Due to a Bug in PyTorch we cannot evaluate before training begin.
  • Workaround: Disabled eval before training
  • Solution: Issue and PR opened for PyTorch 2.9
  • See: Pipeline parallelism does not correctly initialize backwards stages when evaluating before training. pytorch/pytorch#162822
• Issue: Normal distributed dataloader/sampler use would loose batches
  • Solution: Use (Resumable)DistributedMultiDimSamplerConfig
  • TODO: Implement DistributedMultiDimSamplerConfig for eval dataloader
  • TODO??: Remove data_parallel_key param from SamplerFactory.create_resumable_distributed_multi_dim_sampler()
    • Instead use BOTH ParallelismDegrees.DP_REPLICATE and ParallelismDegrees.DP_SHARD (multiplied)
• Issue: Gradient clipping needs to sync over all stages
  • Solution: Implemented same procedure using all_reduce() as in TorchTitan.
  • Torch Issue: Gradient Clipping in Pipeline Parallelism Schedules pytorch/pytorch#162638
  • TorchTitan implementation: https://github.com/pytorch/torchtitan/blob/b291ad662493b63d25b038a30a915082d3617baf/torchtitan/distributed/utils.py#L298
• Issue: Instead of calling forward() and backward() on the model and executing the loss_fct(), PP requires to call step() on the pipeline schedule.
  • Solution: Integrated scheduled pipeline into Trainer and Evaluator
• Issue: Want to run evaluation with torch.no_grad() and without doing a backwards() pass in the pipeline schedule step().
  • Solution: This is only supported from PyTorch 2.9 (current nightly) on. There the pipeline schedule gets a eval() method to be used instead of step().
  • ⚠️Warning⚠️: Consequently, we can only use PP in training with PyTorch 2.9 Nightly installed.
• Note: Weight tying is currently not supported and probably not possible with PP.
• Question: Do we need to consider something regarding model init so that all corresponding model copies are initialized the same?
  • I.e.: How can we assure that parallel stages are initialized the same?
  • Other way around: Will stages containing similar model parts be initialized the same?
  • Answer: Since we first build that PP stages and then the FSDP2 parallelization and then init, this is fine.
  • Follow up question: What happens, if we use full replica data parallelization?
• Issue: Want to compare non-PP forward pass with PP forward pass in unit test
  • Goal: Compare losses
  • Solution: Added a non-PP config used to compute a comparison loss
  • Issue: Need non-PP forward pass loss on all ranks that contain a last stage of the PP forward pass.
    • Solution: Run a normal FSDP2 forward pass on all ranks and ignore the output on non-last stage ranks.
    • Remark: Running the FSDP2 pass only on the last stage ranks led to hanging. This might be a known issue with FSDP2.
  • Issue: Need the model to effectively compute the same forward pass.
    • Solution:
      • In the test config, run model model initialization before sharding/staging (don't do this in a real training!).
      • Fix torch seed before each initialization.
      • Fix torch seed before generating each input sequence. Note that we thus have identical data parallel batches.
  • TODO: Numerical instability can still be observed and needs to be investigated further (in FSDP only runs as well).
• TP + PP
  • Issue: What is the correct sharding/staging order?
    • Solution: 1. PP 2. TP 3. FSDP
  • Issue: TP initialization modifies the layers of the model and runs into problems if those have been deleted for PP.
    • Solution: Adapted TP code to be able to handle missing/None layers.
  • Issue: PP + TP test hangs on first stages, first microbatches at the start of attention.
    • Solution: The test used sequence length 255. It seems, this not being divisible by two caused issues with TP's sequence parallelism.
    • ⚠️Warning⚠️: This needs to be kept in mind when using PP + TP.
• TODO: Wherever "dist.get_world_size()" is used, check if we should use number of data parallel ranks instead.
• TODO: MFU is 0 on dgx2, but also without PP
• TODO: Integrate and test other pipeline schedules.
• TODO: Check if checkpointing still works.

Checklist before submitting final PR

• My PR is minimal and addresses one issue in isolation
• I have merged the latest version of the target branch into this feature branch
• I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
• I have run a sample config for model training
• I have checked that all tests run through (python tests/tests.py)
• I have updated the internal changelog (CHANGELOG_DEV.md)

[le1nux]

I need to check one more test but wanted to share some comments already.

[le1nux]

LGTM :) I think we can merge this! 🥳
Amazing work and loved the in-depth testing. 🚀