File size: 4,145 Bytes
9f69669 | 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 | # ==============================================================================
# VIU-MINI-DENSE-360M: RTX 5090 / 4090 CLOUD NOTEBOOK MASTER RUNNER
# Works out-of-the-box on: RunPod, Vast.ai, Colab Pro, Lambda, Paperspace, Local PC
# ==============================================================================
import os, sys, shutil, subprocess, glob
# 1. Environment & Auth Setup
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["PYTHONUNBUFFERED"] = "1"
os.environ["HF_TOKEN"] = os.environ.get("HF_TOKEN", "")
print("=" * 70)
print(" 🚀 ViuMini-Dense-360M: RTX 5090 High-Throughput Trainer")
print("=" * 70)
# 2. Hardware Audit
import torch
if not torch.cuda.is_available():
raise RuntimeError("CUDA is not available. Please ensure GPU runtime is enabled!")
gpu_name = torch.cuda.get_device_name(0)
props = torch.cuda.get_device_properties(0)
vram_gb = props.total_memory / (1024**3)
num_gpus = torch.cuda.device_count()
print(f"[GPU] Detected: {gpu_name} ({vram_gb:.1f} GB VRAM x {num_gpus} GPUs)")
# 3. Dependencies
print("\n[1/4] Ensuring required dependencies...")
subprocess.run([
sys.executable, "-m", "pip", "install", "-q", "--upgrade",
"transformers", "datasets", "tokenizers", "bitsandbytes>=0.43.0",
"pyarrow", "pyyaml", "tqdm", "huggingface_hub", "pandas"
], check=True)
# 4. Clone or Pull Latest Codebase
if os.path.exists("/workspace"):
repo_dir = "/workspace/ViuMini-Dense-360M"
elif os.path.exists("/content"):
repo_dir = "/content/ViuMini-Dense-360M"
elif os.path.exists("/kaggle/working"):
repo_dir = "/kaggle/working/ViuMini-Dense-360M"
else:
repo_dir = os.path.abspath("ViuMini-Dense-360M")
print(f"\n[2/4] Setting up repository at {repo_dir}...")
if not os.path.exists(os.path.join(repo_dir, ".git")):
subprocess.run(["git", "clone", "https://huggingface.co/ViuAI/ViuMini-Dense-360M", repo_dir], check=True)
else:
subprocess.run(["git", "-C", repo_dir, "pull"], check=False)
os.chdir(repo_dir)
# 5. Check Tokenizer
tok_path = os.path.join(repo_dir, "tokenizer", "outputs", "tokenizer.json")
if not os.path.exists(tok_path) or os.path.getsize(tok_path) < 100000:
from huggingface_hub import hf_hub_download
print("[Tok] Downloading 4.25 MB tokenizer.json from Hub...")
hf_hub_download(
repo_id="ViuAI/ViuMini-Dense-360M",
filename="tokenizer/outputs/tokenizer.json",
local_dir=repo_dir,
token=os.environ["HF_TOKEN"]
)
print(f"[Tok] Tokenizer verified ({os.path.getsize(tok_path):,} bytes)")
# 6. Check Latest Checkpoint on Hugging Face to Auto-Resume
print("\n[3/4] Checking checkpoint history on Hugging Face Hub...")
from huggingface_hub import HfApi
api = HfApi(token=os.environ["HF_TOKEN"])
remote_ckpts = []
try:
files = api.list_repo_files(repo_id="ViuAI/ViuMini-Dense-360M", repo_type="model")
for f in files:
if f.startswith("model/checkpoints/step_") and f.endswith(".pt"):
st = int(f.split("step_")[1].split(".pt")[0])
remote_ckpts.append((st, f))
remote_ckpts.sort(key=lambda x: x[0])
except Exception as e:
print(f"[HF] Note: {e}")
if remote_ckpts:
latest_step, latest_file = remote_ckpts[-1]
print(f"[Resume] Found latest checkpoint on Hub: Step {latest_step} ({latest_file})")
print(f"[Resume] Model will seamlessly continue from Step {latest_step} onwards!")
else:
print("[Resume] No previous checkpoint found. Starting from Step 0.")
# 7. Select Hardware Config
cfg_file = "model/configs/train_rtx5090.yaml"
if not os.path.exists(cfg_file):
cfg_file = "model/configs/train_rtx4090.yaml"
print(f"\n[4/4] Launching High-Throughput Training with {cfg_file}...")
cmd = [
sys.executable, "model/scripts/train.py",
"--config", cfg_file,
"--resume", "auto"
]
# Launch with unbuffered real-time stdout streaming (step 10, 20, 30...)
proc = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
universal_newlines=True
)
for line in iter(proc.stdout.readline, ''):
print(line, end='', flush=True)
proc.stdout.close()
proc.wait()
|