Upload code/train_sft.py with huggingface_hub

code/train_sft.py

@@ -1,12 +1,18 @@
 # ==============================================================================
-# 🚀 ViuAI Sarus-500M — SFT v16 Master Training Engine (OOM-Resistant Edition)
+# 🚀 ViuAI Sarus-500M — SFT v16 Universal Auto-Adaptive Training Engine
 # ==============================================================================
-# KEY IMPROVEMENTS:
-#   • OOM-PROOF: Default micro-batch set to 2/4 with Grad Accum (Effective Batch = 128)
-#   • ZERO MEMORY SPIKES: torch.cuda.empty_cache() & efficient gradient accumulation
-#   • DYNAMIC VOCAB LOSS: Loss dynamically adapts to shift_logits.size(-1)
-#   • ACCURATE SPEED: Padded tokens excluded from throughput metrics
-#   • 100% RESUME-READY: Full state dict preservation & Hugging Face Auto-Sync
+# UNIVERSAL HARDWARE SUPPORT (Zero Manual Tuning Needed!):
+#   • 👑 Ultra-Tier (H200 141GB, H100 80GB, GH200)       -> Batch 16 x Accum 8  = 128
+#   • 🚀 High-Tier  (RTX 5090 32GB, A100 40GB, A6000)     -> Batch 8  x Accum 16 = 128
+#   • 💎 Pro-Tier   (RTX 4090 24GB, RTX 3090 24GB, L4)    -> Batch 4  x Accum 32 = 128
+#   • ⚡ Entry-Tier (T4 16GB, V100 16GB, RTX 4080 16GB)   -> Batch 2  x Accum 64 = 128
+#   • 💻 Budget     (<12GB VRAM / CPU)                    -> Batch 1  x Accum 128 = 128
+#
+# KEY FEATURES:
+#   • 100% OOM-PROOF: Auto-probes VRAM, sets expandable segments & clears cache.
+#   • TARGET EFFECTIVE BATCH = 128 (Mathematical invariance across all hardware).
+#   • NATIVE BF16 / FP16 AUTO-SWITCH: Uses bfloat16 on Ampere/Ada/Hopper/Blackwell.
+#   • RESUME & HF AUTO-PUSH: Seamless checkpointing and Hub synchronization.
 # ==============================================================================
 
 import os
@@ -23,6 +29,9 @@ import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 from huggingface_hub import HfApi, hf_hub_download
 
+# Optimize CUDA allocator to prevent fragmentation
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
 # Fix output encoding for Windows & Cloud terminals
 if hasattr(sys.stdout, "reconfigure"):
     sys.stdout.reconfigure(encoding="utf-8", errors="replace")
@@ -36,7 +45,68 @@ if cur_dir not in sys.path:
 from model import Transformer, ModelArgs
 
 
-# --- 1. Memory-Mapped High Performance Dataset ---
+# ------------------------------------------------------------------------------
+# 1. Universal Hardware Prober & Auto-Tuner
+# ------------------------------------------------------------------------------
+def auto_profile_hardware():
+    """Auto-detects GPU model, VRAM capacity, compute capability and selects golden parameters."""
+    if not torch.cuda.is_available():
+        return {
+            "tier": "CPU", "device_name": "CPU", "vram_gb": 0.0,
+            "micro_batch": 1, "grad_accum": 128, "dtype": torch.float32,
+            "desc": "CPU fallback mode"
+        }
+        
+    props = torch.cuda.get_device_properties(0)
+    device_name = props.name
+    vram_gb = props.total_memory / (1024 ** 3)
+    major, minor = props.major, props.minor
+    bf16_supported = torch.cuda.is_bf16_supported()
+    dtype = torch.bfloat16 if bf16_supported else torch.float16
+
+    # Tier Classification for 2048 Context Length
+    if vram_gb >= 75:  # H200 (141GB), H100 (80GB), GH200, A100-80GB
+        tier = "Ultra-Tier"
+        micro_batch = 16
+        grad_accum = 8
+        desc = "NVIDIA Hopper / Datacenter Beast"
+    elif vram_gb >= 30:  # RTX 5090 (32GB), A100 (40GB), A6000 (48GB), RTX 6000 Ada
+        tier = "High-Tier"
+        micro_batch = 8
+        grad_accum = 16
+        desc = "NVIDIA Blackwell / High-End Workstation"
+    elif vram_gb >= 20:  # RTX 4090 (24GB), RTX 3090 (24GB), L4 (24GB), A10G (24GB)
+        tier = "Pro-Tier"
+        micro_batch = 4
+        grad_accum = 32
+        desc = "NVIDIA Ada Lovelace / Ampere Pro (24GB)"
+    elif vram_gb >= 12:  # T4 (16GB), V100 (16GB), RTX 4080 (16GB), RTX 4070Ti (12GB)
+        tier = "Entry-Tier"
+        micro_batch = 2
+        grad_accum = 64
+        desc = "Standard Cloud GPU / 16GB"
+    else:  # < 12GB VRAM
+        tier = "Budget-Tier"
+        micro_batch = 1
+        grad_accum = 128
+        desc = "Budget GPU (< 12GB)"
+
+    return {
+        "tier": tier,
+        "device_name": device_name,
+        "vram_gb": vram_gb,
+        "compute_cap": f"{major}.{minor}",
+        "micro_batch": micro_batch,
+        "grad_accum": grad_accum,
+        "effective_batch": micro_batch * grad_accum,
+        "dtype": dtype,
+        "desc": desc
+    }
+
+
+# ------------------------------------------------------------------------------
+# 2. Memory-Mapped High Performance Dataset
+# ------------------------------------------------------------------------------
 class SFTDataset(Dataset):
     def __init__(self, ids_path: str, labels_path: str, offsets_path: str, domains_path: str = None):
         self.tokens_mmap = np.load(ids_path, mmap_mode="r")
@@ -52,7 +122,6 @@ class SFTDataset(Dataset):
         start_idx = int(self.offsets[idx])
         end_idx = int(self.offsets[idx + 1])
         
-        # Load slices from memory-mapped files
         tokens = torch.from_numpy(self.tokens_mmap[start_idx:end_idx].astype(np.int64))
         labels = torch.from_numpy(self.labels_mmap[start_idx:end_idx].astype(np.int64))
         domain_id = int(self.domains[idx]) if self.domains is not None else 0
@@ -60,7 +129,9 @@ class SFTDataset(Dataset):
         return tokens, labels, domain_id
 
 
-# --- 2. Dynamic Padding & Strict 2048 Bound Collate Function ---
+# ------------------------------------------------------------------------------
+# 3. Dynamic Padding & Strict 2048 Bound Collate Function
+# ------------------------------------------------------------------------------
 def sft_collate_fn(batch, pad_token_id=64000, ignore_index=-100):
     inputs, labels, domain_ids = zip(*batch)
     
@@ -79,7 +150,9 @@ def sft_collate_fn(batch, pad_token_id=64000, ignore_index=-100):
     return batch_inputs, batch_labels, torch.tensor(domain_ids, dtype=torch.long)
 
 
-# --- 3. Cosine Learning Rate Schedule with Warmup ---
+# ------------------------------------------------------------------------------
+# 4. Cosine Learning Rate Schedule with Warmup
+# ------------------------------------------------------------------------------
 def get_lr(it, warmup_steps, total_steps, max_lr, min_lr):
     if it < warmup_steps:
         return max_lr * (it + 1) / max(1, warmup_steps)
@@ -90,35 +163,14 @@ def get_lr(it, warmup_steps, total_steps, max_lr, min_lr):
     return min_lr + coeff * (max_lr - min_lr)
 
 
-# --- 4. Auto-detect OOM-Safe Batch Size by Hardware VRAM ---
-def auto_detect_batch_params():
-    if not torch.cuda.is_available():
-        return 2, 64
-    vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
-    if vram_gb >= 100:     # H100, H200 (80GB-140GB)
-        return 16, 8       # 16 micro batch * 8 accum = 128 effective
-    elif vram_gb >= 40:    # A100-40GB/80GB, A6000
-        return 8, 16       # 8 micro batch * 16 accum = 128 effective
-    elif vram_gb >= 20:    # RTX 3090, RTX 4090, L4, A10G (24GB)
-        return 4, 32       # 4 micro batch * 32 accum = 128 effective (OOM-Proof!)
-    else:                  # T4, V100, RTX 4080 (16GB)
-        return 2, 64       # 2 micro batch * 64 accum = 128 effective (OOM-Proof!)
-
-
-# --- 5. Self-Contained Cloud Download Helper ---
+# ------------------------------------------------------------------------------
+# 5. Cloud Auto-Sync (Dataset & Checkpoint Auto-Download)
+# ------------------------------------------------------------------------------
 def is_valid_checkpoint(path: str) -> bool:
-    if not os.path.exists(path):
-        return False
-    if os.path.getsize(path) < 10 * 1024 * 1024:
-        return False
-    return True
+    return os.path.exists(path) and (os.path.getsize(path) >= 10 * 1024 * 1024)
 
 def is_valid_data_file(path: str) -> bool:
-    if not os.path.exists(path):
-        return False
-    if os.path.getsize(path) < 1000:
-        return False
-    return True
+    return os.path.exists(path) and (os.path.getsize(path) >= 1000)
 
 def ensure_cloud_data_and_checkpoint(version: str, data_dir: str, ckpt_path: str, token: str = None):
     stage_subfolder = f"sft_{version}"
@@ -167,17 +219,19 @@ def ensure_cloud_data_and_checkpoint(version: str, data_dir: str, ckpt_path: str
             print(f"⚠️ Error downloading base checkpoint: {e}")
 
 
-# --- 6. Main Training Runner ---
+# ------------------------------------------------------------------------------
+# 6. Main Universal Training Runner
+# ------------------------------------------------------------------------------
 def main():
-    auto_b, auto_ga = auto_detect_batch_params()
+    hw = auto_profile_hardware()
 
-    parser = argparse.ArgumentParser(description="ViuAI Sarus-500M — SFT v16 Master Training Engine")
+    parser = argparse.ArgumentParser(description="ViuAI Sarus-500M — SFT v16 Universal Auto-Adaptive Training Engine")
     parser.add_argument("--version", type=str, default="v16", help="Dataset/Checkpoint version: v16 (default)")
     parser.add_argument("--data_dir", type=str, default=None, help="Root directory containing tokenized_data")
     parser.add_argument("--init_ckpt", type=str, default=None, help="Initial checkpoint path")
     parser.add_argument("--output_dir", type=str, default=None, help="Directory to save checkpoints")
-    parser.add_argument("--batch_size", type=int, default=auto_b, help=f"Micro batch size (Auto-detected: {auto_b})")
-    parser.add_argument("--grad_accum", type=int, default=auto_ga, help=f"Gradient accumulation steps (Auto-detected: {auto_ga})")
+    parser.add_argument("--batch_size", type=int, default=None, help="Micro batch size (Auto-configured if omitted)")
+    parser.add_argument("--grad_accum", type=int, default=None, help="Gradient accumulation steps (Auto-configured if omitted)")
     parser.add_argument("--epochs", type=int, default=3, help="Number of training epochs (Default: 3 for SFT v16)")
     parser.add_argument("--max_lr", type=float, default=3.2e-5, help="Peak learning rate for Cosine Schedule")
     parser.add_argument("--min_lr", type=float, default=2.0e-6, help="Minimum learning rate")
@@ -189,6 +243,11 @@ def main():
     parser.add_argument("--hf_token", type=str, default=None, help="Hugging Face API token")
     args = parser.parse_args()
 
+    # Determine batch parameters: explicit user override OR auto-detected golden values
+    micro_b = args.batch_size if args.batch_size is not None else hw["micro_batch"]
+    grad_acc = args.grad_accum if args.grad_accum is not None else hw["grad_accum"]
+    eff_batch = micro_b * grad_acc
+
     root_dir = os.path.abspath(os.path.join(cur_dir, ".."))
     data_dir = args.data_dir or os.path.join(root_dir, "tokenized_data")
     output_dir = args.output_dir or os.path.join(root_dir, "sft_checkpoints", f"sft_{args.version}")
@@ -214,18 +273,15 @@ def main():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         
-    print("=" * 80)
-    print(f"🚀 ViuAI Sarus-500M — SFT {args.version.upper()} Unified Master Training Runner")
-    print(f"   • Device:            {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
-    print(f"   • Data Folder:       {stage_data_dir}")
-    print(f"   • Init Checkpoint:   {init_ckpt}")
-    print(f"   • Save Checkpoint:   {save_path}")
-    print(f"   • Micro Batch Size:  {args.batch_size} (Grad Accum {args.grad_accum} -> Effective Batch {args.batch_size * args.grad_accum})")
-    print(f"   • Epochs:            {args.epochs}")
-    print(f"   • Max / Min LR:      {args.max_lr} / {args.min_lr} (Cosine Schedule)")
-    print(f"   • Push to HF:        {args.push_to_hf}")
-    print(f"   • Resume Mode:       {args.resume}")
-    print("=" * 80)
+    print("=" * 85)
+    print(f"🚀 ViuAI Sarus-500M — SFT {args.version.upper()} Universal Auto-Adaptive Training Runner")
+    print(f"   • Hardware Tier:     {hw['tier']} ({hw['desc']})")
+    print(f"   • Device Name:       {hw['device_name']} | VRAM: {hw['vram_gb']:.2f} GB")
+    print(f"   • Auto-Tuned Batch:  Micro-Batch {micro_b} × Grad-Accum {grad_acc} = Effective Batch {eff_batch}")
+    print(f"   • Precision Mode:    {hw['dtype']}")
+    print(f"   • Checkpoint Path:   {save_path}")
+    print(f"   • Target Epochs:     {args.epochs}")
+    print("=" * 85)
 
     # 1. Load Data
     train_dataset = SFTDataset(
@@ -243,7 +299,7 @@ def main():
 
     train_loader = DataLoader(
         train_dataset,
-        batch_size=args.batch_size,
+        batch_size=micro_b,
         shuffle=True,
         collate_fn=sft_collate_fn,
         num_workers=2,
@@ -251,7 +307,7 @@ def main():
     )
     val_loader = DataLoader(
         val_dataset,
-        batch_size=args.batch_size,
+        batch_size=micro_b,
         shuffle=False,
         collate_fn=sft_collate_fn,
         num_workers=2,
@@ -318,10 +374,10 @@ def main():
     ]
     optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.max_lr, betas=(0.9, 0.95), eps=1e-8)
 
-    steps_per_epoch = math.ceil(len(train_loader) / args.grad_accum)
+    steps_per_epoch = math.ceil(len(train_loader) / grad_acc)
     total_steps = steps_per_epoch * args.epochs
     warmup_steps = max(10, int(total_steps * args.warmup_ratio))
-    print(f"📊 Training Steps: {steps_per_epoch} steps/epoch | Total: {total_steps} steps | Warmup: {warmup_steps} steps")
+    print(f"📊 Training Plan: {steps_per_epoch} steps/epoch | Total: {total_steps} steps | Warmup: {warmup_steps} steps")
 
     start_epoch = 1
     global_step = 0
@@ -339,12 +395,9 @@ def main():
 
     # Mixed Precision Setup
     if torch.cuda.is_available():
-        dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
-        autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=dtype)
+        autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=hw["dtype"])
     else:
-        dtype = torch.float32
         autocast_ctx = contextlib.nullcontext()
-    print(f"⚡ Mixed Precision: {dtype}")
 
     # Evaluation Helper
     @torch.no_grad()
@@ -379,9 +432,9 @@ def main():
     total_tokens_trained = 0
     model.train()
 
-    print("\n" + "=" * 80)
-    print(f"🏁 STARTING SFT {args.version.upper()} MASTER TRAINING")
-    print("=" * 80)
+    print("\n" + "=" * 85)
+    print(f"🏁 STARTING SFT {args.version.upper()} MASTER TRAINING (AUTO-ADAPTIVE GPU ENGINE)")
+    print("=" * 85)
 
     for epoch in range(start_epoch, args.epochs + 1):
         print(f"\n--- Epoch {epoch}/{args.epochs} ---")
@@ -405,14 +458,14 @@ def main():
                 shift_logits = logits[..., :-1, :].contiguous()
                 shift_labels = labels[..., 1:].contiguous()
                 loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
-                loss_scaled = loss / args.grad_accum
+                loss_scaled = loss / grad_acc
 
             loss_scaled.backward()
             accum_loss += loss.item()
             epoch_loss += loss.item()
             epoch_batches += 1
 
-            if (micro_idx + 1) % args.grad_accum == 0:
+            if (micro_idx + 1) % grad_acc == 0:
                 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                 
                 lr = get_lr(global_step, warmup_steps, total_steps, args.max_lr, args.min_lr)
@@ -423,7 +476,7 @@ def main():
                 optimizer.zero_grad(set_to_none=True)
                 global_step += 1
                 
-                step_avg_loss = accum_loss / args.grad_accum
+                step_avg_loss = accum_loss / grad_acc
                 accum_loss = 0.0
 
                 if global_step % 10 == 0 or global_step == 1:
@@ -448,7 +501,6 @@ def main():
                         best_val_loss = v_loss
                         print(f"  🏆 New Best Validation Loss: {best_val_loss:.4f}! Saving checkpoint...")
                     
-                    # Save checkpoint
                     save_payload = {
                         "model_state_dict": model.state_dict(),
                         "optimizer_state_dict": optimizer.state_dict(),
@@ -463,7 +515,7 @@ def main():
                     print(f"  💾 Saved checkpoint -> {save_path}\n")
 
         # End of Epoch Handling
-        if (micro_idx + 1) % args.grad_accum != 0:
+        if (micro_idx + 1) % grad_acc != 0:
             torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
             lr = get_lr(global_step, warmup_steps, total_steps, args.max_lr, args.min_lr)
             for param_group in optimizer.param_groups:
@@ -477,13 +529,13 @@ def main():
 
     # Final Evaluation & Save
     final_val_loss, final_val_ppl = evaluate()
-    print("\n" + "=" * 80)
+    print("\n" + "=" * 85)
     print(f"🎉 SFT {args.version.upper()} MASTER TRAINING COMPLETED!")
     print(f"   • Final Validation Loss: {final_val_loss:.4f}")
     print(f"   • Final Perplexity:      {final_val_ppl:.2f}")
     print(f"   • Total Active Tokens:   {total_tokens_trained:,}")
     print(f"   • Total Time Taken:      {(time.time() - start_time)/60:.2f} minutes")
-    print("=" * 80)
+    print("=" * 85)
 
     # Final Save
     final_payload = {