Skip to content

[feature] intra-microbatch reordering for MegatronMIMO (+ sequence packing, scalable DP) #4609

Open
Open
[feature] intra-microbatch reordering for MegatronMIMO (+ sequence packing, scalable DP)#4609
Assignees
liding-nv
Labels
area:dataDataset builders, preprocessing, and samplerscommunity-requestfeatureNew capabilities, enhancements, or enablement workx-naver

Description

@sailor1493
sailor1493
opened on Jul 1, 2026

User problem

Non-colocated MegatronMIMO training (e.g. Qwen3.5-VL: separate vision encoder + language model wired by BridgeCommunicator) has a per-step data-parallel straggler: per-micro-batch multimodal load
(vision patch count) is uneven across samples, so a DP rank that draws a heavy-image shard stalls the whole group every step. Throughput is bounded by the slowest rank.

Two supporting gaps make this worse / block the fix:

  1. Sequence packing is unused. MegatronMIMO on Megatron-Core supports packed sequences, but in Megatron Bridge, there is no logic to actually feed a packed language shard, so padding is computed.
  2. Data parallelism is not scalable (vision DP = 1 only). Today every DP worker reads the full global micro-batch and slices its own part locally in the forward step → CPU/host-IO overhead that scales with DP, and vision DP > 1 is effectively unsupported. When the vision data parallel degress scale linearly, the total i/o volume across the cluster will scale quadratically.

Desired outcome

Main — Intra-microbatch Reordering (cf. DistTrain §5.2, https://arxiv.org/abs/2408.04275):
rebalance each micro-batch's vision load across the module DP group, so per-rank vision
load is even and the straggler tail is removed. Works with heterogeneous DP (vision_dp != language_dp)
via canonical n_groups pairing (vision replica r ↔ language replica r) and a variable number of
images per sample (0 / 1 / N).

Auxiliary 1 — Sequence Packing: add the logic to pack a language DP shard's real tokens into a
single [1, T] THD sequence so the LM skips padding compute (MCore already accepts packed input).

Auxiliary 2 — Scalable Data Parallelism: let each rank's sampler read only its disjoint 1/dp
shard instead of every rank reading the full batch and slicing locally — reducing the per-rank CPU/IO
overhead and enabling vision DP > 1. This is the read-sharding base the reorder rebalances on top of.

All three are off by default and gated by MegatronMIMOFeatureConfig. Expected effect (16-GPU
validation): median throughput dominated by the read-sharding win, plus a large straggler-tail
reduction; same total work / same images processed.

The name/implementation of the configuration class can be a subject of discussion.

Alternatives considered

  • Status quo — full-batch read + forward-time local slice (slice_batch_for_megatron_mimo):
    wastes host IO, supports vision DP = 1 only, cannot rebalance per-rank cost.

Affected area

area:data (primary), area:training

Urgency / use case

Important but not blocking.

Use case: non-colocated MegatronMIMO VLM finetuning (Qwen3.5-VL + receipt/doc datasets) where vision load is the per-step straggler.

Extra context

Activity

  1. sailor1493
    added
    featureNew capabilities, enhancements, or enablement work
    needs-triageNew item needs classification and ownership
    on Jul 1, 2026
  2. github-actions
    added
    community-request
    on Jul 1, 2026
  3. yaoyu-33
    added
    area:dataDataset builders, preprocessing, and samplers
    waiting-on-maintainersWaiting on maintainers to respond
    and removed
    needs-triageNew item needs classification and ownership
    on Jul 1, 2026
  4. cuichenx
    assigned
    liding-nv
    on Jul 2, 2026

  5. liding-nv commented on Jul 2, 2026

    @liding-nv
    liding-nv
    on Jul 2, 2026
    Contributor

    @sailor1493 thank you for the contribution. i will review the PR and we can take it from there

  6. liding-nv commented on Jul 9, 2026

    @liding-nv
    liding-nv
    on Jul 9, 2026
    Contributor

    hello @sailor1493 following up after checking the PR and the issue. This is a good feature and we definitely want to have them. I have several high-level points I'd like to align on first

    • the PR actually accomplished several independently useful and independently risky capabilities. Could we split it into a few small focused PR, such as (1) seq packing support for mimo (2) scalable data loading (3) intra-mb vision token balancing. this should make the PR much easier to review and validate
    • re: "Data parallelism is not scalable (vision DP = 1 only)", encoder DP>1 is currently supported. The actual gap appears to be data-loading scalability. Each rank loads the full microbatch before local slicing. Did you encouter a separate encoder-DP failure?
    • Is this intended to be a generic megatron mimo feature or specifically a qwen3.5 implementation? The config and forward step changes are shared, and data-loading and reordering wiring appears largely qwen specific.

  7. sailor1493 commented on Jul 10, 2026

    @sailor1493
    sailor1493
    on Jul 10, 2026
    ContributorAuthor

    Good day! @liding-nv

    It seems that our views are largely aligned. However, since today is my last day at NAVER Cloud—I’m moving on to a new opportunity :)—I may not be able to contribute further to this issue. Yoonsik or Gayeon will make further comments.

    Yoonsik (@terryyoo), Gayeon (@going-song), and other contributors will continue working on it. Thank you for your understanding.

    Wishing both NVIDIA and NAVER Cloud the very best!

  8. huvunvidia commented on Jul 10, 2026

    @huvunvidia
    huvunvidia
    on Jul 10, 2026
    Contributor

    Hi @sailor1493 ,
    Thank you for letting us know.
    Since I am not able to tag Yoonsik and Gayeon, I wonder if you can engage them into the conversation to move forward with @liding-nv 's discussion above?

  9. huvunvidia
    added
    waiting-on-customerWaiting on the original author to respond
    and removed
    waiting-on-maintainersWaiting on maintainers to respond
    on Jul 10, 2026

  10. terryoo commented on Jul 12, 2026

    @terryoo
    terryoo
    on Jul 12, 2026
    Contributor

    Hi @liding-nv , this is Yoonsik from NAVER Cloud. I'll be taking over this discussion going forward along with Gayeon.

  11. terryoo commented on Jul 14, 2026

    @terryoo
    terryoo
    on Jul 14, 2026
    Contributor

    Good day! @liding-nv

    • Agreed that the three pieces are independently useful and should be reviewed separately. One clarification (apologies, as a non-native speaker I want to make sure I read this right): would you like us to split and resubmit, or is this something the maintainers would handle? If it's on our side, happy to do it. We'd follow your suggested split: (1) sequence packing for MegatronMIMO, (2) scalable data loading, (3) intra-mb vision token balancing on top of (2).
    • No, we did not encounter a separate encoder-DP failure. The gap we meant is data-loading: as you'd expect, having every rank read the full micro-batch and slice locally becomes the dominant cost as DP grows. In our measurements, fixing this gave a larger throughput win than sequence packing itself.
    • This is intended as a generic MegatronMIMO feature. The config and forward-step changes are model-agnostic. Some data-loading and reordering wiring may look Qwen-specific because Qwen3.5-VL was the first requested target, but we're already running it internally on a non-Qwen model without issues.

  12. 2 remaining items

  13. huvunvidia
    added
    waiting-on-maintainersWaiting on maintainers to respond
    on Jul 15, 2026
  14. sajadn
    added
    waiting-on-customerWaiting on the original author to respond
    and removed
    waiting-on-maintainersWaiting on maintainers to respond
    on Jul 17, 2026
  15. going-song
    added 2 commits that reference this issue on Jul 21, 2026

    feat(training): add MegatronMIMO in-batch sequence packing

    21d53ee

    feat(training): add MegatronMIMO in-batch sequence packing

    885070d

  16. going-song commented on Jul 28, 2026

    @going-song
    going-song
    on Jul 28, 2026
    Contributor

    Hi @liding-nv — first split is up (#5131) : MegatronMIMO in-batch sequence packing

    The second (scalable data loading) and third (intra-microbatch reordering) are
    in progress and will follow as separate PRs on top of it.
    (cc. @sailor1493 @terryoo )

  17. liding-nv
    added
    x-naver
    on Jul 29, 2026
  18. going-song
    added 4 commits that reference this issue on Aug 12, 2026

    feat(data): scalable canonical-grid reads for MegatronMIMO data loading

    ee1db0e

    feat(data): scalable canonical-grid reads for MegatronMIMO data loading

    314e58a

    feat(data): scalable canonical-grid reads for MegatronMIMO data loading

    958575f

    feat(data): scalable canonical-grid reads for MegatronMIMO data loading

    2f6f339
  19. going-song
    added a commit that references this issue on Aug 27, 2026

    feat(data): scalable canonical-grid reads for MegatronMIMO data loading

    30280ca

  20. going-song commented on Oct 3, 2026

    @going-song
    going-song
    on Oct 3, 2026
    Contributor

    Hi @liding-nv — the remaining two splits are up:

    #6306 also carries two small unrelated fixes needed for MegatronMIMO training on current main; details in the PR description.
    (cc. @sailor1493 @terryoo)

  21. cuichenx
    removed
    waiting-on-customerWaiting on the original author to respond
    on Oct 3, 2026

  22. going-song commented on Oct 13, 2026

    @going-song
    going-song
    on Oct 13, 2026
    Contributor

    All three splits from #4609 (comment) are merged:

    The original PR #4608 is already closed, so this issue can be closed as well.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:dataDataset builders, preprocessing, and samplerscommunity-requestfeatureNew capabilities, enhancements, or enablement workx-naver

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions