Fix virtual pipeline support

[cjwystz]

PR Category

Train

PR Types

New Features

PR Description

Modifications

• Add vp_stage parameter to get_qwen35_language_model_spec() and implement full parameter transmission
• Introduce vp_stage parameter into get_qwen35_mtp_block_spec() to complete parallel adaptation for MTP module
• Pass vp_stage to experimental attention variant structure function to ensure correct layer offset calculation
• Parse and inject vp_stage into model instances within model_provider
• Add judgment logic for intermediate VPP chunks in batch acquisition to avoid excessive data iterator reading
• Extract virtual pipeline stage identifier in forward_step and pass it down to data interface
• Remove VPP restrictions for vision encoder and implement dynamic layer number calculation under VPP mode
• Supplement vp_stage argument for train_valid_test_dataloaders_provider to unify global calling interfaces

Background

Virtual Pipeline Parallelism (VPP) is an efficient mainstream parallel optimization in Megatron-LM. It splits each physical pipeline stage into multiple virtual stages and adopts interleaved scheduling to reduce pipeline bubbles, thus significantly improving overall training throughput in distributed large model scenarios.

The original Qwen35 training framework lacked support for the vp_stage parameter, leading to multiple adaptation defects. The language model failed to compute accurate layer offsets under VPP and triggered assertion errors. The MTP module could not determine the correct deployment rank without virtual stage identification. Intermediate virtual chunks kept consuming data iterators and caused data disorder. Besides, the vision encoder was incompatible with VPP scheduling, making full VPP acceleration unavailable for the whole multimodal model.

Implementation Details

1. Adaptation in layer_specs.py
   Add optional parameter vp_stage=None to language model and MTP specification functions. Transmit the parameter down to transformer block definition functions to calculate precise layer offsets when VPP is enabled. Meanwhile, judge whether to initialize MTP branches according to current virtual pipeline stage.

2. Vision Encoder Optimization in transformer_config.py
   Delete original assertions that disabled VPP for vision modules. Add logic to obtain virtual pipeline parallel world size and dynamically adjust total layer numbers of vision encoder, enabling unified VPP support for both text and vision branches.

3. Training Pipeline Upgrade in train_qwen35.py
   Extract vp_stage from keyword arguments and pass it into models during initialization. Add identification rules for intermediate VPP chunks, which skip loading new data while retaining position id and attention mask calculation required by RoPE. Deliver virtual stage information in forward process and skip redundant vision data assembly. Update dataloader provider function signature to be fully compatible with VPP scheduler calling standards.

Tests

• Verified stable training on 16-GPU cluster with no runtime exceptions
• Verified correct layer allocation with 10 layers configured per virtual pipeline stage
• Confirmed MTP blocks are only constructed on the last virtual chunk of final pipeline stage
• Fixed excessive data reading issue and ensured consistent batch data alignment across all chunks
• Verified accurate layer calculation and normal operation of vision encoder under VPP mode
• Training throughput meets expected performance indicators

Impact Scope

• Compatibility: vp_stage is set to None by default. All original logic remains unchanged when VPP is disabled.
• Affected Modules: Qwen35 language backbone, MTP multi-token prediction module, multimodal vision encoder

Contributors

• Yinxi Zhao (zhaoyinxi0@gmail.com)
• Jiawei Chen (2451427796@qq.com)

[CLAassistant]

All committers have signed the CLA.