Upload folder using huggingface_hub

train.py

@@ -86,6 +86,18 @@ def train(args):
     os.makedirs(ckpt_dir, exist_ok=True)
     cfg = build_cfg(args)
 
+    # ---- device resolution (cuda / xla / cpu) ----
+    if getattr(args, "device", "cuda") == "xla":
+        import torch_xla.core.xla_model as xm
+        device = xm.xla_device()
+        print(f"[train] device = TPU:XLA ({device})")
+    elif getattr(args, "device", "cuda") == "cpu":
+        device = torch.device("cpu")
+        print("[train] device = CPU")
+    else:
+        device = torch.device("cuda")
+        print(f"[train] device = {device}")
+
     tok = YKTokenizer.load(os.path.join(data_dir, "tokenizer.json"))
     arr, seq_len, vocab, n_tokens = load_data(data_dir)
     cfg["vocab_size"] = vocab
@@ -93,7 +105,7 @@ def train(args):
     print(f"[train] data n_tokens={n_tokens:,} seq_len={seq_len} vocab={vocab}")
     print(f"[train] model params = {sum(p.numel() for p in YKDiff(cfg).parameters())/1e6:.1f}M")
 
-    model = YKDiff(cfg).cuda()
+    model = YKDiff(cfg).to(device)
     n_params = sum(p.numel() for p in model.parameters())
     print(f"[train] allocated params = {n_params/1e6:.1f}M")
 
@@ -108,13 +120,13 @@ def train(args):
     ckpts = sorted([f for f in os.listdir(ckpt_dir) if f.endswith(".pt")])
     if ckpts and not args.fresh:
         path = os.path.join(ckpt_dir, ckpts[-1])
-        sd = torch.load(path, map_location="cuda")
+        sd = torch.load(path, map_location=device)
         model.load_state_dict(sd["model"]); optim.load_state_dict(sd["optim"])
         step0 = sd["step"]
         print(f"[train] resumed from {path} step={step0}")
 
     model.train()
-    amp = torch.cuda.amp.autocast(dtype=torch.bfloat16)
+    amp = torch.amp.autocast(device_type=device.type, dtype=torch.bfloat16)
     t0 = time.time()
     limit = args.hours * 3600.0
     step = step0
@@ -127,20 +139,20 @@ def train(args):
 
         optim.zero_grad(set_to_none=True)
         mode_ar = (torch.rand(1).item() < 0.5)
-        idx = sample_batch(arr, seq_len, args.batch).cuda()
+        idx = sample_batch(arr, seq_len, args.batch).to(device)
 
         with amp:
             if mode_ar:
-                m = torch.zeros(args.batch, dtype=torch.long, device="cuda")
+                m = torch.zeros(args.batch, dtype=torch.long, device=device)
                 logits = model(idx, m, t=None)                 # causal
                 loss = F.cross_entropy(
                     logits[:, :-1].reshape(-1, V),
                     idx[:, 1:].reshape(-1), ignore_index=pad_id)
                 mname = "AR"
             else:
-                m = torch.ones(args.batch, dtype=torch.long, device="cuda")
-                r = torch.rand(args.batch, device="cuda")       # per-sample ratio
-                is_mask = torch.rand(args.batch, seq_len, device="cuda") < r[:, None]
+                m = torch.ones(args.batch, dtype=torch.long, device=device)
+                r = torch.rand(args.batch, device=device)       # per-sample ratio
+                is_mask = torch.rand(args.batch, seq_len, device=device) < r[:, None]
                 not_pad = idx != pad_id
                 masked = idx.clone(); masked[is_mask] = mask_id
                 logits = model(masked, m, t=r)
@@ -154,6 +166,8 @@ def train(args):
         loss.backward()
         nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         optim.step()
+        if device.type == "xla":
+            xm.mark_step()
 
         running = running * 0.9 + float(loss.item()) * 0.1
         step += 1
@@ -186,6 +200,8 @@ if __name__ == "__main__":
     ap.add_argument("--log-every", type=int, default=25)
     ap.add_argument("--ckpt-every", type=int, default=500)
     ap.add_argument("--fresh", action="store_true")
+    ap.add_argument("--device", default="cuda", choices=["cuda", "xla", "cpu"],
+                    help="training device (cuda default; xla for TPU; cpu for tests)")
     ap.add_argument("--data-dir", default=None)
     ap.add_argument("--ckpt-dir", default=None)
     ap.add_argument("--hf-repo", default=None,