| """Full benchmarks of the REAP-pruned checkpoints with apex_moe.lmeval, the dense full-run recipe: |
| MATH-500, GPQA-Diamond, HumanEval+, MBPP+ with thinking and a 64k budget, MMLU-Pro without thinking |
| and a 16k budget; seeds 0, 31323, 67210 (the dense full runs' seeds); tag full-reap-official; no routes saved. Qwen3.6 (NVFP4) jobs take |
| one idle GPU, Qwen3.8 (FP8, TP 2) jobs any two idle GPUs (NVSwitch: any pair works), DeepSeek-V4.1-Flash |
| (TP 4) any four, with its full-run recipe (deep_gemm, batch 36, MMLU-Pro batch 128, |
| VLLM_USE_BREAKABLE_CUDAGRAPH=0); --gpus limits the queue to some GPUs; a model's jobs start once its config exists |
| (make_config.py). Jobs with an ok ledger row or a live process (a restarted queue) are skipped; a |
| failed job is retried once. Seed 0 of every model first, in order: while an earlier two-GPU job |
| waits for a second idle GPU, later jobs do not take the idle one. Prints every launch and exit. |
| |
| cd apex-moe && setsid nohup .venv/bin/python external/reap/experiments/apex/eval_queue.py \ |
| > data/logs/reap-official-eval-queue.log 2>&1 & |
| """ |
|
|
| import json |
| import os |
| import subprocess |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path("/localhome/mkleineg/dgx-moe/apex-moe") |
| TAG = "full-reap-official" |
| MODELS = { |
| "qwen3.6-35b-a3b-nvfp4-reapofficial-keep50": (1, 76), |
| "qwen3.6-35b-a3b-nvfp4-reapofficial-keep25": (1, 76), |
| "qwen3.8-flash-next-fp8-reapofficial-keep50": (2, 80), |
| "qwen3.8-flash-next-fp8-reapofficial-keep25": (2, 80), |
| "deepseek-v4.1-flash-reapofficial-keep50": (4, 36), |
| "deepseek-v4.1-flash-reapofficial-keep25": (4, 36), |
| } |
| DEEPSEEK = { |
| "sets": ["model.model_args.moe_backend=deep_gemm"], |
| "mmlu_batch": 128, |
| "env": {"VLLM_USE_BREAKABLE_CUDAGRAPH": "0"}, |
| } |
| TASKS = { |
| "minerva": "minerva_math500", |
| "gpqa": "gpqa_diamond_cot_zeroshot", |
| "humaneval": "humaneval_plus_think", |
| "mbpp": "mbpp_plus_think", |
| "mmlu": "mmlu_pro", |
| } |
| SEEDS = (0, 31323, 67210) |
|
|
|
|
| def gpu_memory() -> dict[int, int]: |
| out = subprocess.run( |
| ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader,nounits"], |
| capture_output=True, text=True, check=True, |
| ).stdout |
| return {int(i): int(m) for i, m in (line.split(",") for line in out.strip().splitlines())} |
|
|
|
|
| def done() -> set[tuple[str, int]]: |
| ok = set() |
| for line in (ROOT / "experiments/ledger.jsonl").open(): |
| row = json.loads(line) |
| if row.get("tag") == TAG and row.get("status") == "ok": |
| ok.add((row["run_id"].rsplit("-", 2)[0], row["seed"])) |
| return ok |
|
|
|
|
| def live() -> set[tuple[str, str, int]]: |
| """(model, task, seed) of every apex_moe.lmeval process of this tag already running.""" |
| out = set() |
| for proc in Path("/proc").iterdir(): |
| try: |
| argv = (proc / "cmdline").read_bytes().decode().split("\0") |
| except (OSError, ValueError): |
| continue |
| if "apex_moe.lmeval" not in argv or TAG not in argv: |
| continue |
| model = Path(argv[argv.index("--model") + 1]).stem |
| task = next(k for k, v in TASKS.items() if v == argv[argv.index("--tasks") + 1]) |
| out.add((model, task, int(argv[argv.index("--seed") + 1]))) |
| return out |
|
|
|
|
| def name_of(model: str, task: str) -> str: |
| return f"{model}-lmeval-{task}-docs20000-{'nothink-' if task == 'mmlu' else ''}dense" |
|
|
|
|
| def command(model: str, task: str, seed: int) -> list[str]: |
| batch = MODELS[model][1] |
| deepseek = model.startswith("deepseek") |
| mmlu_batch = DEEPSEEK["mmlu_batch"] if deepseek else 140 |
| cmd = [ |
| ".venv/bin/python", "-m", "apex_moe.lmeval", "--model", f"configs/models/{model}.yaml", |
| "--docs", "20000", "--seed", str(seed), "--tasks", TASKS[task], "--tag", TAG, |
| ] |
| if task == "mmlu": |
| cmd += [ |
| "--set", "model.model_args.enable_thinking=false", |
| "--set", "model.gen_kwargs.chat_template_kwargs.enable_thinking=false", |
| "--set", "model.gen_kwargs.max_gen_toks=16384", |
| "--set", "model.model_args.max_model_len=18432", |
| "--set", f"model.model_args.max_num_seqs={mmlu_batch}", |
| "--set", f"name={name_of(model, task)}", |
| "--note", f"official REAP pruned, full MMLU-Pro without thinking, batched {mmlu_batch}", |
| ] |
| else: |
| cmd += [ |
| "--set", "model.gen_kwargs.max_gen_toks=65536", |
| "--set", "model.model_args.max_model_len=69632", |
| "--set", f"model.model_args.max_num_seqs={batch}", |
| "--note", f"official REAP pruned, full test set, 64k budget, batched {batch}", |
| ] |
| if task in ("humaneval", "mbpp"): |
| cmd.append("--allow-code-execution") |
| if deepseek: |
| for setting in DEEPSEEK["sets"]: |
| cmd += ["--set", setting] |
| return cmd |
|
|
|
|
| def main() -> None: |
| import argparse |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--gpus", default="0,1,2,3,4,5,6,7", help="GPUs the queue may use") |
| allowed = {int(g) for g in parser.parse_args().gpus.split(",")} |
| os.chdir(ROOT) |
| jobs = [(m, t, s) for s in SEEDS for m in MODELS for t in TASKS] |
| running: dict[tuple, tuple[subprocess.Popen, tuple[int, ...]]] = {} |
| failures: dict[tuple, int] = {} |
| print(f"{time.strftime('%T')} queue: {len(jobs)} jobs", flush=True) |
| while True: |
| for job, (proc, gpus) in list(running.items()): |
| if proc.poll() is not None: |
| print(f"{time.strftime('%T')} exit {proc.returncode}: {job} on GPUs {gpus}", flush=True) |
| if proc.returncode != 0: |
| failures[job] = failures.get(job, 0) + 1 |
| del running[job] |
| finished = done() |
| elsewhere = live() - set(running) |
| pending = [ |
| j for j in jobs |
| if (name_of(j[0], j[1]), j[2]) not in finished and j not in running and failures.get(j, 0) < 2 |
| and j not in elsewhere |
| and (ROOT / f"configs/models/{j[0]}.yaml").exists() |
| ] |
| waiting = [j for j in jobs if not (ROOT / f"configs/models/{j[0]}.yaml").exists()] |
| if not pending and not running and not waiting and not elsewhere: |
| break |
| busy = {g for _, gpus in running.values() for g in gpus} |
| memory = gpu_memory() |
| idle = [g for g, used in sorted(memory.items()) if g in allowed and g not in busy and used < 2000] |
| for job in pending: |
| width = MODELS[job[0]][0] |
| gpus = tuple(idle[:width]) if len(idle) >= width else None |
| if gpus is None: |
| break |
| idle = [g for g in idle if g not in gpus] |
| env = os.environ | { |
| "PATH": f"{ROOT}/.venv/bin:{os.environ['PATH']}", |
| "OMP_NUM_THREADS": "8", |
| "CUDA_VISIBLE_DEVICES": ",".join(map(str, gpus)), |
| "APEX_SAVE_ROUTES": "0", |
| "HF_ALLOW_CODE_EVAL": "1", |
| } | (DEEPSEEK["env"] if job[0].startswith("deepseek") else {}) |
| log = ROOT / f"data/logs/reapoff-{job[0].split('-reapofficial-')[1]}-{job[0][:8]}-{job[1]}-s{job[2]}.log" |
| proc = subprocess.Popen( |
| command(*job), env=env, stdout=log.open("w"), stderr=subprocess.STDOUT, |
| stdin=subprocess.DEVNULL, start_new_session=True, |
| ) |
| running[job] = (proc, gpus) |
| where = ",".join(map(str, gpus)) |
| print(f"{time.strftime('%T')} GPUs {where}: {job} (PID {proc.pid}) -> {log.name}", flush=True) |
| time.sleep(20) |
| time.sleep(30) |
| print(f"{time.strftime('%T')} queue done; failed twice: {[j for j, n in failures.items() if n >= 2]}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
