Skip to content

Add encoder prefetch for heterogeneous MIMO training - #5833

Merged
yashaswikarnati merged 8 commits into
NVIDIA:mainfrom
yashaswikarnati:ykarnati/mimo-encoder-prefetch-public
Jul 24, 2026
Merged

Add encoder prefetch for heterogeneous MIMO training#5833
yashaswikarnati merged 8 commits into
NVIDIA:mainfrom
yashaswikarnati:ykarnati/mimo-encoder-prefetch-public

Conversation

@yashaswikarnati

@yashaswikarnati yashaswikarnati commented Jul 15, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What changed

  • add opt-in, configurable-depth prefetch for frozen MIMO encoder features with any compatible batch iterator.
  • keep completed features in a bounded FIFO and cap unfinished encoder GPU work at one batch.
  • overlap CPU generation of the next batch with the current GPU encode using one staged CPU batch.
  • retain encoder-rank input_ids on CPU, removing an unnecessary H2D copy and pre-bridge GPU-to-host synchronization.
  • keep the trainable projection on the schedule stream and inside autograd.

Validation

  • final cleanup bundle: Black, isort, Ruff, and 76 focused/surrounding MIMO tests passed on CW DFW 14223654
  • final 8-GPU heterogeneous mock smoke completed two iterations with finite loss/gradients and zero skipped/NaN iterations on CW DFW 14223688; all eight rank logs were checked

@copy-pr-bot

copy-pr-bot Bot commented Jul 15, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@yashaswikarnati

yashaswikarnati commented Jul 15, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 10cd3e1

@yashaswikarnati
yashaswikarnati force-pushed the ykarnati/mimo-encoder-prefetch-public branch from 10cd3e1 to b130c54 Compare July 16, 2026 20:57
@yashaswikarnati

yashaswikarnati commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test b130c54

@yashaswikarnati
yashaswikarnati force-pushed the ykarnati/mimo-encoder-prefetch-public branch from b130c54 to a6c0904 Compare July 16, 2026 21:49
@yashaswikarnati

yashaswikarnati commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a6c0904

1 similar comment
@yashaswikarnati

yashaswikarnati commented Jul 16, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a6c0904

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Jul 22, 2026
@yashaswikarnati
Explain encoder capture hook
704be29 
Signed-off-by: ykarnati <ykarnati@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Jul 23, 2026
@yashaswikarnati

yashaswikarnati commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 704be29

@yashaswikarnati

yashaswikarnati commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 704be29

@yashaswikarnati
Preserve CUDA device in prefetch tests
2019209 
Signed-off-by: ykarnati <ykarnati@nvidia.com>
@yashaswikarnati

yashaswikarnati commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 2019209

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30071620729

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30092280246

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30116065878

@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30117533510

This was referenced Jul 24, 2026
Closed
Merged
This was referenced Aug 17, 2026
Closed
Open
@svcnemo-autobot svcnemo-autobot mentioned this pull request Oct 7, 2026
Merged

This branch was previously deployed

2 inactive deployments
public — 2019209b Deployed Jul 23, 2026 by copy-pr-bot[bot] via finalize / publish-docs #4451
test — 2019209b Deployed Jul 23, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #17495
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@deepakn94 deepakn94 deepakn94 approved these changes
@yaoyu-33 yaoyu-33 yaoyu-33 approved these changes

Assignees

No one assigned

Labels

Approved All necessary approvals have been made complexity: high Run functional tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@yashaswikarnati @svcnvidia-nemo-ci @deepakn94 @yaoyu-33 @ko3n1g