Update Kaggle Dual T4 runner with DDP 4x32 batch sizing

runners/KAGGLE_CONTINUATION_PRETRAIN_GPU.txt

@@ -1,26 +1,30 @@
 # ==============================================================================
-# 🚀 ViuAI Sarus-500M — Kaggle 2x T4 GPU Training Runner (Continuation V2)
+# 🚀 ViuAI Sarus-500M — Kaggle 2x T4 Dual-GPU Pretraining Runner (Continuation V2)
 # ==============================================================================
 # Instructions for Kaggle:
-# 1. In Kaggle Notebook Settings:
+# 1. In Kaggle Notebook Settings (Right Panel):
 #    • Accelerator: Select "GPU T4 x2" (Dual GPU)
-#    • Internet: ON
-# 2. Paste this entire script into a single cell and click "Run All".
-# 3. It will:
-#    • Download base checkpoint (Step 40,000) from ViuAI/ViuAI-500M.
-#    • Stream the new 10B token shards from ViuAI/viuai-500m-data/shards_v2/.
-#    • Train on both T4 GPUs in parallel using PyTorch DDP.
-#    • Auto-save and upload updated ckpt_latest.pt back to Hugging Face.
+#    • Internet: Always ON
+#    • Environment: Always use latest environment
+# 2. Paste this entire script into a single cell and click Run (Shift + Enter).
+# 3. What it does automatically:
+#    • Resumes seamlessly from latest HF checkpoint (Step 42,750+).
+#    • Uses PyTorch DDP across both T4 GPUs in parallel (world_size=2).
+#    • Effective Batch: MICRO_BATCH=4, GRAD_ACCUM=32 -> 524,288 tok/step.
+#    • Streams 200 shards (10B tokens) from ViuAI/viuai-500m-data/shards_v2/.
+#    • Background Checkpoint Uploader: Uploads ckpt_latest.pt every 250 steps to HF.
 # ==============================================================================
 
-!pip install -q torch transformers datasets tokenizers huggingface_hub numpy
+!pip install -q transformers datasets tokenizers huggingface_hub numpy
 
 import os, sys, shutil
+import torch
 from huggingface_hub import HfApi, hf_hub_download
 
-# Ensure unbuffered real-time logs and memory fragmentation protection
+# Real-time unbuffered logging & memory configuration
 os.environ["PYTHONUNBUFFERED"] = "1"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 
 HF_TOKEN = os.environ.get("HF_TOKEN") or ''.join(['hf_','ssyCVhuny','XxjGdqKp','VLPpkmWK','FrrMOIFbg'])
 os.environ["HF_TOKEN"] = HF_TOKEN
@@ -29,25 +33,47 @@ REPO_ID = "ViuAI/ViuAI-500M"
 DATA_REPO = "ViuAI/viuai-500m-data"
 
 print("=" * 80)
-print("🚀 ViuAI Sarus-500M — Launching Dual T4 GPU Training (Continuation V2)")
+print("🚀 ViuAI Sarus-500M — Launching Kaggle Dual T4 (2x 15GB) DDP Pretraining")
 print("=" * 80)
 
-# Continuation Pretraining Hyperparameters
-os.environ["TOTAL_STEPS"] = "60000"       # 20,000 continuation steps (~10B tokens)
+num_gpus = torch.cuda.device_count()
+print(f"🖥️ Detected GPUs: {num_gpus}x CUDA Device(s)")
+for i in range(num_gpus):
+    name = torch.cuda.get_device_name(i)
+    vram = torch.cuda.get_device_properties(i).total_memory / (1024**3)
+    print(f"   [GPU {i}] {name} ({vram:.1f}GB VRAM)")
+
+if num_gpus < 2:
+    print("⚠️ Warning: Less than 2 GPUs detected! Make sure 'GPU T4 x2' is selected in Kaggle settings.")
+
+# Hyperparameters for continuation up to 60,000 steps
+os.environ["TOTAL_STEPS"] = "60000"       # Full run to 60,000 steps (~10B tokens)
+os.environ["MAX_STEPS"] = "60000"         # Unrestricted run
 os.environ["MAX_LR"] = "1.0e-4"           # Calibrated continuation LR to protect base weights
 os.environ["MIN_LR"] = "1.0e-5"
 os.environ["WARMUP_STEPS"] = "500"
-os.environ["SAVE_EVERY"] = "250"          # Frequent checkpoint saves every ~20 mins for Kaggle safety
+os.environ["SAVE_EVERY"] = "250"          # Checkpoints upload in background thread (ZERO GPU pause)
 os.environ["EVAL_EVERY"] = "250"
-os.environ["MICRO_BATCH"] = "4"           # Perfectly sized for Kaggle Dual T4 (15GB VRAM)
+
+# Kaggle 2x T4 Optimal Configuration (524,288 tokens/step)
+# Formula: MICRO_BATCH(4) * GRAD_ACCUM(32) * WORLD_SIZE(2) * CONTEXT(2048) = 524,288 tokens/step
+os.environ["USE_CHECKPOINT"] = "1"        # Activation Checkpointing ON (~9.5GB / 15GB per GPU -> 0% OOM)
+os.environ["MICRO_BATCH"] = "4"
 os.environ["GRAD_ACCUM"] = "32"
 
+print(f"⚙️ Config: MICRO_BATCH=4, GRAD_ACCUM=32, WORLD_SIZE=2 -> 524,288 tok/step")
+
 # Fetch latest verified code and tokenizer directly from HF
 for f in ["code/config.py", "code/model.py", "code/train.py", "tokenizer/tokenizer.json"]:
     os.makedirs(os.path.dirname(f), exist_ok=True)
     print(f"📥 Downloading latest {f}...")
     p = hf_hub_download(repo_id=REPO_ID, filename=f, token=HF_TOKEN, force_download=True)
     shutil.copy(p, f)
+    # Also update working directory copy
+    work_copy = f"/kaggle/working/{f}"
+    if os.path.exists(os.path.dirname(work_copy)):
+        try: shutil.copy(p, work_copy)
+        except OSError: pass
 
-# Launch 2x GPU Distributed Training
+# Launch PyTorch DDP Distributed Training across both T4 GPUs
 !torchrun --nproc_per_node=2 code/train.py