| |
| |
| |
| |
| import os, sys, shutil, subprocess, glob |
|
|
| |
| 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) |
|
|
| |
| 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)") |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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)") |
|
|
| |
| 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.") |
|
|
| |
| 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" |
| ] |
|
|
| |
| 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() |
|
|
