feat(launch): opt-in cuda13 workspace venv (pd-lm --cuda-extra cuda13)#959
Closed
danbraunai-goodfire wants to merge 1 commit into
Closed
feat(launch): opt-in cuda13 workspace venv (pd-lm --cuda-extra cuda13)#959danbraunai-goodfire wants to merge 1 commit into
danbraunai-goodfire wants to merge 1 commit into
Conversation
cuBLAS >= 13.2 fixes the Blackwell TMEM double-free (silent-corruption hazard cuBLAS 12.9 warns about on B200). Same jax pin, cu13 runtime libs. cuda (cu12) stays the default: and-gmi's r570 driver can't run CUDA 13. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FdeKYLCnk7E8BhKZoNvhAc
ocg-goodfire
added a commit
that referenced
this pull request
Jul 7, 2026
Port the cuda13 extra from #959 (jax[cuda13]==0.10.1 + the Blackwell-TMEM-fixed cuBLAS >= 13.2) and drop its --cuda-extra flag: now that the venv is built job-side on the compute node, the driver IS observable at build time, so each node picks cuda13 vs cuda off its own driver major (>= 580 -> cuda13; and-gmi's r570 falls back to cu12). #959's stated reason for keeping cu12 the default was only the and-gmi driver floor — this gives every driver-capable cluster the corruption-fixed cuBLAS automatically. Co-Authored-By: Dan Braun <dan.braun@goodfire.ai> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm
Collaborator
|
Subsumed by #960: with job-side node venvs the driver is observable at build time, so the node picks |
ocg-goodfire
added a commit
that referenced
this pull request
Jul 7, 2026
…d-FS build (#960) * refactor(pd-lm): job-side node workspaces, drop the submit-time shared-FS build Return the trainer launch to the repo's job-side setup pattern (the one harvest/autointerp never left): each node clones the snapshot into node-local /tmp and builds the CUDA venv at job start, instead of pd-lm materializing an immutable shared-FS workspace per run at submit. - Config delivery moves to the run dir: submit pins the (group/tags-stamped) config as runs/<id>/launch_config.yaml — the file the trainer already pins — and the job (and every requeue) reads that copy. The config no longer needs to live inside the repo, and the inline (dp=null) path no longer mutates the submitted config file in place. - Nodes clone/fetch from ~/param-decomp (durable; worktrees share its refs), not the submitting checkout, which may be an ephemeral worktree a requeue would no longer find; visibility of the snapshot ref is asserted at submit. - The srun task still execs the trainer (SIGTERM-save delivery unchanged), so cleanup is a batch-script EXIT trap sweeping /tmp workspaces, with a start-of-task sweep as the backstop for hard kills. Trainer jobs hold whole nodes, so the sweep cannot race a concurrent job. - pd-pretrain keeps the submit-time build: its 8-tasks-per-node layout would race a per-node job-side build. Smoke: p-a5769c4b (job 171732, 1 node). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm * fix(pd-lm): clone from the submitting checkout's common git dir, not ~/param-decomp The $HOME/param-decomp hardcoding was inherited from infra's shared fragment, where it exists only because torch-era async slow-eval generated scripts from inside a SLURM job (REPO_ROOT = node-local /tmp there). pd-lm always submits from a login-node checkout, so resolve `git rev-parse --git-common-dir` at submit instead: the main checkout's .git even under a worktree — durable across worktree deletion for requeues, carries the snapshot ref by construction, and assumes nothing about where the checkout lives. .env rides from its parent. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm * refactor(pd-lm): fetch snapshots from origin — remotely pushed refs are ground truth The async-slow-eval-era local-checkout plumbing isn't blessed; origin is the source of truth for snapshot refs. create_git_snapshot's origin push becomes a hard requirement (a failed push fails the submit), and each node now git-inits an empty workspace and shallow-fetches the snapshot ref straight from the origin URL — no job dependency on any local checkout. .env (secrets, not in git) is staged into the run dir at submit alongside the pinned launch config, mode-preserved, and copied into the workspace by the node setup. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm * feat(launch): driver-gated cuda13 node venvs (subsumes #959) Port the cuda13 extra from #959 (jax[cuda13]==0.10.1 + the Blackwell-TMEM-fixed cuBLAS >= 13.2) and drop its --cuda-extra flag: now that the venv is built job-side on the compute node, the driver IS observable at build time, so each node picks cuda13 vs cuda off its own driver major (>= 580 -> cuda13; and-gmi's r570 falls back to cu12). #959's stated reason for keeping cu12 the default was only the and-gmi driver floor — this gives every driver-capable cluster the corruption-fixed cuBLAS automatically. Co-Authored-By: Dan Braun <dan.braun@goodfire.ai> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm * refactor(pd-lm): fetch snapshots from local shared-FS git, keep origin push as provenance Separate the two things "remotely pushed refs are ground truth" bought: the hard origin push stays (durable provenance record for every run), but the job-side fetch source returns to local git state like main — the submitting checkout's common git dir over shared FS (file:// so --depth survives the local transport). Requeues no longer depend on GitHub reachability or SSH auth from compute nodes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm * refactor(git): snapshot origin push back to best-effort, like main Jobs fetch snapshots from local shared-FS git, so origin is a provenance backup — a GitHub outage or missing push perms shouldn't block submits. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PvoYnzxByyYP4dmQtXWSmm --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Dan Braun <dan.braun@goodfire.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
cuda13optional-dependency extra:jax[cuda13]==0.10.1+nvidia-cublas>=13.2.0.9(locks to 13.6.0.2). Same jax/jaxlib pin — only the CUDA runtime libs and plugin change.pd-lm --cuda-extra cuda13selects it for the workspace venv (default remainscuda/cu12). Submit-time machine decision, like the venv-derivedLD_LIBRARY_PATH.Motivation and Context
On B200, jax's cu12 wheels load cuBLAS 12.9, which has a known TMEM double-free: "Executing a cuBLAS kernel concurrently with another kernel (e.g. on another stream) can lead to silent data corruption" (the warning every current B200 run prints at startup — and our steps run cuBLAS GEMMs concurrently across streams with the latency-hiding scheduler). cuBLAS ≥ 13.2 fixes it; NVIDIA's JAX release notes confirm. cu12 stays the default because and-gmi's r570 driver is below the CUDA-13 floor (r580); cw-east/cw-rno2 (595.71) and btdr (580.126) can all run it.
How Has This Been Tested?
Two 4-node full32L runs on cw-east B200 with
--cuda-extra cuda13(jobs 171656/171657): workspace builds cleanly, the TMEM warning is gone, and bf16 throughput is unchanged vs cu12 (5.62 vs 5.65–5.7 s/step).uv lockresolves;make checkclean.Does this PR introduce a breaking change?
No — default behavior is byte-identical; cuda13 is opt-in per launch.
🤖 Generated with Claude Code
https://claude.ai/code/session_01FdeKYLCnk7E8BhKZoNvhAc