exp108: 3B Qwen3 contacts-v1 sweep on CoreWeave GPU (batch priority)#111
Open
timodonnell wants to merge 13 commits into
Open
exp108: 3B Qwen3 contacts-v1 sweep on CoreWeave GPU (batch priority)#111timodonnell wants to merge 13 commits into
timodonnell wants to merge 13 commits into
Conversation
…riority) First MarinFold GPU/CoreWeave training run (issue #108). Scale-up of Eric's #75 tuning sweep: Qwen3 at #75's 1.5B width with layers doubled 24->48 (~2.9B), LR sweep {5e-4,1e-3,2e-3} at wd 0.2 / 10% warmup, 16 epochs, batch 128, seq 8192. - dispatch_train.py: direct grug-style Fray dispatch of each training gang at iris batch priority (JobRequest(priority=3)), reusing marin's run_levanter_train_lm + _build_train_lm_config with concrete S3 paths and on-the-fly tokenize (auto_build_caches). The executor can't set a priority band, which #108 requires for all work. - train_qwen_3b_contacts_v1_sweep.py: one batch-priority job per LR. - contacts_v1_train_common.py: shared S3/tokenizer/GPU constants + HF export. - stage_data_to_coreweave.py: GCS-mirror -> CoreWeave S3 corpus staging (done). - models/: add a `gpu` extra (marin-core[gpu]) mirroring `tpu`. Everything writes under a single s3://marin-us-east-02a/MarinFold/ prefix. Pre-launch: pin marinfold-models rev; smoke run to validate the direct-dispatch path live (cloudpickle, cache build, batch band, checkpoint layout). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
62ba9db carries the models/ `gpu` extra the CoreWeave run needs. Pins the git source so the iris worker resolves marinfold-models[gpu] reproducibly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
marin 0.2.38 refactored its execution framework — the old executor surface (ExecutorStep/this_output_path/executor_main/versioned/...) that the vendored marinfold_models imported is gone, and marin.experiment.train.train_lm is now the shipped assembler. exp85's committed uv.lock hid this by freezing old marin; a fresh sync today breaks. (Chosen fix: refresh the shared harness, not a one-off.) - models/marinfold_models/defaults.py: full rewrite. Dropped the dead executor-based default_train/default_tokenize/_build_train_lm_config. New StepContext-free `build_train_lm_on_pod_config` reproduces train_lm's TrainerConfig/mesh/checkpointer/mp assembly but takes a concrete output_path and a levanter LmDataConfig directly — so the caller can direct-dispatch it at batch priority. __init__ exports it + SimpleTrainConfig + MARIN_PRECISION. - exp108/dispatch_train.py: use the new builder (+ build the AdamW cosine optimizer in the driver); keep the JobRequest(priority=3) dispatch + the tokenize-on-the-fly S3 data config. - exp108/pyproject.toml: fix stale/conflicting pins surfaced by uv sync — torch 2.10->2.11 (match marin-core[cpu]), python >=3.12 (marin dev dropped 3.11), drop huggingface-hub>=1.5 (conflicts with marin's transformers), Linux-only environments. - exp108/export: WIP stub (marin's HF-export API also moved; not needed until a run produces a checkpoint). Validated live in the venv: build_train_lm_on_pod_config assembles a TrainLmOnPodConfig against 0.2.38; all exp108 modules import clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Bumps the marinfold-models git rev to 1f5653f (the modern-marin refresh) and commits uv.lock to FREEZE the resolution at marin 0.2.38.dev202607060946 — the root cause of this whole detour was the unpinned `<0.3` range silently drifting to a marin build that broke the vendored code. The lock (like exp67/exp85's) stops that. Verified: `uv sync` from the lock fetches marinfold-models@1f5653f from git and the TrainLmOnPodConfig builds clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Smoke run surfaced two launch issues: - The training gangs are CHILDREN of the driver job; with wait=False the driver exited immediately and iris finalized (killed) the gang. Fix: submit all gangs, then block on each (they run concurrently, driver stays alive for the run). - The workstation launcher needs marin-iris[controller] to reach the CoreWeave k8s controller (CloudK8sService); added it + re-locked. (Also needs kubectl on PATH for the controller tunnel — installed out-of-band, not a repo change.) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Left a wait={WAIT} in the summary print after removing the WAIT knob; the driver
crashed at import before dispatching. Now validated by running main() locally
with a stubbed dispatch (full driver path, no cluster submit).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Throughput investigation from the smoke run (~52k tok/s, ~15% MFU on 8xH100): - Root cause of low attention throughput: transformer_engine isn't in the --extra gpu env, so levanter's default GPU NVTE attention silently fell back to the VANILLA O(seq^2) kernel — brutal at seq 8192. Added an EXP108_ATTN knob (default jax_flash = levanter's pallas flash, no TE dep). NB: measured no net speedup from flash alone — attention is only ~15% of FLOPs here; the bottleneck is systemic (FSDP collectives / gradient-checkpointing recompute, likely worsened by the CUDA fabric-memory fallback). - dispatch_train now forwards XLA_FLAGS/NCCL_/JAX_ from the driver to the gang (like grug's dispatch) — the gang is a separate pod and wasn't inheriting launcher-set flags, so throughput was untunable via -e XLA_FLAGS. This is a correctness fix + the hook for further tuning. - EXP108_RUN_SUFFIX isolates probe runs' S3 output paths. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Injects a ProfilerConfig (window past compile) into the trainer when EXP108_PROFILE=1. Works (captures a perfetto trace), but note levanter uploads only jaxpr/HLO to W&B, not the profiler dir — retrieving the timing trace off the ephemeral pod needs extra plumbing (log_dir is a local PosixPath). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Probe confirmed grad-checkpointing-off OOMs (692 GiB activations for 48 layers at seq 8192, even at 8 seq/GPU) — so it stays ON for the real run. Knob kept for the record / future selective-checkpointing work. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Captures everything learned taking the sweep live on cw-rno2a: - Workstation prereqs (kubectl, marin-iris[controller], wandb-from-netrc). - Operational findings: driver must wait on gangs; node-count ceiling ~4 (8-node JAX coordination bootstrap fails; likely NCCL env); launch runs as separate staggered drivers (3-from-one-driver collides); benign FABRIC warns. - Throughput/scaling: ~20 s/step 1-node (~15% MFU), near-linear to ~5 s/step at 4 nodes (~4-day full run); gradient_checkpointing mandatory (off = 692 GiB OOM); attention backend irrelevant. - Results: smoke run validated the full path; real 3-LR sweep launched at 4 nodes each (v2 runs), all three training. Target: beat #75's 2.7566 eval loss. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e agents Durable, cluster-general lessons from exp108's first MarinFold GPU run: the kubectl + marin-iris[controller] launcher prereqs; S3 (not GCS) storage with virtual-hosted addressing; batch priority not propagating to executor children (dispatch directly); drivers must wait on child gangs; the ~4-node multi-node ceiling + separate-staggered-drivers pattern; benign FABRIC warnings + the missing-Transformer-Engine attention fallback; and always pin + uv.lock marin. Recipe specifics stay in the exp108 README. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… diverged The v2 sweep (depth-doubled 24->48 layers) went unstable at all three LRs — diff vs #75's stable 1.5B run showed the ONLY change was num_layers (48 vs 24); everything else (LR/wd/warmup/init/clip/z_loss) was identical. LR transfers under width scaling, not depth, so 48 layers lowered the stable-LR ceiling. Fix: widen #75's 24-layer model instead (hidden 2048->2816, ff 8192->11264, GQA 4:1) to ~2.78B, and turn on Qwen3 QK-norm (was off) for deep-attention stability. #75's recipe/LRs should now transfer. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e fix Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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.
First MarinFold GPU/CoreWeave training run — issue #108. A scale-up of Eric's #75 tuning sweep to see whether we beat #75's best eval loss (~2.7).
What this is
Llama3RotaryEmbeddingsConfig) with layers doubled 24→48 → ~2.9B params. Depth-only scaling keeps exp: run contacts-v1 1.5B tuning sweep #75's width so its tuned LR/wd transfer.cw-rno2a(CoreWeave RNO2A, 512× H100). All artifacts under a singles3://marin-us-east-02a/MarinFold/prefix (removable later, per exp: train a 3B qwen model contacts-v1 on coreweave GPUs #108).Key design decision — batch priority via direct dispatch
#108 requires iris batch priority for all work.
--priority batchonly sets the driver band; the marin executor submits child jobs without a band (→ interactive), and there's no executor knob for it. Sodispatch_train.pysubmits each training gang itself as afray.JobRequest(priority=3)— grug-style — while still reusing marin'srun_levanter_train_lm+_build_train_lm_config(AdamW/TrainerConfig/mesh/checkpointer). Data is tokenized on the fly from raw S3 parquet (auto_build_caches=True), so there's no separate executor tokenize step.Files
dispatch_train.py— direct batch-priority Fray dispatch (the crux).train_qwen_3b_contacts_v1_sweep.py— sweep entry point.contacts_v1_train_common.py— shared S3/tokenizer/GPU constants + HF export.export_qwen_3b_contacts_v1.py,stage_data_to_coreweave.py.models/pyproject.toml— newgpuextra (marin-core[gpu]).Status
marinfold-modelspinned to 62ba9db (carries thegpuextra).Closes #108 is intentionally not set — experiments are closed by humans.
🤖 Generated with Claude Code