| """2-process CPU DDP smoke test (gloo): no hang across an epoch boundary, and resume works.""" |
| import json |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
|
|
|
|
| def _run(tmp: Path, *extra: str, port: str) -> subprocess.CompletedProcess: |
| cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node", "2", "--master_port", port, "-m", "groot1.train", |
| "--tiny", "--train", str(tmp / "data" / "train.jsonl"), "--val", str(tmp / "data" / "val.jsonl"), "--out", str(tmp / "run"), |
| "--batch-size", "4", "--epochs", "2", "--log-every", "1", "--eval-every", "100", "--eval-rows", "4", "--final-eval-rows", "8", |
| "--save-every", "3", "--warmup", "1", *extra] |
| return subprocess.run(cmd, capture_output=True, text=True, timeout=900, env={"CUDA_VISIBLE_DEVICES": "", "PATH": "/usr/bin:/bin", |
| "HOME": str(Path.home()), "TOKENIZERS_PARALLELISM": "false"}, cwd=ROOT) |
|
|
|
|
| @pytest.mark.slow |
| def test_ddp_epoch_boundary_and_resume(tmp_path): |
| subprocess.run([sys.executable, str(ROOT / "scripts" / "make_tiny_data.py"), "--out", str(tmp_path / "data"), "--n", "37"], check=True, |
| capture_output=True) |
| |
| first = _run(tmp_path, "--max-steps", "3", port="29531") |
| assert first.returncode == 0, first.stdout[-2000:] + first.stderr[-2000:] |
| state = json.loads((tmp_path / "run" / "ckpt" / "trainer_state.json").read_text()) |
| assert state["step"] == 3 |
| second = _run(tmp_path, "--max-steps", "6", "--resume", port="29532") |
| assert second.returncode == 0, second.stdout[-2000:] + second.stderr[-2000:] |
| assert '"resumed"' in second.stdout |
| logs = [json.loads(l) for l in (tmp_path / "run" / "train_log.jsonl").read_text().splitlines()] |
| steps = [r["step"] for r in logs if "step" in r and "loss" in r] |
| assert max(steps) == 6 and any(r.get("epoch") == 1 for r in logs if "loss" in r) |
| assert (tmp_path / "run" / "final" / "model.safetensors").exists() |
|
|
