Support multimodal payloads in Dynamic-CP sequence packing scheduler
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
Tag @NVIDIA/mcore-oncall
to get oncall's attention to this issue.
Hi Megatron team,
I have implemented plain-text Dynamic Context Parallelism integration in MS-Swift using the Megatron-LM **dev** branch. For packed/padding-free text training, Dynamic-CP works and we observed speedups in several training configurations.
I then tried to extend this integration to multimodal training, but found that the current Dynamic-CP sequence packing path appears to assume text-only THD fields.
Currently, `data_schedule.py` and `data_schedule_utils.py` mainly process and broadcast these fields:
- `tokens`
- `labels`
- `loss_mask`
- `position_ids`
- sequence metadata such as `cu_seqlens`, `cu_seqlens_padded`, `max_seqlen`, `local_cp_size`
For multimodal training, frameworks also need to preserve and reroute payloads such as:
- `pixel_values`
- `pixel_values_videos`
- `image_grid_thw`
- `video_grid_thw`
- other modality metadata, e.g. video timestamps or per-media metadata
These fields do not have the same shape semantics as token-aligned THD tensors. For example, `pixel_values` is media-aligned rather than sequence-token-aligned, while `image_grid_thw/video_grid_thw` describe media layout. Therefore, simply adding these fields to the current batch dict is not enough; the current unpack/reroute/pack logic would need to understand different payload categories.
In MS-Swift, this means that if we want to support multimodal Dynamic-CP without upstream support, we likely need to reimplement or heavily wrap the reroute and packing logic on the Swift side, while only reusing the scheduling/grouping algorithm from Megatron.
Would Megatron consider supporting multimodal Dynamic-CP upstream?
A possible direction may be to split the current Dynamic-CP data path into two layers:
1. A scheduling/planning layer that only decides:
- sample grouping
- target DPxCP ranks
- local CP size per packed microbatch
2. A payload movement/packing layer that understands field semantics:
- token-aligned fields, such as `tokens`, `labels`, `loss_mask`, `position_ids`
- sample-level metadata
- media-aligned tensors, such as `pixel_values` / `pixel_values_videos`
- media metadata, such as `image_grid_thw` / `video_grid_thw`
- fields that should be broadcast/copied rather than token-sliced
Megatron could provide default payload policies for the current text-only path and extend them for multimodal/VLM batches. Downstream frameworks could also reuse the scheduling result without duplicating the Dynamic-CP grouping algorithm.
Is multimodal Dynamic-CP support on the roadmap? If not, would the Megatron team be open to an API refactor that makes the scheduler extensible for non-text payloads?
Thanks!
Contributor guide
Assessment
This issue has not been assessed yet.