Megatron-LM Base Image Bump

Bump the PyTorch base image

End-to-end workflow for moving Megatron-LM's CI to a newer nvcr.io/nvidia/pytorch:<YY.MM>-py3 container. The most common failure mode is forgetting that GitHub CI and GitLab CI have separate pins — a bump that only touches the former lands green, then breaks GitLab CI on main and forces an immediate follow-up PR. Always update both in the same PR.

Answer-First Pattern: dev Pin Sync

For a dev-only base-image bump, lead with the synchronization rule:

• docker/.ngc_version.dev is only the GitHub/local Dockerfile pin.
• GitLab CI has separate hardcoded BASE_IMAGE rows in .gitlab/stages/01.build.yml; update both IMAGE_TYPE: dev rows, one PLATFORM: amd64 and one PLATFORM: arm64.
• Leave docker/.ngc_version.lts and all IMAGE_TYPE: lts rows unchanged unless the user explicitly asks for an LTS bump.
• Verify before review with cat docker/.ngc_version.dev plus rg -n '^\s*BASE_IMAGE: nvcr\.io/nvidia/pytorch:' .gitlab/stages/01.build.yml | rg -B1 'IMAGE_TYPE: dev' | rg 'BASE_IMAGE'.

Inputs to gather from the user

1. Target tag, e.g. 26.04-py3. NVIDIA NGC PyTorch containers are released as nvcr.io/nvidia/pytorch:YY.MM-py3.
2. Scope — usually dev only. The lts pin (docker/.ngc_version.lts, plus the FILE: Dockerfile.ci.lts rows in GitLab) is bumped on a different cadence; only touch it if the user explicitly asks.
3. Workflow run ID (optional but typical) — after the first CI run, the user will provide a GitHub Actions run ID for golden-value refresh.

Workflow

- [ ] Step 1: Update the GitHub CI pin (docker/.ngc_version.dev)
- [ ] Step 2: Update the GitLab CI pin (.gitlab/stages/01.build.yml)
- [ ] Step 3: Open the PR with the `Run functional tests` label
- [ ] Step 4: Re-run failing tests via `/ok to test <commit-sha>`
- [ ] Step 5: For golden-value drift → refresh with the `update-golden-values` skill
- [ ] Step 6: For hangs / real regressions → mark tests `mr-broken` and file tracking issues
- [ ] Step 7: Verify both pins are in sync before merging

Step 1 — GitHub CI pin

docker/.ngc_version.dev is a single-line file consumed by docker/Dockerfile.ci.dev (via FROM_IMAGE_NAME=$(cat docker/.ngc_version.dev)). Overwrite it:

echo 'nvcr.io/nvidia/pytorch:<YY.MM>-py3' > docker/.ngc_version.dev

The file has no trailing newline historically; preserving or adding one is fine — the build args treat the value as $(cat ...). Do not touch docker/.ngc_version.lts unless bumping LTS too.

Step 2 — GitLab CI pin

GitLab CI does not read docker/.ngc_version.dev. It hardcodes BASE_IMAGE in a parallel: matrix: block. Update the two IMAGE_TYPE: dev rows (one per platform):

# .gitlab/stages/01.build.yml — under test:pre_build_image -> parallel.matrix
- IMAGE: CI_MCORE_DEV_IMAGE
  FILE: Dockerfile.ci.dev
  IMAGE_TYPE: dev
  BASE_IMAGE: nvcr.io/nvidia/pytorch:<YY.MM>-py3   # amd64 row
  PLATFORM: amd64
- IMAGE: CI_MCORE_DEV_IMAGE
  FILE: Dockerfile.ci.dev
  IMAGE_TYPE: dev
  BASE_IMAGE: nvcr.io/nvidia/pytorch:<YY.MM>-py3   # arm64 row
  PLATFORM: arm64

Leave the FILE: Dockerfile.ci.lts rows alone. Quick sanity check before commit:

rg -n '^\s*BASE_IMAGE: nvcr\.io/nvidia/pytorch:' .gitlab/stages/01.build.yml
# expect:  lts pin × 2 unchanged, dev pin × 2 == new tag

Step 3 — Open the PR

• Title convention: chore: Update Docker image version to <YY.MM>-py3 (see #4611).
• Apply the Run functional tests label before the first push. This unlocks the full functional matrix on the PR; without it the bump only runs the standard GH PR checks and you'll miss the drift.
• Push as draft first if you're still iterating; the bot will auto-draft otherwise.

Step 4 — Re-running CI on a new commit

For PRs from forks (the typical contributor case), each new commit needs an explicit /ok to test <commit-sha> PR comment to authorize NVIDIA runners (see the copy-pr-bot flow in #4611). One comment per commit. If copy-pr-bot reports "had a problem deploying to test", just push another commit (or re-issue the comment after the next push); the deploy is per-commit, not per-comment.

Step 5 — Golden-value drift

Container bumps shift CUDA / cuBLAS / cuDNN / kernel autotuning, which moves lm loss, num-zeros, iteration-time, and mem-* metrics on a large fraction of functional tests. This is expected and is not a correctness regression — refresh the golden values rather than chasing each test.

Hand off to the update-golden-values skill with:

• --source github
• --pipeline-id <WORKFLOW_RUN_ID> from the failing CI run
• --only-failing (refresh just the trajectories that drifted)

PR #4611 refreshed 78 golden-value files across dev_dgx_h100 and dev_dgx_gb200 for GPT / MoE / MIMO / hybrid suites in a single pass via this exact flow. The per-metric relative-difference summary the skill produces is the recommended PR description blurb — reviewers expect to see it.

Step 6 — Real regressions: mark broken, don't block the bump

A small number of tests will genuinely break (hangs, OOM, real numerical regressions). Don't gate the base-image bump on fixing them — that conflates two changes. Instead:

1. File a GitHub issue describing the failure mode and linking the failing CI run.

2. Flip the test's scope to the -broken variant in the recipe YAML under tests/test_utils/recipes/<arch>/, with an inline comment that references the issue. Pattern:

   - test_case: [hybrid_dynamic_inference_tp1_ep8_nanov3_chunked_prefill]
     products:
       - environment: [dev]
         # Broken: hangs on repeat iter 3, exceeds 1h job limit — see issue #<N>.
         scope: [mr-broken, mr-github-broken]      # was: [mr, mr-github]
         platforms: [dgx_h100]

   Scope mapping (replace, don't append):

   Before	After
   mr	mr-broken
   mr-github	mr-github-broken
   nightly	nightly-broken

   The recipe still runs in the -broken scope, but failures stop blocking PR merges.