exp112: 3B Qwen3 on CoreWeave via NVIDIA NeMo — MFU comparison vs Levanter#113
Open
timodonnell wants to merge 9 commits into
Open
exp112: 3B Qwen3 on CoreWeave via NVIDIA NeMo — MFU comparison vs Levanter#113timodonnell wants to merge 9 commits into
timodonnell wants to merge 9 commits into
Conversation
First non-JAX training path on the CoreWeave cw-rno2a cluster: dispatch a foreign NVIDIA NeMo container (nvcr.io/nvidia/nemo:25.04.02) as a batch-priority Fray job with a torchrun BINARY entrypoint, to benchmark MFU on the same 3B Qwen3 contacts-v1 model as exp108's Levanter run (~15% MFU baseline). Approach (benchmark-first, per #112 scope): - Single 8xH100 node, torchrun --standalone (sidesteps iris's missing multi-node torchrun rendezvous API; only DP is needed at 2.9B). - Mock/synthetic data for the MFU number (compute-bound, data-independent), so the tokenize->bin/idx pipeline is deferred to a real run. - dispatch_nemo.py base64-inlines the in-container train script (no repo checkout in the NeMo image); create_environment(docker_image=...) so the launcher workspace is NOT synced into the container; JobRequest(priority=3) batch band. - nemo_train_qwen3.py builds the exact exp108 geometry as a Megatron llm.GPTConfig (2048h/48L/32h.8kv GQA/qk_layernorm/RMSNorm/SwiGLU/RoPE, vocab 2845) and reports throughput (tokens/s, s/step) + MFU. Throughput leads (formula-free, directly comparable to exp108); MFU is the cross-check. Validated live in the NeMo 2.3.2 container on a local GPU (--tiny): tokenizer load, full model/data/trainer/optimizer construction, 5+ training steps, and the throughput/MFU summary all work. Launcher dry-run builds the batch-band JobRequest (image + nproc=8 + inlined bootstrap). Not yet run on the cluster. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
Single 8xH100, batch 128 x seq 8192, bf16, grad-ckpt on, mock data (MFU is
compute-bound, so mock is faithful). Micro-batch sweep {1,2,4}:
mbs s/step tok/s MFU
1 11.69 89,737 25.4%
2 11.34 92,465 26.2%
4 11.29 92,861 26.3% <- best
vs exp108 Levanter ~20 s/step / ~52k tok/s / ~15% MFU => ~1.75-1.79x.
The whole novel path worked first try on cw-rno2a: nvcr.io image pull, torchrun
binary entrypoint at batch band, base64-inlined script, 8-GPU bringup, S3 result
upload. Remaining MFU headroom (to ~35-40%) is grad-ckpt recompute + bf16 (no fp8)
+ DP-replicated params — all documented as deferred levers.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
…(issue #112) Extends the benchmark scaffold toward the actual 16-epoch trained model (Tim's ask: a trained model, sooner). Validated locally in the NeMo 2.3.2 container. - prepare_megatron_data.py: contacts-v1 parquet `document` -> JSONL -> preprocess_data.py (HF tokenizer, --append-eod) -> Megatron .bin/.idx on S3. Ran full corpus: 4,129,682 docs -> 9.35 GB train bin (+ val). - nemo_train_qwen3.py (--data real): PreTrainingDataModule on bin/idx with reset_attention_mask/reset_position_ids (block cross-doc attn); periodic val loss; NeMo ModelCheckpoint; an S3CheckpointSync callback that syncs the sharded (DCP) checkpoint per-node to S3 (no shared FS on cw-rno2a) + prunes/commits latest.txt; a resume preamble that mirrors the latest S3 checkpoint back before AutoResume. Callbacks moved to MODULE level (NeMo io-dumps the trainer context via fiddle, which can't pyref a main.<locals> class). - dispatch_nemo.py: replicas=N gang; bootstrap does S3 torchrun rendezvous (rank-0 publishes IRIS_ADVERTISE_HOST to an attempt-scoped key; peers poll) for multi-node, --standalone for single; downloads bin/idx per node; max_retries_preemption=100 so a preempted gang auto-restarts and resumes. - common.py: replicas param, data/ckpt/rdzv S3 prefixes, num_train_steps (71,312 @ bs128 = 16 epochs, matches exp108). Locally validated (single-GPU, container): real-data training, checkpoint save + S3 sync (per-node shard upload, prune, latest.txt), and RESUME — a fresh container pulled the S3 checkpoint and continued from the saved step with LR schedule intact. Multi-node rendezvous is validated next on a 2-node cluster smoke. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
Validated 16-rank (2×8 H100) training + val + checkpoint sync live. Fixes for the real multi-node / no-shared-FS gotchas found on the cluster (each was a distinct blocker uncovered in sequence): - Rendezvous: IRIS_TASK_ID has no ":attempt" suffix in-pod -> use a gang-common (run_name) key, not a per-node one; publish/poll rank-0's IP over the CONSISTENT object endpoint (cwobject.com), because the injected LOTA cache negative-caches a cross-node poll-before-write. - Gloo: pin GLOO_SOCKET_IFNAME to the host-eth iface — found via Python stdlib SIOCGIFADDR (the NeMo container has no working `ip`), matching IRIS_ADVERTISE_HOST. Without it Gloo (Megatron CPU groups + distributed optimizer) connectFullMesh fails across nodes. - Dataset index (no shared FS): Megatron builds the .npy on GLOBAL rank-0 into a node-local dir; patch the builder so DURING build get_rank() reports LOCAL rank -> each node's local-0 builds its own cache before the collective barrier. - Checkpoint sync: detect a NEW checkpoint by its consumed_samples (NeMo's step- count trigger fired before ModelCheckpoint wrote the dir); each node uploads its own DCP shards (union = complete), rank-0 prunes + commits latest.txt; also sync on_train_end for the final checkpoint. - set -e footgun in the bootstrap (`[ ] && cmd` -> if/then). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
Validated 2-node resume live (restore sharded ckpt from S3, continue with model state preserved — loss picked up at 5.84 exactly where it left off). - Rendezvous freshness: a preemption-restart reuses the run_name key, so a peer could read the previous attempt's DEAD master IP and hang 900s in torchrun. The poller now stamps a probe to read the S3 server clock and only trusts `master` if written this attempt (age < 120s), else waits for rank-0's fresh write. (No attempt-unique iris env var is available in-pod.) - Checkpoint sync throttle: NeMo bumps the `-last` ckpt's consumed_samples every step, so the detect-new-ckpt sync uploaded ~40 GB every step. Throttle to one upload per checkpoint_every global steps (deterministic across ranks -> barrier-safe); still catches the final ckpt on_train_end. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
…rope theta Replicate the best contacts-v1 model from #75 (Qwen3 1.47B: 2048h/8192ffn/32h·8kv/ 24 layers, qk-norm, Llama3 rope theta 500000, lr 1e-3/wd 0.2) but 16 epochs. The earlier 48-layer ~2.9B config (from a stale exp108 README) was a by-depth 3B, not #75; the actual #75 1.5B is 24 layers at this width (exp67/exp85 protein_llama_1_5b + Qwen3 qk-norm/rope per marin eac/plm-exp75). Set rotary_base 500000 (the #75 Llama3 theta; NeMo's generic GPTConfig lacks the llama3 freq-scaling knob, a context-extension feature dubious at train seq==8192 — one documented divergence). Run renamed plm-exp112-cv1-1_5b-nemo-e16-lr1e-3-wd0p2. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
…2:0') crash) IRIS_TASK_ID is inconsistent: some gangs get '/user/job/gang/<rank>', others '/user/job/gang/<rank>:<attempt>'. An earlier simplification dropped the ':attempt' strip, so a 4-node gang whose id carried ':0' crashed on int(NODE_RANK). Restore the two-step parse (strip ':attempt' then take the last path component). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
The 1.5B run failed at step ~35.5k when one coscheduled sibling node died and gang semantics killed the rest with no retry (max_retries_failure=0). The code is validated (35k steps), so a mid-run failure is transient infra, not a bug — set max_retries_failure=20 so a failed gang auto-retries and resumes from the S3 checkpoint (each retry re-runs the bootstrap -> restore -> AutoResume). Also bump the S3 client max_pool_connections to 32 (>= the 16-thread checkpoint uploader) to stop the 'connection pool is full' warning flood during checkpoint sync. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf
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
1 participant
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.
Answers issue #112: same 3B Qwen3 contacts-v1 model as exp108, but on NVIDIA NeMo / Megatron-Core instead of JAX/Levanter, to see whether we beat exp108's poor ~15% MFU.
Result — yes, ~1.75×
Single 8×H100, batch 128 × seq 8192, bf16, grad-ckpt on, mock data (MFU is compute-bound → mock is faithful; it's how NVIDIA publishes its perf tables). NeMo 2.3.2 /
nvcr.io/nvidia/nemo:25.04.02.~26% MFU vs ~15%, ~93k vs ~52k tok/s, ~11.3 vs ~20 s/step on identical model/shape/hardware, out of the box. Micro-batch barely matters (memory-bound with grad checkpointing).
What this is / how it works
The novelty vs exp108 is the payload, not the submission: exp108 dispatched a levanter callable into marin's iris-task image; exp112 dispatches a
torchrunbinary into the foreign NeMo container at batch priority.ResourceConfig.image) and run a binary entrypoint (Entrypoint.from_binary). The only iris gap — no torchrun rendezvous API — bites multi-node only; single-nodetorchrun --standalonesidesteps it. Batch band viaJobRequest.priority=3(as exp108).dispatch_nemo.pybase64-inlines the training script into the bootstrap bash;create_environment(docker_image=…)(notworkspace) so the launcher env is not synced into the container.llm.GPTConfig— 2048h/48L/32h·8kv GQA / QK-norm / RMSNorm / SwiGLU / RoPE, vocab 2845.Validated locally in the NeMo container (
--tiny) before any cluster time, then the full path ran first-try oncw-rno2a.Files
common.py(constants),nemo_train_qwen3.py(in-container, base64-inlined),dispatch_nemo.py(batch-priority Fray dispatch), launcher-onlypyproject.toml/uv.lock. Everything under one removables3://…/MarinFold/exp112_qwen_3b_nemo_mfu/prefix.Deferred (documented in README)
Remaining MFU headroom (→ ~35–40%) = grad-ckpt recompute + bf16 (no fp8) + DP-replicated params — all known levers. A full run would need the real
.bin/.idxdata path + multi-node torchrun rendezvous.Does not close #112 — experiments are closed by humans.
🤖 Generated with Claude Code
https://claude.ai/code/session_01SH6m5LcVoDRKsDYVD89Fwf