File size: 4,683 Bytes
61ffaac | 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 | """Kaggle GPU job (GPU T4 x2): trains two Nirnay models at once, one per GPU, with train/train_gpu.py.
Input: the private dataset <user>/nirnay-train-data made by jobs/kaggle_gpu.py (code/, records_v31.jsonl,
records_ml.jsonl, v2_model_best.pt). Output in /kaggle/working: nirnay-v31-gpu/, nirnay-multilingual/, one log per
model and the training run logs. With a single GPU the two models train one after the other.
"""
import glob
import os
import shutil
import subprocess
import sys
import time
import zipfile
W = "/kaggle/working"
SRC = os.path.join(W, "src")
def find(name):
hits = sorted(glob.glob("/kaggle/input/**/" + name, recursive=True))
return hits[0] if hits else None
def stage_code():
script = find("train_gpu.py")
if script is None: # code uploaded as an archive that was not unpacked
for z in glob.glob("/kaggle/input/**/*.zip", recursive=True):
zipfile.ZipFile(z).extractall(os.path.join(W, "unzipped"))
script = (sorted(glob.glob(W + "/unzipped/**/train_gpu.py", recursive=True)) or [None])[0]
assert script, "train_gpu.py not found in the attached dataset"
root = os.path.dirname(os.path.dirname(script)) # .../code (holds nirnay/ and train/)
shutil.copytree(root, SRC, dirs_exist_ok=True)
return os.path.join(SRC, "train", "train_gpu.py")
COMMON = {"NIRNAY_FUSION": "deep", "NIRNAY_SPLIT": "8", "NIRNAY_BATCH": "64", "NIRNAY_WARMUP": "300",
"NIRNAY_EVAL_EVERY": "1000", "NIRNAY_MAX_STATE": "256", "NIRNAY_TRAIN_STATE_TOKENS": "256",
"NIRNAY_THREADS": "2", "HF_HUB_DISABLE_TELEMETRY": "1", "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"}
JOBS = [
{"name": "multilingual", "NIRNAY_RECORDS": "records_ml.jsonl", "NIRNAY_BACKBONE": "intfloat/multilingual-e5-base",
"NIRNAY_POOL": "mean", "NIRNAY_PREFIX": "query: ", "NIRNAY_EPOCHS": "2.5", "NIRNAY_LR": "2e-4",
"NIRNAY_ADAPTER_LR": "5e-5", "NIRNAY_RUN": "multilingual", "NIRNAY_EXPORT": W + "/nirnay-multilingual"},
{"name": "v31-gpu", "NIRNAY_RECORDS": "records_v31.jsonl", "NIRNAY_BACKBONE": "BAAI/bge-base-en-v1.5",
"NIRNAY_POOL": "cls", "NIRNAY_PREFIX": "none", "NIRNAY_EPOCHS": "2.0", "NIRNAY_LR": "1e-4",
"NIRNAY_ADAPTER_LR": "3e-5", "NIRNAY_INIT": "v2_model_best.pt", "NIRNAY_RUN": "v31-gpu",
"NIRNAY_EXPORT": W + "/nirnay-v31-gpu"},
]
def launch(job, gpu, script):
env = dict(os.environ, **COMMON)
for k, v in job.items():
if k.startswith("NIRNAY_"):
env[k] = (find(v) or v) if k in ("NIRNAY_RECORDS", "NIRNAY_INIT") else v
env["CUDA_VISIBLE_DEVICES"] = str(gpu)
log = open(os.path.join(W, job["name"] + ".log"), "w")
print("launch %s on GPU %s: records %s" % (job["name"], gpu, env["NIRNAY_RECORDS"]), flush=True)
return subprocess.Popen([sys.executable, "-u", script], env=env, stdout=log, stderr=subprocess.STDOUT, cwd=SRC), log
def last_line(name):
try:
lines = [l for l in open(os.path.join(W, name + ".log"), encoding="utf-8", errors="replace").read().splitlines() if l.strip()]
return lines[-1][:160] if lines else ""
except OSError:
return ""
def main():
print(subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True).stdout, flush=True)
import torch
ngpu = torch.cuda.device_count()
script = stage_code()
t0 = time.time()
queue = list(JOBS)
running = []
while queue or running:
while queue and len(running) < max(1, ngpu):
free = sorted(set(range(max(1, ngpu))) - {g for _, _, _, g in running})[0]
job = queue.pop(0)
p, log = launch(job, free, script)
running.append((job, p, log, free))
time.sleep(120)
for job, p, log, g in list(running):
state = "running" if p.poll() is None else "exit %s" % p.returncode
print("[%6.1f min] %-12s %-9s %s" % ((time.time() - t0) / 60, job["name"], state, last_line(job["name"])), flush=True)
if p.poll() is not None:
log.close()
running.remove((job, p, log, g))
for job in JOBS:
print("=" * 30, job["name"], "-> exported:", os.path.isdir(job["NIRNAY_EXPORT"]), flush=True)
print(open(os.path.join(W, job["name"] + ".log"), encoding="utf-8", errors="replace").read()[-4000:], flush=True)
for run in glob.glob(SRC + "/train/runs/*"): # keep the run logs and best checkpoints next to the models
shutil.copytree(run, os.path.join(W, "runs", os.path.basename(run)), dirs_exist_ok=True)
shutil.rmtree(SRC, ignore_errors=True)
shutil.rmtree(os.path.join(W, "unzipped"), ignore_errors=True)
if __name__ == "__main__":
main()
|