Upload train.py

code/train.py

@@ -72,15 +72,17 @@ def load_tokenizer_robust():
             except Exception:
                 pass
     try:
-        tok = AutoTokenizer.from_pretrained(MODEL_REPO, subfolder=TOKENIZER_SUBFOLDER, token=HF_TOKEN)
+        tok = AutoTokenizer.from_pretrained(MODEL_REPO, subfolder=TOKENIZER_SUBFOLDER, token=HF_TOKEN, use_fast=False)
         if IS_MAIN:
             print(f"✅ Tokenizer loaded from HF ({MODEL_REPO}/{TOKENIZER_SUBFOLDER})")
         return tok
     except Exception as e:
-        raise RuntimeError(f"Could not load tokenizer: {e}")
+        print(f"⚠️ Tokenizer bypassed.")
+        return None
 
 tokenizer = load_tokenizer_robust()
-VOCAB_SIZE = len(tokenizer)
+VOCAB_SIZE = len(tokenizer) if tokenizer is not None else 64000
+if IS_MAIN: print(f"ℹ️ VOCAB_SIZE set to {VOCAB_SIZE}")
 
 if IS_MAIN:
     create_repo(CKPT_REPO, repo_type="model", token=HF_TOKEN, exist_ok=True, private=True)
@@ -266,13 +268,9 @@ model = ViuResonance100M(vocab_size=VOCAB_SIZE, d_model=1024, n_layers=12, n_hea
 
 if IS_MAIN:
     log(f"Model params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M | world_size={world_size}")
-    log("Compiling model (if it fails, will fallback to eager)...")
+    log("Bypassing torch.compile (Disabled for Kaggle T4 compatibility)...")
 
-try:
-    model = torch.compile(model)
-except Exception as e:
-    if IS_MAIN:
-        log(f"⚠️ torch.compile failed, falling back to eager mode: {e}")
+# torch.compile removed to prevent Kaggle T4 IndentationError/freeze
 model = DDP(model, device_ids=[local_rank])
 
 optimizer = torch.optim.AdamW(model.parameters(), lr=MAX_LR, weight_decay=WEIGHT_DECAY, betas=(0.9, 0.95), fused=True)
@@ -288,6 +286,7 @@ val_shards = sorted(f for f in files if f.startswith("shards/") and f.endswith("
 
 shard_pool = ShardPool(train_shards, DATA_REPO, HF_TOKEN, f"{WORK}/train_shards_r{rank}", seed=rank + start_step)
 val_tokens = load_val_tokens(val_shards, DATA_REPO, HF_TOKEN, f"{WORK}/val_shards") if IS_MAIN else None
+dist.barrier() # Sync ranks after downloading validation data
 
 @torch.no_grad()
 def evaluate():