[Draft] Support pp in XTuner - #2036
Open
HAOCHENYE wants to merge 7 commits into
Open
Conversation
…rallel Split the single-seq_ctx MoE forward into _prepare_forward / _embed_step / _layers_step / _head_step, with _forward as a thin orchestrator. _layers_step takes an optional layer_indices range so a later pipeline stage can run only a subset of decoder layers. Behavior is unchanged for the non-split path; the intra-layer micro-batch path (_micro_batch_forward) is untouched. Add a regression test asserting the orchestrator equals an explicit manual composition of the four helpers and is run-to-run deterministic with both auxiliary losses active.
Introduce PipelineParallelConfig (pp_size / schedule / layer_split), MoE.split_for_pipeline for trimming a model down to one stage's modules, and MoE.pipeline_forward for staged execution over the PR's forward helpers. Verified by a single-process split equivalence test that feeds each stage's output into the next and compares the final loss against the unsplit forward.
Thread a data-parallel reduce group through the loss/aux collectives so they no longer assume WORLD (which deadlocks when only the last stage computes loss), add PPEngine driving torch.distributed.pipelining PipelineStage + Schedule1F1B, materialize AsyncCollectiveTensor outputs of the all2all dispatcher before they are saved for backward, and handle mixed plain/DTensor parameters (ep>1) in grad clipping and optimizer construction.
…allel A pipeline-split model holds only its stage's parameters, which broke HF save/load: - split_for_pipeline marks the model with _pipeline_split so the checkpoint paths know it is a subset of the full model. - Save (replicated 'others'): the others-save rank is now the first rank of each pipeline stage (via dp_group) instead of global rank 0, so every stage's embed/norm/lm_head/layer params get written; the weight maps are still gathered and merged into one index covering the full model. - Save (fused experts): distribute the per-expert keys across the load_spec's own group (the stage's ep group; a single owner when ep_size==1) instead of across the whole world, otherwise ranks that do not own a layer silently drop their slice of that layer's experts. - Load: a pipeline stage legitimately sees the other stages' keys in the checkpoint, so extra (unexpected) keys are no longer an error when split; missing keys still are. - PPEngine.save_hf delegates to the model. Non-pipeline behavior is unchanged (all branches guard on _pipeline_split / a set dp_group). Verified end-to-end: pp=2/ep=1 and pp=2/ep=2 save_hf produce a complete (1575/1575 keys) checkpoint that reloads into a full unsplit model with no missing or unloaded keys.
Select PPEngine in Trainer.build_engine when pp_size > 1 (and forward pp_cfg through Trainer.from_config), build the (pp, dp, sp, tp) data mesh with data replicated along pp, return TrainStepInfo from PPEngine.train_step for the trainer loop, correct throughput accounting and align the profiling window across ranks under pipeline parallel, and add distributed engine tests (train + checkpoint round-trip for pp2/ep2 and pp4/ep1).
Activation recompute was applied only inside MoE.fully_shard, so the pipeline engine (which never calls fully_shard) ran with no recomputation, inflating activation memory. Extract the layer-wrapping into MoE.apply_activation_checkpointing(recompute_ratio) and call it from both fully_shard and PPEngine.build_model (after split_for_pipeline; only the stage's owned layers are wrapped, keyed by global layer index so the recompute policy matches the unsplit model). _should_recompute now takes recompute_ratio explicitly instead of reading fsdp_config, decoupling recompute from FSDP. PPEngine gains a recompute_ratio argument, threaded from fsdp_cfg.recompute_ratio in the trainer. FSDP behavior is unchanged (same ratio, same wrapping, just hoisted above the shard loop). Verified under pipeline parallel: with recompute_ratio=1.0 the stage's layers are CheckpointWrapper-wrapped and a train step runs. Note: torch.compile already applies under pipeline parallel — for MoE it is method-level (@maybe_compile) enabled at model build, independent of fully_shard, so PPEngine picks it up via model_cfg.build() with no extra wiring.
Add PipelineParallelConfig.num_virtual_stages: layers are split into pp_size * num_virtual_stages chunks assigned round-robin, each owned virtual stage gets its own PipelineStage, and ScheduleInterleaved1F1B drives them. layer_split is honored per virtual stage. MoE aux loss is rejected under interleaving with a clear error because its global tokens-per-expert accumulator mixes microbatches across the two tail-rank virtual stages.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TODO