Upload train_300m.py with huggingface_hub

train_300m.py

@@ -256,7 +256,22 @@ def probe_max_batch(model, device, dtype, seq_len, vocab_size, max_search=4096,
             tmp_opt.zero_grad(set_to_none=True)
             hi = mid - 1
         except RuntimeError as e:
-            if "out of memory" in str(e).lower():
+            err_str = str(e).lower()
+            if ("out of memory" in err_str
+                    or "no valid triton configs" in err_str
+                    or "out of resource" in err_str
+                    or "inductorerror" in type(e).__name__.lower()):
+                try: del x, t, loss
+                except: pass
+                torch.cuda.empty_cache()
+                tmp_opt.zero_grad(set_to_none=True)
+                hi = mid - 1
+            else:
+                raise
+        except Exception as e:
+            err_str = str(e).lower()
+            if ("out of resource" in err_str
+                    or "no valid triton configs" in err_str):
                 try: del x, t, loss
                 except: pass
                 torch.cuda.empty_cache()
@@ -536,28 +551,13 @@ def train(cfg: dict):
     unique_params = model.num_params
     print(f"  Parameters   : {actual_params:,} total  ({unique_params:,} unique, tied embeddings)")
 
-    # ── torch.compile ─────────────────────────────────────────────────────────
-    compiled_model = False
-    if cfg.get("compile", True) and sys.platform == "linux" and device.type == "cuda":
-        try:
-            print("  Compiling model (torch.compile)...")
-            model = torch.compile(model)
-            compiled_model = True
-            print("  torch.compile: enabled")
-        except Exception as e:
-            print(f"  torch.compile: failed ({e}), continuing without")
-    else:
-        reason = "non-linux" if sys.platform != "linux" else "disabled"
-        print(f"  torch.compile: skipped ({reason})")
-
-    # ── Batch sizing ──────────────────────────────────────────────────────────
+    # ── Batch sizing (BEFORE torch.compile to avoid Triton kernel issues) ────
     if cfg["auto_config"] and device.type == "cuda":
         print(f"\n  Probing max micro_batch_size (VRAM search)...")
         max_mbs = probe_max_batch(
             model, device, dtype, cfg["seq_len"], cfg["vocab_size"]
         )
-        raw_model = model._orig_mod if hasattr(model, "_orig_mod") else model
-        raw_model.apply(raw_model._init_weights)
+        model.apply(model._init_weights)
         torch.cuda.empty_cache(); gc.collect()
         grad_accum = max(1, math.ceil(cfg["global_batch"] / max_mbs))
         effective_batch = max_mbs * grad_accum
@@ -573,6 +573,20 @@ def train(cfg: dict):
     tokens_per_step = effective_batch * cfg["seq_len"]
     print(f"  Tokens/step : {tokens_per_step:,}")
 
+    # ── torch.compile (after batch probe to avoid Triton shared-memory errors)
+    compiled_model = False
+    if cfg.get("compile", True) and sys.platform == "linux" and device.type == "cuda":
+        try:
+            print("  Compiling model (torch.compile)...")
+            model = torch.compile(model)
+            compiled_model = True
+            print("  torch.compile: enabled")
+        except Exception as e:
+            print(f"  torch.compile: failed ({e}), continuing without")
+    else:
+        reason = "non-linux" if sys.platform != "linux" else "disabled"
+        print(f"  torch.compile: skipped ({reason})")
+
     # ── Optional init from existing weights ───────────────────────────────────
     init_model_path = cfg.get("init_model_path")
     if init_model_path: