Skip to content

perf(megatron): skip redundant weight load - #3397

Merged
terrykong merged 3 commits into
NVIDIA-NeMo:mainfrom
tdene:tde/skip_inference_weight_load
Aug 5, 2026
Merged

perf(megatron): skip redundant weight load#3397
terrykong merged 3 commits into
NVIDIA-NeMo:mainfrom
tdene:tde/skip_inference_weight_load

Conversation

@tdene

@tdene tdene commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

#2355 updated Megatron Inference to be a full peer of vllm inference. Model conversion still happens twice, once on training workers, once on inference workers. It should happen only once, and just be streamed across.

Issues

List issues that this PR closes (syntax):

#2615

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

@tdene
tdene requested review from a team as code owners July 29, 2026 04:53
@copy-pr-bot

copy-pr-bot Bot commented Jul 29, 2026

Copy link
Copy Markdown

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@tdene
tdene force-pushed the tde/skip_inference_weight_load branch from bc65313 to 6507e43 Compare July 29, 2026 06:12
@tdene tdene added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 29, 2026
@tdene

tdene commented Jul 29, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 6507e43

shanmugamr1992
shanmugamr1992 previously approved these changes Jul 29, 2026
@tdene
tdene force-pushed the tde/skip_inference_weight_load branch from 6507e43 to aa96ed3 Compare August 4, 2026 05:33
@tdene

tdene commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 5b80bd4

@tdene
tdene requested a review from a team as a code owner August 4, 2026 09:26
cspades
cspades previously approved these changes Aug 4, 2026
Comment thread tests/unit/models/generation/test_megatron_generation.py
shanmugamr1992
shanmugamr1992 previously approved these changes Aug 4, 2026
@tdene tdene added CI:L1 Run doctests, unit tests, and functional tests and removed CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) labels Aug 4, 2026
@tdene
tdene dismissed stale reviews from shanmugamr1992 and cspades via 95d5a81 August 4, 2026 18:55
@tdene

tdene commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 95d5a81

@tdene
tdene force-pushed the tde/skip_inference_weight_load branch from 95d5a81 to e19d73c Compare August 4, 2026 22:59
@tdene

tdene commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test e19d73c

tdene added 3 commits August 4, 2026 22:19
Closes NVIDIA-NeMo#2615

Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Signed-off-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
@tdene
tdene force-pushed the tde/skip_inference_weight_load branch from e19d73c to 91ccce7 Compare August 5, 2026 03:23
@tdene

tdene commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 91ccce7

@terrykong
terrykong merged commit 80a84eb into NVIDIA-NeMo:main Aug 5, 2026
98 of 99 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L1 Run doctests, unit tests, and functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants