NVIDIA / NVIDIA/Megatron-LM

Support multimodal payloads in Dynamic-CP sequence packing scheduler

Open
#5,683 4 comments 0 reactions 1 assignee Assigned to @ilml View on GitHub
community-request enhancement waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.