| """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: |
| 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) |
|
|
