Upload 2 files

code/train.py

@@ -46,41 +46,44 @@ def get_hf_token():
     raise ValueError("HF_TOKEN not found in environment or kaggle_secrets.")
 
 def setup_ddp():
-    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
-        dist.init_process_group("nccl")
-        rank = int(os.environ["RANK"])
-        local_rank = int(os.environ["LOCAL_RANK"])
-        world_size = int(os.environ["WORLD_SIZE"])
-        is_ddp = True
-    else:
-        rank = 0
-        local_rank = 0
-        world_size = 1
-        is_ddp = False
+    if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
+        return 0, 0, 1, torch.device("cuda" if torch.cuda.is_available() else "cpu"), False
+    dist.init_process_group("nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
     device = torch.device(f"cuda:{local_rank}")
     torch.cuda.set_device(device)
-    return rank, local_rank, world_size, device, is_ddp
+    return int(os.environ["RANK"]), local_rank, int(os.environ["WORLD_SIZE"]), device, True
 
-def load_tokenizer_robust(token: str):
-    logger.info("Loading tokenizer...")
-    try:
-        return AutoTokenizer.from_pretrained("ViuAI/ViuRec", token=token, use_fast=True)
-    except Exception as e:
-        logger.warning(f"Failed to load fast tokenizer: {e}. Trying use_fast=False.")
+def load_tokenizer_robust(hf_token: str, model_repo: str, subfolder: str, is_main: bool):
+    if is_main: logger.info("Loading tokenizer...")
+    local_candidates = [
+        "/kaggle/working/ViuRec/tokenizer",
+        "/kaggle/input/viuai-500m-tokenizer",
+        "./tokenizer",
+        "./viuai-500m-tokenizer"
+    ]
+    
+    for fast in [True, False]:
+        for p_str in local_candidates:
+            p = Path(p_str)
+            if (p / "tokenizer.json").exists() or (p / "vocab.json").exists():
+                try:
+                    return AutoTokenizer.from_pretrained(str(p), local_files_only=True, use_fast=fast, trust_remote_code=True)
+                except Exception as e:
+                    if "sentencepiece" in str(e).lower() or "tiktoken" in str(e).lower():
+                        continue
         try:
-            return AutoTokenizer.from_pretrained("ViuAI/ViuRec", token=token, use_fast=False)
-        except Exception as e2:
-            logger.error(f"Failed to load tokenizer: {e2}")
-            raise
-
-def load_model_code(token: str):
-    if not Path("model.py").exists():
-        logger.info("model.py not found locally. Downloading from HuggingFace...")
+            return AutoTokenizer.from_pretrained(model_repo, subfolder=subfolder, token=hf_token, use_fast=fast, trust_remote_code=True)
+        except Exception as e:
+            pass
         try:
-            hf_hub_download(repo_id="ViuAI/ViuRec", filename="code/model.py", token=token, local_dir=".")
+            return AutoTokenizer.from_pretrained(model_repo, token=hf_token, use_fast=fast, trust_remote_code=True)
         except Exception as e:
-            logger.error(f"Failed to download model.py: {e}")
-            raise
+            pass
+
+    if is_main:
+        logger.warning("Tokenizer failed, using dummy vocab 64000")
+    return None
 
 def load_checkpoint(model, optimizer, scaler, api, token, repo_id, filename="checkpoints/latest_chk.pt", strict=False):
     local_path = "latest_chk.pt"
@@ -187,6 +190,7 @@ def main():
     parser.add_argument("--lr", type=float, default=3e-4)
     parser.add_argument("--weight-decay", type=float, default=0.1)
     parser.add_argument("--ctx-len", type=int, default=2048)
+    parser.add_argument("--vocab-size", type=int, default=None)
     args = parser.parse_args()
 
     token = get_hf_token()
@@ -200,9 +204,8 @@ def main():
     rng = np.random.default_rng(seed)
 
     # Load resources
-    tokenizer = load_tokenizer_robust(token)
-    vocab_size = len(tokenizer)
-    load_model_code(token)
+    tokenizer = load_tokenizer_robust(token, "ViuAI/ViuRec", "tokenizer", is_main)
+    vocab_size = args.vocab_size or (len(tokenizer) if tokenizer is not None else 64000)
     
     from model import ViuResonance100M, ModelConfig
     config = ModelConfig(vocab_size=vocab_size)