Skip to content

Fix concurrent VMM allocator extension builds - #7494

Merged
yashaswikarnati merged 1 commit into
NVIDIA:mainfrom
yashaswikarnati:codex/gtp-vmm-build-race-main
Oct 20, 2026
Merged

Fix concurrent VMM allocator extension builds#7494
yashaswikarnati merged 1 commit into
NVIDIA:mainfrom
yashaswikarnati:codex/gtp-vmm-build-race-main

Conversation

@yashaswikarnati

@yashaswikarnati yashaswikarnati commented Sep 18, 2026

Copy link
Copy Markdown
Contributor
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Prevent ranks from rewriting the VMM allocator's C++ source while another rank compiles it. PyTorch's load_inline() publishes main.cpp before taking its build lock; on a cold cache, competing writers can truncate the source while the compiler is reading it, potentially producing a shared library without the expected Python module initializer.

Move the existing C++ into megatron/core/allocator/csrc/vmm_symm_allocator.cpp and build it with cpp_extension.load(). Include the source in wheels and source distributions. Preserve the allocator's behavior, extension name, build directory, and success/failure caches. Add regression coverage for packaged-source loading and failure caching, and include the C++ path in the allocator's existing determinism exemption.

Validation

  • Two allocator unit tests passed through torch.distributed.run with the normal repository conftest.
  • Four simultaneous processes completed a cold build/load of the real extension; output showed one compiler/linker invocation and successful loads on all four ranks.
  • Built a source distribution and a wheel from it; both contain the exact C++ source.
  • C++ token comparison against the former inline source passed, excluding the explicit <cstdlib> include.
  • Black, isort, Ruff, Pylint, the determinism coverage gate, and git diff --check passed.

Local validation used the existing megatron-lm:local development image (PyTorch 2.13.0a0+8145d630e8.nv26.06) and its bundled CUDA compatibility libraries. Current main's CI image pin is 26.08, so this is not exact CI-image parity. No GPU memory allocation or NCCL collective was exercised. The unit-test fixture attempted an unavailable dataset download; these two tests do not use that dataset.

Issue tracking

Small build-race fix; no issue linked.

Contribution process

Pre-checks

  • I have added relevant unit tests
  • I have added relevant functional tests (not added; local concurrent build smoke test performed)
  • Determinism coverage: retain the existing allocation-only exemption and register the extracted C++ source; no numerical kernel change
  • I have added proper typing to my code (no public API added)
  • I have added relevant documentation (explain the source-publication race at the build call)
  • I have run tools/autoformat.sh in check mode (Black, isort, Ruff, and Pylint passed; optional mypy was unavailable in the container)

All PRs start as drafts.

@yashaswikarnati
Fix concurrent VMM allocator extension builds
c217021 
Signed-off-by: ykarnati <ykarnati@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 18, 2026

Copy link
Copy Markdown

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

Contributors can view more details about this message here.

@yashaswikarnati
yashaswikarnati marked this pull request as ready for review September 18, 2026 21:29
@yashaswikarnati
yashaswikarnati requested review from a team as code owners September 18, 2026 21:29
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Sep 18, 2026
@deepakn94 deepakn94 closed this Sep 19, 2026
@fanshiqing

fanshiqing commented Oct 19, 2026

Copy link
Copy Markdown
Member

Curious why this MR is closed but not merged? It's also required for the MLM/main branch. I ran into this bug for the MLM/main branch today.
@yashaswikarnati @deepakn94

@yashaswikarnati

yashaswikarnati commented Oct 19, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test c217021

@fanshiqing fanshiqing linked an issue Oct 19, 2026 that may be closed by this pull request
[feat] GTP+Symmetric Memory Registration buffer #6540
Closed
@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 Oct 19, 2026
@yashaswikarnati
yashaswikarnati added this pull request to the merge queue Oct 19, 2026
@nemo-automation-bot

nemo-automation-bot Bot commented Oct 19, 2026

Copy link
Copy Markdown

🔄 Merge queue validation started!

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

@fanshiqing fanshiqing linked an issue Oct 20, 2026 that may be closed by this pull request
GTP + DSV4 perf optimization #8338
Closed
Merged via the queue into NVIDIA:main with commit 6c8f14a Oct 20, 2026
209 of 214 checks passed
@yashaswikarnati
yashaswikarnati deleted the codex/gtp-vmm-build-race-main branch October 20, 2026 01:01
@FDecaYed FDecaYed mentioned this pull request Oct 22, 2026
Merged
nemo-autobot-origin Bot pushed a commit to svcnemo-autobot/Megatron-LM that referenced this pull request Oct 22, 2026
@svcnemo-autobot
fix: post-CI corrections
98b43a0 
Repair main-to-dev sync compatibility for the 2026-10-20 slice
(main 13a9d46 merged into
dev ff4fbc9):

- moe/router: InferenceTopKRouter._compiled_topk_routing accepts and
  forwards precomputed_indices, matching the base router call site.
- tensor_parallel/random: MHCCheckpointManager subclasses
  CheckpointWithoutOutputManager so main's checkpoint consumers accept it.
- moe/token_dispatcher: _HybridEPManager.setup_metadata clears the stale
  num_permuted_tokens budget when no capacity factor is configured.
- tests: align dev-only tests with main API changes (test_training
  TrainingConfig and TrainState, test_mhc_schedule _model_chunk_state,
  test_hybrid_hash_routing tensor_parallel import, test_route_metadata_cpu
  token_dispatcher fake, MTP argmax acceptance-count test).
- tests: test_train_step_schedule_plumbing provides the cfg.train,
  cfg.validation, cfg.ddp and TrainState inputs main's train_step and
  setup_model_and_optimizer now read.
- tests: test_virtual_experts quantile cases select the micro_batch
  estimation scope that main's quantile estimator implements; the merged
  default keeps dev's global_batch scope.
- tests/test_utils/recipes/gb200/unit-tests.yaml: drop the duplicated
  tests/unit_tests/ssm/test_gated_delta_net.py bucket produced by the
  merge, restoring the one-bucket-per-file unit-test scope contract.
- tests: test_hybrid_distillation_entry expects the multimodal_dev entry
  point to leave runtime services and tokenizer vocabulary resolution to
  initialize_megatron, which main NVIDIA#8272 made their owner.
- examples/mimo_bagel/train.py: read the training iteration from
  get_train_state() for the PP>1 data seed and the loss-alignment audit
  log, because main NVIDIA#7932 no longer sets args.curr_iteration or
  args.iteration (review 5609274485).
- pyproject.toml: carry main NVIDIA#7494's [tool.setuptools.package-data]
  entry so wheel and sdist installs ship
  megatron/core/allocator/csrc/vmm_symm_allocator.cpp (review 5609274485).

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
nemo-autobot-origin Bot pushed a commit to svcnemo-autobot/Megatron-LM that referenced this pull request Oct 22, 2026
@svcnemo-autobot
fix: post-CI corrections
73c827d 
Repair main-to-dev sync compatibility for the 2026-10-20 slice
(main 0156c21 merged into
dev ae76c60):

- moe/router: InferenceTopKRouter._compiled_topk_routing accepts and
  forwards precomputed_indices, matching the base router call site.
- tensor_parallel/random: MHCCheckpointManager subclasses
  CheckpointWithoutOutputManager so main's checkpoint consumers accept it.
- moe/token_dispatcher: _HybridEPManager.setup_metadata clears the stale
  num_permuted_tokens budget when no capacity factor is configured.
- tests: align dev-only tests with main API changes (test_training
  TrainingConfig and TrainState, test_mhc_schedule _model_chunk_state,
  test_hybrid_hash_routing tensor_parallel import, test_route_metadata_cpu
  token_dispatcher fake, MTP argmax acceptance-count test).
- tests: test_train_step_schedule_plumbing provides the cfg.train,
  cfg.validation, cfg.ddp and TrainState inputs main's train_step and
  setup_model_and_optimizer now read.
- tests: test_virtual_experts quantile cases select the micro_batch
  estimation scope that main's quantile estimator implements; the merged
  default keeps dev's global_batch scope.
- tests/test_utils/recipes/gb200/unit-tests.yaml: drop the duplicated
  tests/unit_tests/ssm/test_gated_delta_net.py bucket produced by the
  merge, restoring the one-bucket-per-file unit-test scope contract.
- tests: test_hybrid_distillation_entry expects the multimodal_dev entry
  point to leave runtime services and tokenizer vocabulary resolution to
  initialize_megatron, which main NVIDIA#8272 made their owner.
- examples/mimo_bagel/train.py: read the training iteration from
  get_train_state() for the PP>1 data seed and the loss-alignment audit
  log, because main NVIDIA#7932 no longer sets args.curr_iteration or
  args.iteration (review 5609274485).
- pyproject.toml: carry main NVIDIA#7494's [tool.setuptools.package-data]
  entry so wheel and sdist installs ship
  megatron/core/allocator/csrc/vmm_symm_allocator.cpp (review 5609274485).

Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

This branch was previously deployed

2 inactive deployments
public — c2170212 Deployed Oct 19, 2026 by copy-pr-bot[bot] via build-test-publish-wheels / publish-wheels (megatron-fsdp, amd64) #12344
test — c2170212 Deployed Oct 19, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #24602
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
@fanshiqing fanshiqing fanshiqing approved these changes
@ko3n1g ko3n1g ko3n1g approved these changes
@cspades cspades cspades approved these changes

Assignees

No one assigned

Labels

Approved All necessary approvals have been made complexity: medium nemotron

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

GTP + DSV4 perf optimization [feat] GTP+Symmetric Memory Registration buffer

6 participants

@yashaswikarnati @fanshiqing @deepakn94 @ko3n1g @cspades @svcnvidia-nemo-ci