From f5954dc51c81300610c20dee90e851e0706d5720 Mon Sep 17 00:00:00 2001 From: zhenyulincs Date: Sun, 2 Aug 2026 23:53:57 -0700 Subject: [PATCH] chore(scripts): vast.ai sync + debug toolkit from the M11 offload audit Track the vast-compute tooling used throughout the offload audit and smoke sessions (force-added past the scripts/ ignore, same as the already-tracked run_smoke_*.sh): - sync_to_vast.sh rsync rlix + miles working trees to an instance (env-overridable paths/key), push tooling to /root, optional --setup one-shot env bootstrap - vast_setup.sh instance setup per docs/smoke-test-runbook.md (deps, model, datasets, torch_dist ckpt, patches) - vast_debug_run.sh on-instance debug runner (single|dual driver) - debug_pipeline.py editable/IDE-debuggable pipeline launcher - apply_sglang_patches.py runbook Step 8 SGLang compat patches, idempotent - audit_gpu_sampler.sh 1 Hz whole-GPU + per-process memory sampler - run_smoke_audit.sh single-pipeline low-host-RAM audit smoke - mock_offload_test.py / mock_sglang_offload_test.py / test_offload_simple.py standalone tms/engine offload repros No library code changes; scripts only. --- scripts/apply_sglang_patches.py | 108 +++++++++++ scripts/audit_gpu_sampler.sh | 30 +++ scripts/debug_pipeline.py | 287 ++++++++++++++++++++++++++++ scripts/mock_offload_test.py | 125 ++++++++++++ scripts/mock_sglang_offload_test.py | 71 +++++++ scripts/run_smoke_audit.sh | 97 ++++++++++ scripts/sync_to_vast.sh | 87 +++++++++ scripts/test_offload_simple.py | 37 ++++ scripts/vast_debug_run.sh | 160 ++++++++++++++++ scripts/vast_setup.sh | 47 +++++ 10 files changed, 1049 insertions(+) create mode 100644 scripts/apply_sglang_patches.py create mode 100644 scripts/audit_gpu_sampler.sh create mode 100644 scripts/debug_pipeline.py create mode 100644 scripts/mock_offload_test.py create mode 100644 scripts/mock_sglang_offload_test.py create mode 100644 scripts/run_smoke_audit.sh create mode 100755 scripts/sync_to_vast.sh create mode 100644 scripts/test_offload_simple.py create mode 100755 scripts/vast_debug_run.sh create mode 100644 scripts/vast_setup.sh diff --git a/scripts/apply_sglang_patches.py b/scripts/apply_sglang_patches.py new file mode 100644 index 0000000..8426892 --- /dev/null +++ b/scripts/apply_sglang_patches.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Apply runbook Step 8 SGLang compat patches to miles sglang_engine.py. + +Environment-only patches (new SGLang session-based weight-update API + +flush_cache fault tolerance). Idempotent: skips already-applied hunks. +""" +import sys + +PATH = "/root/miles/miles/backends/sglang_utils/sglang_engine.py" + +HUNKS = [ + # Patch 1a: expanded io_struct imports + ( + """ from sglang.srt.entrypoints.http_server import _global_state + from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput + from sglang.srt.utils import MultiprocessingSerializer +""", + """ from sglang.srt.entrypoints.http_server import _global_state + from sglang.srt.managers.io_struct import ( + BeginWeightUpdateReqInput, + EndWeightUpdateReqInput, + UpdateWeightsFromTensorReqInput, + ) + from sglang.srt.utils import MultiprocessingSerializer +""", + ), + # Patch 1b: flush_cache=False in UpdateWeightsFromTensorReqInput + ( + """ obj = UpdateWeightsFromTensorReqInput( + serialized_named_tensors=serialized_named_tensors, + load_format=None, + flush_cache=True, + ) +""", + """ obj = UpdateWeightsFromTensorReqInput( + serialized_named_tensors=serialized_named_tensors, + load_format=None, + flush_cache=False, + ) +""", + ), + # Patch 1c: begin/end weight-update session wrapping + ( + """ try: + success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( + obj, None + ) + except Exception as exc: # noqa: BLE001 +""", + """ try: + await _global_state.tokenizer_manager.begin_weight_update( + BeginWeightUpdateReqInput(), None + ) + success, message = await _global_state.tokenizer_manager.update_weights_from_tensor( + obj, None + ) + await _global_state.tokenizer_manager.end_weight_update( + EndWeightUpdateReqInput(), None + ) + except Exception as exc: # noqa: BLE001 +""", + ), + # Patch 2a: 400 retry inside flush_cache loop + ( + """ response = requests.get(f"http://{self.server_host}:{self.server_port}/flush_cache") + if response.status_code == 200: + break +""", + """ response = requests.get(f"http://{self.server_host}:{self.server_port}/flush_cache") + if response.status_code == 200: + break + if response.status_code == 400: + logger.info("flush_cache returned 400, retrying in 1s...") + time.sleep(1) + continue +""", + ), + # Patch 2b: timeout -> warning + ( + """ else: + raise TimeoutError("Timeout while flushing cache.") +""", + """ else: + logger.warning("flush_cache timed out after 60 attempts, proceeding anyway") +""", + ), +] + +src = open(PATH).read() +applied, skipped = 0, 0 +for old, new in HUNKS: + if new in src: + skipped += 1 + continue + if old not in src: + print(f"FATAL: hunk not found and not applied:\n{old[:120]}...") + sys.exit(1) + if src.count(old) != 1: + print(f"FATAL: hunk not unique ({src.count(old)} occurrences):\n{old[:120]}...") + sys.exit(1) + src = src.replace(old, new) + applied += 1 + +open(PATH, "w").write(src) +print(f"PATCH_OK applied={applied} skipped={skipped}") +import py_compile +py_compile.compile(PATH, doraise=True) +print("COMPILE_OK") diff --git a/scripts/audit_gpu_sampler.sh b/scripts/audit_gpu_sampler.sh new file mode 100644 index 0000000..7164bcb --- /dev/null +++ b/scripts/audit_gpu_sampler.sh @@ -0,0 +1,30 @@ +#!/bin/bash +# GPU memory sampler for the M11 offload audit. +# Emits one block per second to $OUT: +# T= GPU ,; ... (whole-GPU memory.used) +# T= APP gpu= pid= mem_mib= cmd= (per compute process) +OUT=${OUT:-/root/logs/gpu_samples.log} +INTERVAL=${INTERVAL:-1} + +# bus_id -> index map (bus ids in compute-apps output) +declare -A BUS2IDX +while IFS=, read -r idx bus; do + bus=$(echo "$bus" | tr -d ' ') + idx=$(echo "$idx" | tr -d ' ') + BUS2IDX[$bus]=$idx +done < <(nvidia-smi --query-gpu=index,pci.bus_id --format=csv,noheader) + +while true; do + ts=$(date +%s) + { + g=$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | tr '\n' ';' | tr -d ' ') + echo "T=$ts GPU $g" + nvidia-smi --query-compute-apps=gpu_bus_id,pid,used_memory --format=csv,noheader,nounits | + while IFS=, read -r bus pid mem; do + bus=$(echo "$bus" | tr -d ' '); pid=$(echo "$pid" | tr -d ' '); mem=$(echo "$mem" | tr -d ' ') + cmd=$(tr '\0' ' ' < "/proc/$pid/cmdline" 2>/dev/null | cut -c1-160) + echo "T=$ts APP gpu=${BUS2IDX[$bus]:-$bus} pid=$pid mem_mib=$mem cmd=$cmd" + done + } >> "$OUT" + sleep "$INTERVAL" +done diff --git a/scripts/debug_pipeline.py b/scripts/debug_pipeline.py new file mode 100644 index 0000000..76dab6e --- /dev/null +++ b/scripts/debug_pipeline.py @@ -0,0 +1,287 @@ +#!/usr/bin/env python3 +"""rlix training-pipeline debug launcher — a plain Python file you can open, +edit, and run (on the vast instance), including under an IDE debugger. + +Usage on the instance: + python /root/debug_pipeline.py # run with Config below + python /root/debug_pipeline.py --dry-run # print argv/env, run nothing + python -m pdb /root/debug_pipeline.py # step into the driver + +Everything you would tweak lives in `Config` right below — edit and rerun. +The driver runs IN-PROCESS (via runpy), so breakpoints set in miles/rlix +code work when you launch this file from a debugger. +""" + +from __future__ import annotations + +import math +import os +import runpy +import subprocess +import sys +import time +from dataclasses import dataclass, field + + +# ====================================================================== +# Edit me +# ====================================================================== +@dataclass +class Config: + # --- what to run ------------------------------------------------- + mode: str = "single" # "single" or "dual" + num_rollout: int = 2 # training cycles + extra_args: list[str] = field(default_factory=list) # appended to the driver argv + + # --- GPU topology -------------------------------------------------- + # single mode: train GPUs = range(train_gpus), infer GPUs = range(infer_gpus) + # (derived in run_miles_rlix.py::_build_cluster_device_mappings — train may + # be a subset of infer, that overlap is the shrink/grant/expand handoff). + # train_gpus=1, infer_gpus=2 -> train=[0] infer=[0,1] (GPUs 2,3 idle) + # train_gpus=2, infer_gpus=4 -> train=[0,1] infer=[0,1,2,3] (all 4 GPUs, M11.1 topo) + train_gpus: int = 1 + infer_gpus: int = 2 + # dual mode: explicit per-pipeline GPU index lists (comma strings). + # defaults = M11.2 overlap topology, all 4 GPUs, shared infer on [1,2] + dual_p1_train: str = "0" + dual_p1_infer: str = "0,1,2" + dual_p2_train: str = "3" + dual_p2_infer: str = "1,2,3" + + # --- batch sizes ----------------------------------------------------- + # Megatron constraint: global_batch_size % (micro_batch(1) * DP) == 0, + # where DP = number of train GPUs. And each rollout must produce enough + # samples: rollout_batch_size * n_samples_per_prompt >= global_batch_size. + # None -> auto-derived from the topology (smallest legal debug values). + global_batch_size: int | None = None # auto: = train GPUs (dual: per-pipeline train GPUs) + rollout_batch_size: int | None = None # auto: = global_batch_size / n_samples_per_prompt + n_samples_per_prompt: int = 1 + + # --- memory offload knobs ---------------------------------------- + # "auto": preload where it is safe (non-Blackwell, or Blackwell with + # cu13+ torch wheels), torch on Blackwell + cu12.x (the tms + # 0.0.9 preload SIGSEGV combo — RTX 50xx / RTX PRO 6000 / B100 + # on the fork-baseline image). + tms_hook_mode: str = "auto" # "auto" | "preload" | "torch" + residual_threshold_gb: float | None = None # None -> code default (13.0 pre-#31, 7.0 after) + + # --- paths (instance layout) ------------------------------------- + miles: str = "/root/miles" + rlix: str = "/root/rlix" + megatron: str = "/root/Megatron-LM" + model_hf: str = "/root/Qwen2.5-0.5B" + model_torch_dist: str = "/root/Qwen2.5-0.5B_torch_dist" + model_miles: str = "/root/Qwen2.5-0.5B_miles/" + prompt_data: str = "/root/dapo-math-17k/dapo-math-17k.jsonl" + eval_data: str = "/root/aime-2024/aime-2024.jsonl" + + # --- runtime ------------------------------------------------------ + restart_ray: bool = True # kill + restart the local ray head first + num_gpus: int = 4 + # Mirror all output (incl. ray/Megatron C-level writes) to this file + # while still printing to the terminal. Previous log is rotated to + # .prev.log at launch. Set to None to disable. + log_file: str | None = "/root/logs/run.log" + + +CFG = Config() +# ====================================================================== + + +def build_env(cfg: Config) -> None: + """Set process env (inherited by ray workers via raylet).""" + os.environ["PYTHONPATH"] = f"{cfg.miles}:{cfg.rlix}:{cfg.megatron}" + os.environ["RLIX_CONTROL_PLANE"] = "rlix" + os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" + nvlink = subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout.count("NVLink") + os.environ["NCCL_NVLS_ENABLE"] = "1" if nvlink > 0 else "0" + os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + os.environ["MILES_SKIP_NODE_PG_PIN"] = "1" + # This image's ray ships an OTel metrics exporter whose background gRPC + # thread getenv()s while actor startup setenv()s -> glibc environ race -> + # SIGSEGV during actor creation (seen 2026-07-09 on MilesModelUpdateService). + # Metrics are useless for debugging; kill the racing thread at the source. + os.environ["RAY_enable_metrics_collection"] = "false" + hook = cfg.tms_hook_mode + if hook == "auto": + import torch + + cc_major = torch.cuda.get_device_capability()[0] if torch.cuda.is_available() else 0 + cuda = torch.version.cuda or "0" + blackwell_pre_cu13 = cc_major >= 10 and int(cuda.split(".")[0]) < 13 + hook = "torch" if blackwell_pre_cu13 else "preload" + print(f"tms hook auto-selected: {hook} (cc_major={cc_major}, cuda={cuda})") + os.environ["MILES_TMS_HOOK_MODE"] = hook + if hook == "preload": + # Guard escape hatch — needed for preload on Blackwell+cu13 (segfault + # verified gone there, 2026-07-05 audit). NOT set in torch mode: on + # Blackwell + cu12.x the guard's raise is protecting you from a real + # SIGSEGV; bypassing it would crash build_cpu_bucket_cache. + os.environ["MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL"] = "1" + if cfg.residual_threshold_gb is not None: + os.environ["MILES_MAX_RESIDUAL_GPU_MEM_GB"] = str(cfg.residual_threshold_gb) + if cfg.mode == "dual": + os.environ["MILES_INIT_DEFER_ADD_WORKER"] = "1" + os.environ["MILES_DUAL_P1_TRAIN"] = cfg.dual_p1_train + os.environ["MILES_DUAL_P1_INFER"] = cfg.dual_p1_infer + os.environ["MILES_DUAL_P2_TRAIN"] = cfg.dual_p2_train + os.environ["MILES_DUAL_P2_INFER"] = cfg.dual_p2_infer + for p in (cfg.miles, cfg.rlix, cfg.megatron): + if p not in sys.path: + sys.path.insert(0, p) + + +def model_args(cfg: Config) -> list[str]: + """MODEL_ARGS from miles' qwen2.5-0.5B.sh (single source of truth).""" + out = subprocess.run( + ["bash", "-c", f'source {cfg.miles}/scripts/models/qwen2.5-0.5B.sh && printf "%s\\n" "${{MODEL_ARGS[@]}}"'], + capture_output=True, text=True, check=True, + ) + return [line for line in out.stdout.splitlines() if line] + + +def build_argv(cfg: Config) -> tuple[str, list[str]]: + driver = os.path.join( + cfg.miles, "examples/rlix", + "run_miles_dual.py" if cfg.mode == "dual" else "run_miles_rlix.py", + ) + # Resolve batch sizes against the topology. Each dual pipeline is an + # INDEPENDENT Megatron world: its DP = len(its own train mapping) + # (P1=[0] -> DP=1, P2=[3] -> DP=1 — the two never form one DP=2 group). + # Both pipelines receive the SAME global_batch_size, so it must be + # divisible by BOTH pipelines' DP -> use the LCM (max is wrong: with + # train pools of 3 and 2 GPUs, global=3 divides 3 but not 2). + def _n(mapping: str) -> int: + return len([x for x in mapping.split(",") if x.strip()]) + + if cfg.mode == "dual": + dp = math.lcm(_n(cfg.dual_p1_train), _n(cfg.dual_p2_train)) + # CLI topology args must AGREE with the dual mappings: the driver + # overrides them per pipeline when the MILES_DUAL_* envs are present, + # but if the envs ever go missing (e.g. driver launched directly), + # the DISJOINT fallback uses the CLI values — a CLI/mapping mismatch + # is exactly what produced the "DP=2 vs global=1" Megatron assert. + cli_train_gpus = max(_n(cfg.dual_p1_train), _n(cfg.dual_p2_train)) + cli_infer_gpus = max(_n(cfg.dual_p1_infer), _n(cfg.dual_p2_infer)) + else: + dp = cfg.train_gpus + cli_train_gpus = cfg.train_gpus + cli_infer_gpus = cfg.infer_gpus + global_bs = cfg.global_batch_size if cfg.global_batch_size is not None else dp + if global_bs % dp != 0: + raise SystemExit( + f"Config error: global_batch_size={global_bs} not divisible by " + f"micro_batch(1) * DP({dp}) — Megatron will assert. Pick a multiple of {dp}." + ) + rollout_bs = ( + cfg.rollout_batch_size + if cfg.rollout_batch_size is not None + else max(1, -(-global_bs // cfg.n_samples_per_prompt)) # ceil div + ) + if rollout_bs * cfg.n_samples_per_prompt < global_bs: + raise SystemExit( + f"Config error: rollout_batch_size({rollout_bs}) * n_samples_per_prompt" + f"({cfg.n_samples_per_prompt}) = {rollout_bs * cfg.n_samples_per_prompt} " + f"< global_batch_size({global_bs}) — a train step would starve for samples." + ) + + argv = [ + driver, + *model_args(cfg), + "--hf-checkpoint", cfg.model_hf, + "--ref-load", cfg.model_torch_dist, + "--load", cfg.model_miles, + "--save", "", + "--save-interval", "100", "--eval-interval", "100", + "--eval-prompt-data", "aime", cfg.eval_data, + "--n-samples-per-eval-prompt", "1", "--eval-max-response-len", "1024", "--eval-top-p", "1", + "--prompt-data", cfg.prompt_data, + "--input-key", "prompt", "--label-key", "label", "--apply-chat-template", "--rollout-shuffle", + "--rm-type", "deepscaler", + "--num-rollout", str(cfg.num_rollout), + "--rollout-batch-size", str(rollout_bs), "--n-samples-per-prompt", str(cfg.n_samples_per_prompt), + "--rollout-max-response-len", "256", "--rollout-temperature", "1", + "--global-batch-size", str(global_bs), "--balance-data", + "--tensor-model-parallel-size", "1", "--pipeline-model-parallel-size", "1", + "--context-parallel-size", "1", + "--advantage-estimator", "grpo", "--use-kl-loss", "--kl-loss-coef", "0.0", + "--kl-loss-type", "low_var_kl", "--eps-clip", "0.2", "--eps-clip-high", "0.28", + "--optimizer", "adam", "--lr", "1e-6", "--lr-decay-style", "constant", + "--weight-decay", "0.1", "--adam-beta1", "0.9", "--adam-beta2", "0.98", + "--use-dynamic-batch-size", "--max-tokens-per-gpu", "512", + "--sglang-mem-fraction-static", "0.90", + "--rollout-num-gpus", str(cli_infer_gpus), "--rollout-num-gpus-per-engine", "1", + "--use-miles-router", + "--rollout-function-path", "examples.fully_async.fully_async_rollout.generate_rollout_fully_async", + "--offload-train", "--offload-rollout", + "--moe-router-topk", "0", + "--model-update-transport", "cpu_serialize", + "--num-gpus-per-node", str(cfg.num_gpus), + "--actor-num-nodes", "1", "--actor-num-gpus-per-node", str(cli_train_gpus), + "--attention-dropout", "0.0", "--hidden-dropout", "0.0", + "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", + "--attention-backend", "flash", + *cfg.extra_args, + ] + return driver, argv + + +def restart_ray(cfg: Config) -> None: + subprocess.run(["ray", "stop", "--force"], capture_output=True) + for pat in ("sglang", "raylet", "gcs_server", "ray::", "run_miles"): + subprocess.run(["pkill", "-9", "-f", pat], capture_output=True) + time.sleep(5) + subprocess.run(["bash", "-c", "rm -rf /tmp/ray /tmp/raylet* /tmp/plasma*"], capture_output=True) + time.sleep(2) + subprocess.run( + ["ray", "start", "--head", "--node-ip-address", "127.0.0.1", + "--num-gpus", str(cfg.num_gpus), "--disable-usage-stats", + "--dashboard-host", "0.0.0.0", "--dashboard-port", "8265"], + check=True, + ) + + +def main() -> None: + cfg = CFG + dry = "--dry-run" in sys.argv + + build_env(cfg) + driver, argv = build_argv(cfg) + + print(f"mode={cfg.mode} hook={cfg.tms_hook_mode} " + f"threshold={cfg.residual_threshold_gb or ''} num_rollout={cfg.num_rollout}") + print(f"driver: {driver}") + if dry: + print("argv:") + for a in argv[1:]: + print(f" {a}") + return + + # Raise fd limit (two pipelines + SGLang subprocesses exhaust the 1024 default). + import resource + resource.setrlimit(resource.RLIMIT_NOFILE, (65536, 65536)) + + if cfg.log_file: + # fd-level tee: dup stdout/stderr through `tee` so C-level writes + # (ray workers' forwarded logs, Megatron banners) land in the file + # too — sys.stdout redirection alone would miss them. + os.makedirs(os.path.dirname(cfg.log_file), exist_ok=True) + if os.path.exists(cfg.log_file): + os.replace(cfg.log_file, cfg.log_file.replace(".log", "") + ".prev.log") + tee = subprocess.Popen(["tee", cfg.log_file], stdin=subprocess.PIPE) + os.dup2(tee.stdin.fileno(), 1) + os.dup2(tee.stdin.fileno(), 2) + print(f"logging to {cfg.log_file} (previous run -> .prev.log)") + + if cfg.restart_ray: + restart_ray(cfg) + + os.chdir("/root") # ray workers inherit CWD; keep it neutral + sys.argv = argv + # In-process execution: your debugger's breakpoints inside miles/rlix fire. + runpy.run_path(driver, run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/scripts/mock_offload_test.py b/scripts/mock_offload_test.py new file mode 100644 index 0000000..94e0c06 --- /dev/null +++ b/scripts/mock_offload_test.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Minimal memory-offload mock test (~30 s), v2. + +Mirrors MegatronTrainRayActor.sleep()/wake_up() mechanics +(miles/backends/megatron_utils/actor.py:212-258): + - torch_memory_saver.hook_mode set BEFORE first tms call (tms-fixes.md #4) + - allocate tensors inside a tms region + - clear_memory equivalent (empty_cache) BEFORE pause, like actor.sleep() + - pause -> physical pages unmapped; resume -> remapped +Two variants: plain region (data may be discarded) and +enable_cpu_backup=True region (data must survive, Megatron-style). + +Whole-GPU nvidia-smi is the source of truth: on vast containers +--query-compute-apps PIDs live in the HOST pid namespace, so per-process +attribution is expected to fail here (same fail-open path as +miles/utils/gpu_probe.py). + +Run: MILES_TMS_HOOK_MODE=torch CUDA_VISIBLE_DEVICES=0 python mock_offload_test.py [alloc_gib] +""" +import os +import subprocess +import sys +import time + +ALLOC_GIB = float(sys.argv[1]) if len(sys.argv) > 1 else 4.0 +PID = os.getpid() + + +def smi(query, extra): + return subprocess.check_output( + ["nvidia-smi", f"--query-{query}={extra}", "--format=csv,noheader,nounits"], + text=True, + ).strip() + + +def gpu0_used(): + line = smi("gpu", "index,memory.used").splitlines()[0] + return int(line.split(",")[1]) + + +def report(label): + import torch + + torch.cuda.synchronize() + time.sleep(0.5) + used = gpu0_used() + print( + f"[{label:<22}] whole_gpu0={used:6d} MiB torch_reserved={int(torch.cuda.memory_reserved() / 2**20):6d} MiB", + flush=True, + ) + return used + + +def pid_namespace_check(): + lines = [l for l in smi("compute-apps", "pid,used_memory").splitlines() if l.strip()] + pids = [int(l.split(",")[0]) for l in lines] + print(f"compute-apps pids visible: {pids} (this pid={PID}, match={PID in pids})", flush=True) + + +def run_variant(torch_memory_saver, torch, cpu_backup): + name = "cpu_backup" if cpu_backup else "plain" + n_elem_half = int(ALLOC_GIB * 2**30 / 2 / 2) # bf16, 2 tensors + kwargs = {"enable_cpu_backup": True} if cpu_backup else {} + with torch_memory_saver.region(tag="default", **kwargs): + tensors = [torch.empty(n_elem_half, dtype=torch.bfloat16, device="cuda") for _ in range(2)] + for i, t in enumerate(tensors): + t.fill_(float(i + 1)) + # checksum WITHOUT big temporaries: sum a small slice per tensor + before = [float(t[:1024].float().sum()) for t in tensors] + alloc = report(f"{name}:after_alloc") + + # actor.sleep() equivalent: clear_memory (empty_cache) then pause + torch.cuda.empty_cache() + report(f"{name}:after_empty_cache") + torch_memory_saver.pause(tag=None) + paused = report(f"{name}:after_pause") + + torch_memory_saver.resume(tag=None) + resumed = report(f"{name}:after_resume") + after = [float(t[:1024].float().sum()) for t in tensors] + data_ok = before == after + + released = alloc - paused + print( + f" -> {name}: released_by_pause={released} MiB " + f"(expect ~{int(ALLOC_GIB * 1024)}), data={'PRESERVED' if data_ok else 'DISCARDED'}", + flush=True, + ) + del tensors + torch.cuda.empty_cache() + return released, data_ok + + +def main(): + t0 = time.time() + mode = os.environ.get("MILES_TMS_HOOK_MODE", "torch") + + from torch_memory_saver import torch_memory_saver + + torch_memory_saver.hook_mode = mode + print(f"tms hook_mode={mode!r} alloc={ALLOC_GIB} GiB pid={PID}", flush=True) + + import torch + + torch.zeros(1, device="cuda") + base = report("baseline(cuda ctx)") + pid_namespace_check() + + rel_plain, ok_plain = run_variant(torch_memory_saver, torch, cpu_backup=False) + rel_backup, ok_backup = run_variant(torch_memory_saver, torch, cpu_backup=True) + + final = report("final(after cleanup)") + print("-" * 72, flush=True) + print(f"cuda context baseline : {base} MiB", flush=True) + print(f"irreducible residual at end : {final} MiB (context + allocator metadata)", flush=True) + print(f"elapsed: {time.time() - t0:.1f}s", flush=True) + + want = ALLOC_GIB * 1024 * 0.9 + ok = rel_plain > want and rel_backup > want and ok_backup + print(f"OFFLOAD_TEST_{'PASS' if ok else 'FAIL'}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/scripts/mock_sglang_offload_test.py b/scripts/mock_sglang_offload_test.py new file mode 100644 index 0000000..41c2b06 --- /dev/null +++ b/scripts/mock_sglang_offload_test.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""SGLang engine offload micro-test (~2-3 min). + +Launches one SGLang engine with enable_memory_saver=True (same flag +miles passes when --offload-rollout, see tms-fixes.md #1), then: + release_memory_occupation -> measure whole-GPU residual + resume_memory_occupation -> measure recovery +Run: CUDA_VISIBLE_DEVICES=0 python mock_sglang_offload_test.py +""" +import subprocess +import sys +import time + + +def gpu0_used(): + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=index,memory.used", "--format=csv,noheader,nounits"], + text=True, + ) + return int(out.strip().splitlines()[0].split(",")[1]) + + +def report(label): + time.sleep(1.0) + used = gpu0_used() + print(f"[{label:<18}] whole_gpu0={used:6d} MiB", flush=True) + return used + + +def main(): + t0 = time.time() + base = report("baseline") + + import sglang as sgl + + engine = sgl.Engine( + model_path="/root/Qwen2.5-0.5B", + mem_fraction_static=0.30, + enable_memory_saver=True, + disable_cuda_graph=False, + skip_server_warmup=True, + ) + loaded = report("engine_loaded") + + out = engine.generate("1+1=", {"max_new_tokens": 4, "temperature": 0}) + print(f"sanity generate: {out['text']!r}", flush=True) + after_gen = report("after_generate") + + engine.release_memory_occupation() + released = report("after_release") + + engine.resume_memory_occupation() + resumed = report("after_resume") + + engine.shutdown() + final = report("after_shutdown") + + print("-" * 60, flush=True) + print(f"engine footprint : {after_gen - base} MiB", flush=True) + print(f"freed by release : {after_gen - released} MiB", flush=True) + print(f"residual after release : {released} MiB (over baseline {released - base} MiB)", flush=True) + print(f"recovered by resume : {resumed - released} MiB", flush=True) + print(f"elapsed: {time.time() - t0:.1f}s", flush=True) + + ok = (after_gen - released) > (after_gen - base) * 0.5 and released - base < 3000 + print(f"SGLANG_OFFLOAD_TEST_{'PASS' if ok else 'FAIL'}", flush=True) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_smoke_audit.sh b/scripts/run_smoke_audit.sh new file mode 100644 index 0000000..e0a2249 --- /dev/null +++ b/scripts/run_smoke_audit.sh @@ -0,0 +1,97 @@ +#!/bin/bash +# M11 memory-offload AUDIT smoke — single pipeline, low host-RAM variant. +# +# Why single pipeline: this host has 15 GB RAM + 8 GB swap; the dual +# smoke (2 Megatron actors + 6 SGLang engines) peaks ~28 GB host RAM and +# Ray's memory monitor kills workers at 95%. One pipeline with +# train=[0] infer=[0,1] still exercises the audited path: the train GPU +# overlaps the infer pool, so every cycle runs +# shrink_engines -> _wait_for_overlap_engines_offloaded -> train wake_up +# -> train sleep -> expand_engines on GPU 0. +# +# Deltas vs run_smoke_e2e.sh: +# - MILES_SKIP_TMS_PAUSE NOT set: torch_memory_saver.pause() is the +# offload under audit; A5000 is Ampere sm_86 (pre-Blackwell), safe. +# - RAY_memory_monitor_refresh_ms=0: rely on the 8 GB swap instead of +# Ray's 95% OOM killer. +# - ray --object-store-memory 500MB: nothing large goes through plasma. +# - Tiny rollout/batch sizes (from run_smoke_dual.sh) to bound memory. + +set -e +ulimit -n 65536 +cd /root + +export PYTHONPATH=/root/miles:/root/rlix:/root/Megatron-LM +export RLIX_CONTROL_PLANE=rlix +export CUDA_DEVICE_MAX_CONNECTIONS=1 +NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + export NCCL_NVLS_ENABLE=1 +else + export NCCL_NVLS_ENABLE=0 +fi +export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 +export CUBLAS_WORKSPACE_CONFIG=:4096:8 +export MILES_TMS_HOOK_MODE=torch +export MILES_SKIP_NODE_PG_PIN=1 +export RAY_memory_monitor_refresh_ms=0 + +echo "=== branch heads ===" +( cd /root/miles && git rev-parse HEAD 2>/dev/null || echo "miles@(rsync)" ) +( cd /root/rlix && git rev-parse HEAD 2>/dev/null || echo "rlix@(rsync)" ) + +echo "=== env vars ===" +echo "RLIX_CONTROL_PLANE=$RLIX_CONTROL_PLANE NCCL_NVLS_ENABLE=$NCCL_NVLS_ENABLE MILES_TMS_HOOK_MODE=$MILES_TMS_HOOK_MODE MILES_SKIP_TMS_PAUSE=${MILES_SKIP_TMS_PAUSE:-}" +echo "AUDIT_T0=$(date +%s)" + +echo "=== ray cleanup + start ===" +ray stop --force >/dev/null 2>&1 || true +pkill -9 -f sglang 2>/dev/null || true +pkill -9 -f raylet 2>/dev/null || true +pkill -9 -f gcs_server 2>/dev/null || true +pkill -9 -f ray:: 2>/dev/null || true +pkill -9 -f run_miles 2>/dev/null || true +sleep 5 +rm -rf /tmp/ray /tmp/raylet* /tmp/plasma* 2>/dev/null || true +sleep 3 +ray start --head --node-ip-address 127.0.0.1 --num-gpus 4 \ + --object-store-memory 500000000 \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +echo "=== launching run_miles_rlix (audit: train=[0] infer=[0,1]) ===" +source /root/miles/scripts/models/qwen2.5-0.5B.sh + +python /root/miles/examples/rlix/run_miles_rlix.py \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint /root/Qwen2.5-0.5B \ + --ref-load /root/Qwen2.5-0.5B_torch_dist \ + --load /root/Qwen2.5-0.5B_miles/ \ + --save "" \ + --save-interval 100 --eval-interval 100 \ + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl \ + --n-samples-per-eval-prompt 1 --eval-max-response-len 1024 --eval-top-p 1 \ + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl \ + --input-key prompt --label-key label --apply-chat-template --rollout-shuffle \ + --rm-type deepscaler \ + --num-rollout 2 --rollout-batch-size 1 --n-samples-per-prompt 1 \ + --rollout-max-response-len 256 --rollout-temperature 1 \ + --global-batch-size 1 --balance-data \ + --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --advantage-estimator grpo --use-kl-loss --kl-loss-coef 0.0 \ + --kl-loss-type low_var_kl --eps-clip 0.2 --eps-clip-high 0.28 \ + --optimizer adam --lr 1e-6 --lr-decay-style constant \ + --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \ + --use-dynamic-batch-size --max-tokens-per-gpu 512 \ + --sglang-mem-fraction-static 0.30 \ + --rollout-num-gpus 2 --rollout-num-gpus-per-engine 1 \ + --use-miles-router \ + --rollout-function-path examples.fully_async.fully_async_rollout.generate_rollout_fully_async \ + --offload-train --offload-rollout \ + --moe-router-topk 0 \ + --model-update-transport cpu_serialize \ + --num-gpus-per-node 4 \ + --actor-num-nodes 1 --actor-num-gpus-per-node 1 \ + --attention-dropout 0.0 --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 \ + --attention-backend flash diff --git a/scripts/sync_to_vast.sh b/scripts/sync_to_vast.sh new file mode 100755 index 0000000..8c6dc08 --- /dev/null +++ b/scripts/sync_to_vast.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# Sync local rlix + miles working trees (and helper tooling) to a vast.ai instance. +# +# Usage: +# ./scripts/sync_to_vast.sh [ssh-host] [flags] +# +# Examples: +# ./scripts/sync_to_vast.sh 18599 # → root@ssh9.vast.ai:18599 +# ./scripts/sync_to_vast.sh 39469 ssh5.vast.ai # other gateway +# ./scripts/sync_to_vast.sh 18599 ssh9.vast.ai --delete --setup +# +# Flags: +# --delete mirror exactly (remove remote files not present locally) — use for +# a clean state; without it rsync only adds/updates. +# --setup after syncing, run the one-shot environment setup on the instance +# (deps, model, datasets, torch_dist checkpoint, SGLang patches). +# --dry-run show what would transfer, change nothing. +# +# What gets synced: +# $RLIX_LOCAL → /root/rlix (this repo) +# $MILES_LOCAL → /root/miles (miles repo) +# scripts/*.sh|*.py → /root/ (debug/mock/setup tooling) +# +# Override paths/targets via env: RLIX_LOCAL, MILES_LOCAL, SSH_KEY, +# RLIX_REMOTE (default /root/rlix), MILES_REMOTE (default /root/miles). + +set -euo pipefail + +PORT="${1:?usage: sync_to_vast.sh [ssh-host] [--delete] [--setup] [--dry-run]}" +shift +HOST="ssh3.vast.ai" +if [[ $# -gt 0 && "$1" != --* ]]; then HOST="$1"; shift; fi + +DELETE=""; SETUP=0; DRY="" +for arg in "$@"; do + case "$arg" in + --delete) DELETE="--delete" ;; + --setup) SETUP=1 ;; + --dry-run) DRY="--dry-run" ;; + *) echo "unknown flag: $arg" >&2; exit 2 ;; + esac +done + +SSH_KEY="${SSH_KEY:-$HOME/.ssh/general_private_key}" +RLIX_LOCAL="${RLIX_LOCAL:-$HOME/Library/CloudStorage/Dropbox/Python/rlix_miles}" +MILES_LOCAL="${MILES_LOCAL:-$HOME/Dropbox/Python/miles}" +RLIX_REMOTE="${RLIX_REMOTE:-/root/rlix}" +MILES_REMOTE="${MILES_REMOTE:-/root/miles}" + +SSH_OPTS=(-o StrictHostKeyChecking=no -i "$SSH_KEY" -p "$PORT") +RSYNC_SSH="ssh -o StrictHostKeyChecking=no -i $SSH_KEY -p $PORT" +EXCLUDES=(--exclude .git --exclude __pycache__ --exclude '*.pyc' --exclude .venv + --exclude node_modules --exclude videos/ --exclude plans/ --exclude design/ + --exclude joe/ --exclude '.DS_Store' --exclude 'wandb/' --exclude 'outputs/') + +echo "=== sync rlix: $RLIX_LOCAL -> root@$HOST:$RLIX_REMOTE" +rsync -az $DRY $DELETE "${EXCLUDES[@]}" -e "$RSYNC_SSH" "$RLIX_LOCAL/" "root@$HOST:$RLIX_REMOTE/" + +echo "=== sync miles: $MILES_LOCAL -> root@$HOST:$MILES_REMOTE" +rsync -az $DRY $DELETE "${EXCLUDES[@]}" -e "$RSYNC_SSH" "$MILES_LOCAL/" "root@$HOST:$MILES_REMOTE/" + +echo "=== sync tooling: scripts/*.{sh,py} -> root@$HOST:/root/" +rsync -az $DRY -e "$RSYNC_SSH" \ + "$RLIX_LOCAL/scripts/vast_setup.sh" \ + "$RLIX_LOCAL/scripts/apply_sglang_patches.py" \ + "$RLIX_LOCAL/scripts/audit_gpu_sampler.sh" \ + "$RLIX_LOCAL/scripts/mock_offload_test.py" \ + "$RLIX_LOCAL/scripts/mock_sglang_offload_test.py" \ + "$RLIX_LOCAL/scripts/test_offload_simple.py" \ + "$RLIX_LOCAL/scripts/vast_debug_run.sh" \ + "$RLIX_LOCAL/scripts/debug_pipeline.py" \ + "root@$HOST:/root/" + +if [[ $SETUP -eq 1 && -z "$DRY" ]]; then + echo "=== running one-shot environment setup on the instance (background, ~10 min)" + # vast_setup.sh checks out git branches over the synced trees only when they + # are git repos; on a plain rsync tree it just installs deps + model + data. + ssh "${SSH_OPTS[@]}" "root@$HOST" \ + 'chmod +x /root/vast_setup.sh /root/vast_debug_run.sh /root/audit_gpu_sampler.sh 2>/dev/null; + setsid nohup bash /root/vast_setup.sh >/dev/null 2>&1 < /dev/null & + echo "setup launched; follow with: tail -f /root/setup.log (done marker: SETUP_DONE)"' +else + ssh "${SSH_OPTS[@]}" "root@$HOST" \ + 'chmod +x /root/vast_debug_run.sh /root/audit_gpu_sampler.sh 2>/dev/null || true; echo remote-ready' +fi + +echo "=== done. debug run: ssh ${SSH_OPTS[*]} root@$HOST 'bash /root/vast_debug_run.sh single'" diff --git a/scripts/test_offload_simple.py b/scripts/test_offload_simple.py new file mode 100644 index 0000000..4dcdbc9 --- /dev/null +++ b/scripts/test_offload_simple.py @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +"""最简 GPU memory offload 测试 — 分配 2GB → offload → 恢复,~10 秒跑完。 + +用法(在 vast 实例上): + CUDA_VISIBLE_DEVICES=0 python test_offload_simple.py +""" +import subprocess +import torch +from torch_memory_saver import torch_memory_saver + +torch_memory_saver.hook_mode = "torch" # 必须在第一次 tms 调用之前设置 + + +def gpu_used(): + """当前 GPU 0 整卡已用显存 (MiB),问 nvidia-smi,不问 torch。""" + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], text=True + ) + return int(out.splitlines()[0]) + + +torch.zeros(1, device="cuda") # 初始化 CUDA context +print(f"1) 基线 (CUDA context) : {gpu_used():6d} MiB") + +with torch_memory_saver.region(tag="default"): + x = torch.ones(2 * 1024**3 // 2, dtype=torch.bfloat16, device="cuda") # 2 GiB +torch.cuda.synchronize() +print(f"2) 分配 2GB 之后 : {gpu_used():6d} MiB") + +torch.cuda.empty_cache() +torch_memory_saver.pause() # ← offload:物理显存被释放 +torch.cuda.synchronize() +print(f"3) offload (pause) 之后 : {gpu_used():6d} MiB ← 应该回到基线附近") + +torch_memory_saver.resume() # ← 恢复:物理显存重新映射 +torch.cuda.synchronize() +print(f"4) 恢复 (resume) 之后 : {gpu_used():6d} MiB ← 应该回到 2GB 水平") diff --git a/scripts/vast_debug_run.sh b/scripts/vast_debug_run.sh new file mode 100755 index 0000000..9f0f8d4 --- /dev/null +++ b/scripts/vast_debug_run.sh @@ -0,0 +1,160 @@ +#!/bin/bash +# rlix training-pipeline debug runner — run ON the vast instance. +# +# Usage: +# bash /root/vast_debug_run.sh [single|dual] [extra driver args...] +# +# Env knobs (all optional): +# HOOK=preload|torch tms hook mode (default preload) +# THRESH= MILES_MAX_RESIDUAL_GPU_MEM_GB override (default: code default) +# NUM_ROLLOUT= training cycles (default 2) +# SILENCE_LIMIT= watchdog: kill after this much log silence (default 300) +# RUN_LIMIT= watchdog: hard wall-clock cap (default 1800) +# SAMPLER=0|1 1 Hz whole-GPU + per-process memory sampler (default 1) +# MILES=/root/miles RLIX=/root/rlix LOG_DIR=/root/logs tree/log locations +# +# Examples: +# bash vast_debug_run.sh single # minimal 1-pipeline debug loop +# HOOK=torch THRESH=13 bash vast_debug_run.sh dual # rollback-mode dual smoke +# NUM_ROLLOUT=5 bash vast_debug_run.sh single --rollout-max-response-len 512 +# +# Topologies: single = train[0], infer[0,1] (self-overlap on GPU 0) +# dual = P1 train[0]/infer[0,1,2] + P2 train[3]/infer[1,2,3] (overlap [1,2]) +# Ends with EXIT_CODE= as the last line of $LOG_DIR/run.log and a result summary. + +set -uo pipefail + +MODE="${1:-single}"; shift || true +HOOK="${HOOK:-preload}" +NUM_ROLLOUT="${NUM_ROLLOUT:-2}" +SILENCE_LIMIT="${SILENCE_LIMIT:-300}" +RUN_LIMIT="${RUN_LIMIT:-1800}" +SAMPLER="${SAMPLER:-1}" +MILES="${MILES:-/root/miles}" +RLIX="${RLIX:-/root/rlix}" +LOG_DIR="${LOG_DIR:-/root/logs}" +LOG="$LOG_DIR/run.log" + +mkdir -p "$LOG_DIR" +for f in run gpu_samples; do + [ -f "$LOG_DIR/$f.log" ] && mv "$LOG_DIR/$f.log" "$LOG_DIR/$f.prev.log" +done + +# ---- sampler ---- +if [ "$SAMPLER" = "1" ] && [ -f /root/audit_gpu_sampler.sh ]; then + pkill -f "[a]udit_gpu_sampler" 2>/dev/null + OUT="$LOG_DIR/gpu_samples.log" setsid nohup bash /root/audit_gpu_sampler.sh >/dev/null 2>&1 < /dev/null & + echo "sampler pid=$! -> $LOG_DIR/gpu_samples.log" +fi + +# ---- env (mirrors scripts/run_smoke_dual.sh conventions) ---- +ulimit -n 65536 +cd /root +export PYTHONPATH="$MILES:$RLIX:/root/Megatron-LM" +export RLIX_CONTROL_PLANE=rlix +export CUDA_DEVICE_MAX_CONNECTIONS=1 +NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l) +[ "$NVLINK_COUNT" -gt 0 ] && export NCCL_NVLS_ENABLE=1 || export NCCL_NVLS_ENABLE=0 +export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 +export CUBLAS_WORKSPACE_CONFIG=:4096:8 +export MILES_SKIP_NODE_PG_PIN=1 +export MILES_TMS_HOOK_MODE="$HOOK" +export MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1 # harmless on cu13+/pre-Blackwell +[ -n "${THRESH:-}" ] && export MILES_MAX_RESIDUAL_GPU_MEM_GB="$THRESH" + +if [ "$MODE" = "dual" ]; then + export MILES_INIT_DEFER_ADD_WORKER=1 + export MILES_DUAL_P1_TRAIN=0 MILES_DUAL_P1_INFER=0,1,2 + export MILES_DUAL_P2_TRAIN=3 MILES_DUAL_P2_INFER=1,2,3 + DRIVER="$MILES/examples/rlix/run_miles_dual.py" + TOPO_ARGS=(--rollout-num-gpus 2 --actor-num-gpus-per-node 1) +else + DRIVER="$MILES/examples/rlix/run_miles_rlix.py" + TOPO_ARGS=(--rollout-num-gpus 2 --actor-num-gpus-per-node 1) +fi + +echo "=== debug run: mode=$MODE hook=$HOOK thresh=${THRESH:-} num_rollout=$NUM_ROLLOUT" +echo "=== heads: miles=$(git -C "$MILES" rev-parse --short HEAD 2>/dev/null || echo rsync) rlix=$(git -C "$RLIX" rev-parse --short HEAD 2>/dev/null || echo rsync)" + +# ---- ray restart ---- +ray stop --force >/dev/null 2>&1 || true +pkill -9 -f sglang 2>/dev/null; pkill -9 -f raylet 2>/dev/null +pkill -9 -f gcs_server 2>/dev/null; pkill -9 -f "ray::" 2>/dev/null +pkill -9 -f run_miles 2>/dev/null +sleep 5; rm -rf /tmp/ray /tmp/raylet* /tmp/plasma* 2>/dev/null; sleep 2 +ray start --head --node-ip-address 127.0.0.1 --num-gpus 4 \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 >/dev/null + +# ---- launch under watchdog ---- +source "$MILES/scripts/models/qwen2.5-0.5B.sh" +( +python "$DRIVER" \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint /root/Qwen2.5-0.5B \ + --ref-load /root/Qwen2.5-0.5B_torch_dist \ + --load /root/Qwen2.5-0.5B_miles/ \ + --save "" --save-interval 100 --eval-interval 100 \ + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl \ + --n-samples-per-eval-prompt 1 --eval-max-response-len 1024 --eval-top-p 1 \ + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl \ + --input-key prompt --label-key label --apply-chat-template --rollout-shuffle \ + --rm-type deepscaler \ + --num-rollout "$NUM_ROLLOUT" --rollout-batch-size 1 --n-samples-per-prompt 1 \ + --rollout-max-response-len 256 --rollout-temperature 1 \ + --global-batch-size 1 --balance-data \ + --tensor-model-parallel-size 1 --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --advantage-estimator grpo --use-kl-loss --kl-loss-coef 0.0 \ + --kl-loss-type low_var_kl --eps-clip 0.2 --eps-clip-high 0.28 \ + --optimizer adam --lr 1e-6 --lr-decay-style constant \ + --weight-decay 0.1 --adam-beta1 0.9 --adam-beta2 0.98 \ + --use-dynamic-batch-size --max-tokens-per-gpu 512 \ + --sglang-mem-fraction-static 0.30 \ + --rollout-num-gpus-per-engine 1 \ + "${TOPO_ARGS[@]}" \ + --use-miles-router \ + --rollout-function-path examples.fully_async.fully_async_rollout.generate_rollout_fully_async \ + --offload-train --offload-rollout \ + --moe-router-topk 0 \ + --model-update-transport cpu_serialize \ + --num-gpus-per-node 4 --actor-num-nodes 1 \ + --attention-dropout 0.0 --hidden-dropout 0.0 \ + --accumulate-allreduce-grads-in-fp32 --attention-softmax-in-fp32 \ + --attention-backend flash \ + "$@" +echo "EXIT_CODE=$?" +) >"$LOG" 2>&1 & +RUN_PID=$! + +START=$(date +%s); LAST_SIZE=0; LAST_CHANGE=$START +while kill -0 $RUN_PID 2>/dev/null; do + NOW=$(date +%s) + SIZE=$(stat -c%s "$LOG" 2>/dev/null || echo 0) + [ "$SIZE" != "$LAST_SIZE" ] && { LAST_SIZE=$SIZE; LAST_CHANGE=$NOW; } + if [ $((NOW - LAST_CHANGE)) -gt "$SILENCE_LIMIT" ] || [ $((NOW - START)) -gt "$RUN_LIMIT" ]; then + echo "=== WATCHDOG: silent=$((NOW-LAST_CHANGE))s elapsed=$((NOW-START))s — killing ===" | tee -a "$LOG" + kill -9 $RUN_PID 2>/dev/null + pkill -9 -f run_miles 2>/dev/null; pkill -9 -f sglang 2>/dev/null + ray stop --force >/dev/null 2>&1 + echo "EXIT_CODE=124" >> "$LOG" + break + fi + sleep 10 +done +wait $RUN_PID 2>/dev/null + +pkill -f "[a]udit_gpu_sampler" 2>/dev/null + +# ---- summary ---- +echo +echo "================ RESULT ================" +tail -1 "$LOG" +echo "tracebacks: $(grep -c Traceback "$LOG" 2>/dev/null || echo 0)" +echo "--- gate readings ---" +# new gate (PR#17+): "whole-GPU mem used max=..."; old gate: "OS-level GPU mem free min=..." +grep -E "whole-GPU mem used|OS-level GPU mem free" "$LOG" | sed 's/\x1b\[[0-9;]*m//g' | sed 's/.*INFO:rlix/INFO:rlix/' | tail -6 +echo "--- training loops ---" +grep -E "training loop complete|WATCHDOG" "$LOG" | sed 's/\x1b\[[0-9;]*m//g' | tail -4 +echo "--- errors (if any) ---" +grep -E "RuntimeError|OutOfMemoryError|ActorDiedError" "$LOG" | sed 's/\x1b\[[0-9;]*m//g' | head -4 +echo "full log: $LOG gpu samples: $LOG_DIR/gpu_samples.log" diff --git a/scripts/vast_setup.sh b/scripts/vast_setup.sh new file mode 100644 index 0000000..4b47c9f --- /dev/null +++ b/scripts/vast_setup.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# M11 offload-audit instance setup — follows docs/smoke-test-runbook.md +set -x +exec > /root/setup.log 2>&1 + +echo "=== STEP 3: miles branch ===" +cd /root/miles +git remote add rlops https://github.com/rlops/miles.git 2>/dev/null +git fetch rlops zhenyu/m11-mvp-test +git checkout -B zhenyu/m11-offload-audit rlops/zhenyu/m11-mvp-test +pip install -e . --no-deps + +echo "=== STEP 4: rlix clone ===" +cd /root +if [ ! -d /root/rlix ]; then git clone https://github.com/rlops/rlix.git; fi +cd /root/rlix +git fetch origin zhenyu/miles-mvp-e2e +git checkout -B zhenyu/m11-offload-audit origin/zhenyu/miles-mvp-e2e +pip install -e . --no-deps + +echo "=== STEP 5: ROLL ===" +python -c "import roll" 2>/dev/null || pip install "roll @ git+https://github.com/rlops/ROLL.git" --no-deps +pip list 2>/dev/null | grep -i -E '^(roll|tg4perfetto|sglang|torch|ray) ' +python -c "import tg4perfetto" 2>/dev/null || pip install tg4perfetto + +echo "=== STEP 6: model + datasets ===" +[ -f /root/Qwen2.5-0.5B/config.json ] || hf download Qwen/Qwen2.5-0.5B --local-dir /root/Qwen2.5-0.5B +[ -d /root/dapo-math-17k ] && [ -n "$(ls /root/dapo-math-17k 2>/dev/null)" ] || hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k +[ -d /root/aime-2024 ] && [ -n "$(ls /root/aime-2024 2>/dev/null)" ] || hf download --repo-type dataset zhuzilin/aime-2024 --local-dir /root/aime-2024 + +echo "=== STEP 7: checkpoint conversion ===" +if [ ! -d /root/Qwen2.5-0.5B_torch_dist ]; then + cd /root/miles + source scripts/models/qwen2.5-0.5B.sh + PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + "${MODEL_ARGS[@]}" \ + --hf-checkpoint /root/Qwen2.5-0.5B \ + --save /root/Qwen2.5-0.5B_torch_dist +fi + +echo "=== VERSIONS ===" +python -c "import sglang; print('sglang', sglang.__version__)" +python -c "import torch; print('torch', torch.__version__, 'cc', torch.cuda.get_device_capability(0) if torch.cuda.is_available() else 'nocuda')" +nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader + +mkdir -p /root/logs +echo "=== SETUP_DONE ==="