Skip to content

fix: reserve MOPD teacher clusters before Gym startup#3340

Open
yfw wants to merge 2 commits into
mainfrom
yifu/mopd_fix_topology
Open

fix: reserve MOPD teacher clusters before Gym startup#3340
yfw wants to merge 2 commits into
mainfrom
yifu/mopd_fix_topology

Conversation

@yfw

@yfw yfw commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Summary

Reserve topology-aware MOPD teacher placement groups before NeMo Gym starts allocating GPU-backed services. Teacher workers and model loading remain deferred until policy initialization completes, avoiding checkpoint-conversion races.

Problem

Previously, teacher topology was planned before Gym startup, but the corresponding placement groups were not reserved until afterward. Gym could consume GPUs in those selected NVLink domains, leaving a teacher placement group permanently unschedulable despite sufficient cluster-wide capacity. This caused setup to time out after some teachers initialized successfully.

Fix

Split teacher setup into two phases:

  1. Reserve all topology-aware teacher placement groups before Gym startup.
  2. Initialize teacher workers and load their models after policy initialization.

This preserves topology-aware teacher placement while preventing Gym from taking the reserved resources.

Issues

List issues that this PR closes (syntax):

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

  • ...

Reserve topology-aware teacher placement groups before Gym can consume their
GPU resources. Defer teacher worker and model initialization until the policy
is ready to avoid checkpoint conversion races.

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw
yfw requested review from a team as code owners July 24, 2026 18:46
@copy-pr-bot

copy-pr-bot Bot commented Jul 24, 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.

@yfw yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Jul 24, 2026
@yfw

yfw commented Jul 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 23a0c72

create_teacher_worker_groups kept a 'teacher_clusters is None' fallback that reserved the clusters itself. That call runs after NeMo Gym startup, so the fallback silently reintroduced the late reservation this branch fixes, and it left segment_size / teacher_segment_topology reachable only through it.

Make teacher_clusters a required keyword-only argument and drop the fallback along with the two now-dead parameters. The existing alias-mismatch check is now the sole guard on the reservation contract.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw

yfw commented Jul 25, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f39d43b

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

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant