Text Generation
Transformers
Safetensors
OpenPeerLLM
English
Arabic
riemann_openpeer
amd
ntk
code
fsdp
rocm
coding
mi325x
decoder
code gen
directml
openpeer
ryzen-ai
reasoning
rewriting
multimodal
paraphrase
prerelease
ntk-trainer
growth-ladder
style-transfer
text-rewriting
riemann-computing
progressive-growth
custom_code
Eval Results (legacy)
File size: 3,839 Bytes
db3982a | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | from __future__ import annotations
import argparse
import json
import logging
import subprocess
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from riemann.paths import CHECKPOINT_ROOT
logger = logging.getLogger("riemann.supervise")
def current_step(checkpoint_dir: Path) -> int:
p = checkpoint_dir / "training_state.json"
if not p.exists():
return 0
try:
return int(json.loads(p.read_text(encoding="utf-8")).get("global_step", 0))
except Exception:
return 0
def main_gpu_present() -> bool:
try:
import torch_directml
if torch_directml.device_count() == 0:
return False
name = torch_directml.device_name(0)
return "880M" not in name
except Exception:
try:
import torch
return torch.cuda.is_available()
except Exception:
return False
def wait_for_gpu(poll_seconds: int) -> None:
waited = 0
while not main_gpu_present():
if waited % (poll_seconds * 6) == 0:
logger.warning(
"Main GPU not present; waiting for it to reconnect (%ds so far)", waited
)
time.sleep(poll_seconds)
waited += poll_seconds
if waited:
logger.info("Main GPU is back after %ds; resuming", waited)
time.sleep(10)
def main() -> None:
logging.basicConfig(
level=logging.INFO, format="%(asctime)s %(levelname)s [supervise] %(message)s"
)
parser = argparse.ArgumentParser()
parser.add_argument("--stage", required=True)
parser.add_argument("--max-steps", type=int, required=True)
parser.add_argument("--max-restarts", type=int, default=200)
parser.add_argument("--gpu-poll-seconds", type=int, default=15)
parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
parser.add_argument("--log-file", default=None)
parser.add_argument("pretrain_args", nargs=argparse.REMAINDER)
args = parser.parse_args()
checkpoint_dir = Path(args.checkpoint_root) / args.stage
script = Path(__file__).resolve().parent / "pretrain.py"
extra = [a for a in args.pretrain_args if a != "--"]
restarts = 0
while True:
step = current_step(checkpoint_dir)
if step >= args.max_steps:
logger.info("Target reached: %s at step %d", args.stage, step)
return
wait_for_gpu(args.gpu_poll_seconds)
cmd = [
sys.executable,
"-u",
str(script),
"--stage",
args.stage,
"--resume",
"--max-steps",
str(args.max_steps),
"--checkpoint-root",
args.checkpoint_root,
*extra,
]
logger.info("Launch %d from step %d: %s", restarts + 1, step, " ".join(cmd[1:]))
if args.log_file:
with open(args.log_file, "a", encoding="utf-8") as fh:
rc = subprocess.call(cmd, stdout=fh, stderr=subprocess.STDOUT)
else:
rc = subprocess.call(cmd)
new_step = current_step(checkpoint_dir)
if rc == 0 and new_step >= args.max_steps:
logger.info("Finished: %s at step %d", args.stage, new_step)
return
restarts += 1
logger.warning(
"pretrain.py exited with code %d at step %d (restart %d/%d)",
rc,
new_step,
restarts,
args.max_restarts,
)
if restarts >= args.max_restarts:
logger.error("Too many restarts; giving up at step %d", new_step)
sys.exit(1)
time.sleep(20)
if __name__ == "__main__":
main()
|