File size: 3,402 Bytes
f0975ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """Train Dory1.0-4B on Beam traces on Modal (torchrun + DDP on B200s).
modal run modal_app/train_dory.py::download_dory # one-time: base weights -> Volume
modal run modal_app/train_dory.py --gpus 1 --run smoke --extra "--max-steps 20"
modal run modal_app/train_dory.py --gpus 8 --run r1 --data sft/sft_verified.jsonl
Data files are paths inside the `beam-dory-data` Volume (upload with `modal volume put`).
Checkpoints land in the `beam-dory-ckpt` Volume under runs/<run>/.
"""
import shlex
import subprocess
import time
import modal
APP = modal.App("beam-dory-train")
weights = modal.Volume.from_name("beam-weights", create_if_missing=True, version=2)
data = modal.Volume.from_name("beam-dory-data", create_if_missing=True, version=2)
ckpt = modal.Volume.from_name("beam-dory-ckpt", create_if_missing=True, version=2)
W, D, C = "/weights", "/data", "/ckpt"
DORY_REPO = "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16"
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install("torch==2.14.1", "safetensors", "transformers>=5.20", "jinja2", "numpy",
"huggingface_hub[hf_xet]>=0.34")
.env({"PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false"})
.add_local_dir("dory", "/proj/dory")
.add_local_dir("scripts", "/proj/scripts")
)
@APP.function(image=image, volumes={W: weights}, cpu=4, memory=16384, timeout=3600)
def download_dory():
from huggingface_hub import snapshot_download
snapshot_download(DORY_REPO, local_dir=f"{W}/dory-4b", allow_patterns=["*.json", "*.jinja", "*.safetensors"])
weights.commit()
print(subprocess.run(["ls", "-la", f"{W}/dory-4b"], capture_output=True, text=True).stdout)
@APP.function(image=image, gpu=["B200:8", "H200:8", "B200:4"], volumes={W: weights, D: data, C: ckpt},
timeout=12 * 3600, cpu=32, memory=262144)
def train(run: str, data_files: list, nproc: int = 0, extra: str = ""):
"""nproc=0 uses every GPU in the container (whichever fallback config was allocated)."""
if not nproc:
out_ = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, text=True)
names = [l for l in out_.stdout.splitlines() if l.strip()]
nproc = len(names)
print(f"got {nproc}x {names[0]}", flush=True)
if "--grad-accum" not in extra: # keep the global batch of 8 GPUs x 1 sequence on smaller fallbacks
extra += f" --grad-accum {max(1, 8 // nproc)}"
out = f"{C}/runs/{run}"
cmd = ["torchrun", "--standalone", f"--nproc_per_node={nproc}", "/proj/scripts/train_dory.py",
"--model", f"{W}/dory-4b", "--out", out, "--data", *[f"{D}/{f}" for f in data_files],
*shlex.split(extra)]
print("cmd:", " ".join(cmd), flush=True)
t0 = time.time()
proc = subprocess.run(cmd, cwd="/proj")
ckpt.commit()
print(f"exit {proc.returncode} after {(time.time() - t0) / 60:.1f} min", flush=True)
if proc.returncode != 0:
raise RuntimeError(f"training failed with exit code {proc.returncode}")
@APP.local_entrypoint()
def main(run: str, data: str = "sft/sft_verified.jsonl", gpus: int = 8, extra: str = ""):
"""Interactive runs (smoke tests). Long runs: deploy, then spawn via scripts/spawn_train.py."""
files = [f.strip() for f in data.split(",") if f.strip()]
train.with_options(gpu=f"B200:{gpus}").remote(run, files, gpus, extra)
|