Upload scripts/fast4gpu_boost.py with huggingface_hub

scripts/fast4gpu_boost.py

@@ -105,7 +105,7 @@ eng = eng.to(device)
 eng.reset_thought(B)
 
 try:
-    eng = torch.compile(eng, mode="reduce-overhead")
+    print("compile disabled for VRAM", flush=True)
     print(f"GPU {GPU}: compile OK", flush=True)
 except Exception as e:
     print(f"GPU {GPU}: compile skip: {e}", flush=True)
@@ -113,11 +113,22 @@ except Exception as e:
 opt = torch.optim.SGD(eng.parameters(), lr=LR, momentum=0.9)
 
 if not SHARD.exists():
-    raise FileNotFoundError(
-        f"Shard not found: {SHARD}\n"
-        "Place tokenized int64 1D shard at data/shard_gpu{id}.pt or set SHARD="
-    )
-tokens = torch.load(SHARD, weights_only=False).to(torch.int64)
+    # allow .npy memmap fallback for large phase2 shards
+    alt = Path(str(SHARD) + ".npy") if not str(SHARD).endswith(".npy") else SHARD
+    npy = SHARD if str(SHARD).endswith(".npy") else Path(str(SHARD).replace(".pt", ".npy"))
+    if not npy.exists():
+        raise FileNotFoundError(
+            f"Shard not found: {SHARD} (also tried {npy})"
+        )
+    SHARD = npy
+
+if str(SHARD).endswith(".npy"):
+    import numpy as np
+    tokens = torch.from_numpy(np.load(str(SHARD), mmap_mode="r")).to(torch.int64)
+    print(f"GPU {GPU}: memmap shard {SHARD} len={len(tokens):,}", flush=True)
+else:
+    tokens = torch.load(SHARD, weights_only=False).to(torch.int64)
+    print(f"GPU {GPU}: loaded shard {SHARD} len={len(tokens):,}", flush=True)
 step_tokens = B * SEQ
 
 start_token = int(os.environ.get("START_TOKEN", "0"))