perf: disable grad_ckpt on T4 and enlarge loss chunk_size to 2048 for 10k+ tok/s

model/configs/train_kaggle_t4.yaml

@@ -44,10 +44,10 @@ weight_decay: 0.1
 betas: [0.9, 0.95]
 grad_clip: 1.0
 
-# Automatic FP16 on T4 + GradScaler + PagedAdamW8bit + Grad Checkpointing (100% Stable)
+# Automatic FP16 on T4 + GradScaler + PagedAdamW8bit (Maximum speed: grad_ckpt off for ~10k+ tok/s)
 amp: auto
 optimizer: paged_adamw_8bit
-grad_ckpt: true
+grad_ckpt: false
 
 compile: false
 fp8: false

model/scripts/train.py

@@ -718,6 +718,8 @@ def main():
     ap.add_argument("--muon_weight_decay", type=float, default=None, help="Muon weight decay override (default: 0.01)")
     ap.add_argument("--compile", action="store_true", help="torch.compile on (10-30%% free speedup, default off)")
     ap.add_argument("--fp8", action="store_true", help="torchao Float8 training (Blackwell/5090 only, EXPERIMENTAL)")
+    ap.add_argument("--grad_ckpt", dest="grad_ckpt", action="store_true", default=None, help="Enable gradient checkpointing")
+    ap.add_argument("--no_grad_ckpt", "--no-grad-ckpt", dest="grad_ckpt", action="store_false", help="Disable gradient checkpointing for maximum speed")
     args = ap.parse_args()
 
     base = Path(__file__).parent
@@ -1505,14 +1507,18 @@ def main():
         {"params": no_decay_params, "weight_decay": 0.0},
     ]
 
-    if bool(C("grad_ckpt", False)):
+    use_gc = args.grad_ckpt if args.grad_ckpt is not None else bool(C("grad_ckpt", False))
+    if use_gc:
         target = model.module if hasattr(model, "module") else model
         if hasattr(target, "gradient_checkpointing_enable"):
             target.gradient_checkpointing_enable()
             if R0:
-                print("[model] gradient checkpointing enabled (VRAM activation memory halved for 48L)")
+                print("[model] gradient checkpointing enabled")
         elif R0:
             print("[model] WARN: grad_ckpt requested but model has no gradient_checkpointing_enable()")
+    else:
+        if R0 and not args.smoke:
+            print("[model] gradient checkpointing DISABLED (pure maximum FLOP speed, zero backward recomputation)")
 
     opt_name = str(args.optimizer or C("optimizer", "adamw")).lower()
     if opt_name == "muon":

model/scripts/viu1_dense.py

@@ -141,7 +141,7 @@ class Attention(nn.Module):
         v = v.transpose(1, 2)
         drop_p = self.drop.p if self.training else 0.0
         if self.sliding:
-            causal = self._get_mask(T, x.device)
+            causal = self._get_mask(T, x.device).to(dtype=q.dtype)
             attn_out = F.scaled_dot_product_attention(q, k, v, attn_mask=causal, dropout_p=drop_p)
         else:
             attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop_p)
@@ -244,7 +244,7 @@ class Viu1Dense(nn.Module):
         if targets is not None:
             flat_h = h.reshape(-1, self.args.dim)
             flat_tgt = targets.reshape(-1)
-            chunk_size = 2048 if (h.is_cuda and torch.cuda.get_device_properties(h.device).total_memory > 16 * 1024**3) else 512
+            chunk_size = 2048 if (h.is_cuda and torch.cuda.get_device_properties(h.device).total_memory >= 14 * 1024**3) else 512
             n_tokens = flat_tgt.numel()
             total_loss = 0.0
             for ci in range(0, n_tokens, chunk_size):