Fix CUDA OOM on T4: cast model to amp_dtype directly, batch_size=2, safe RMSNorm

model/configs/train_kaggle_t4.yaml

@@ -13,23 +13,23 @@ curriculum: true
 curriculum_stages:
   - max_step: 50000
     seq_len: 512
-    batch_size: 2
-    grad_accum: 64  # 512 * 2 * 64 = 65,536 tokens/step per GPU (x2 GPUs = 131,072 tok/step)
+    batch_size: 1
+    grad_accum: 128  # 512 * 1 * 128 = 65,536 tokens/step per GPU (x2 GPUs = 131,072 tok/step)
   - max_step: 200000
     seq_len: 1024
-    batch_size: 2
-    grad_accum: 32  # 1024 * 2 * 32 = 65,536 tokens/step per GPU (x2 GPUs = 131,072 tok/step)
+    batch_size: 1
+    grad_accum: 64   # 1024 * 1 * 64 = 65,536 tokens/step per GPU (x2 GPUs = 131,072 tok/step)
   - max_step: null
     seq_len: 2048
     batch_size: 1
-    grad_accum: 32  # 2048 * 1 * 32 = 65,536 tokens/step per GPU (x2 GPUs = 131,072 tok/step)
+    grad_accum: 32   # 2048 * 1 * 32 = 65,536 tokens/step per GPU (x2 GPUs = 131,072 tok/step)
 
 seq_len: 1024
 train_split: 0.95
 
-# Kaggle T4 16GB VRAM: micro-batch 2 with SDPA and Gradient Checkpointing (~9GB peak VRAM)
-batch_size: 2
-grad_accum: 32
+# Kaggle T4 16GB VRAM: micro-batch 1 with SDPA and Gradient Checkpointing (<11GB peak VRAM)
+batch_size: 1
+grad_accum: 64
 target_global_batch: 128
 auto_tune: false
 

model/scripts/train.py

@@ -1798,6 +1798,8 @@ def main():
             scaler.update()
         else:
             opt.step()
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
         step += 1
 
         # Clean industry-standard pretraining logs: only essential metrics, no terminal garbling

model/scripts/viu_moe.py

@@ -443,25 +443,46 @@ class Viu1MoE(nn.Module):
                 total_aux = total_aux + aux
                 usages.append(usage.detach())
         h = self.norm(x)
-        logits = self.head(h)
-        if self.final_softcap > 0:
-            logits = self.final_softcap * torch.tanh(logits / self.final_softcap)
-
         loss, mt_loss = None, None
         if targets is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
-            del logits
-            logits = None
-            if self.mt_head is not None and T > 1:
-                mt_logits = self.mt_head(h[:, :-1, :])
+            # Chunked cross entropy: computes head projection + loss in 256-token chunks
+            # Eliminates 300MB+ peak allocation at vocabulary projection on 16GB GPUs
+            flat_h = h.reshape(-1, self.args.dim)
+            flat_tgt = targets.reshape(-1)
+            chunk_size = 256
+            n_tokens = flat_tgt.numel()
+            total_loss = 0.0
+            for ci in range(0, n_tokens, chunk_size):
+                ch_h = flat_h[ci : ci + chunk_size]
+                ch_tgt = flat_tgt[ci : ci + chunk_size]
+                ch_logits = self.head(ch_h)
                 if self.final_softcap > 0:
-                    mt_logits = self.final_softcap * torch.tanh(mt_logits / self.final_softcap)
-                mt_tgt = targets[:, 1:]
-                mt_loss = F.cross_entropy(mt_logits.reshape(-1, mt_logits.size(-1)), mt_tgt.reshape(-1))
-                del mt_logits, mt_tgt
+                    ch_logits = self.final_softcap * torch.tanh(ch_logits / self.final_softcap)
+                total_loss = total_loss + F.cross_entropy(ch_logits, ch_tgt, reduction="sum")
+            loss = total_loss / max(n_tokens, 1)
+
+            if self.mt_head is not None and T > 1:
+                flat_h_mt = h[:, :-1, :].reshape(-1, self.args.dim)
+                flat_tgt_mt = targets[:, 1:].reshape(-1)
+                n_tokens_mt = flat_tgt_mt.numel()
+                total_mt_loss = 0.0
+                for ci in range(0, n_tokens_mt, chunk_size):
+                    ch_h_mt = flat_h_mt[ci : ci + chunk_size]
+                    ch_tgt_mt = flat_tgt_mt[ci : ci + chunk_size]
+                    ch_logits_mt = self.mt_head(ch_h_mt)
+                    if self.final_softcap > 0:
+                        ch_logits_mt = self.final_softcap * torch.tanh(ch_logits_mt / self.final_softcap)
+                    total_mt_loss = total_mt_loss + F.cross_entropy(ch_logits_mt, ch_tgt_mt, reduction="sum")
+                mt_loss = total_mt_loss / max(n_tokens_mt, 1)
                 loss = loss + 0.1 * mt_loss
+
             if isinstance(total_aux, torch.Tensor):
                 loss = loss + total_aux
+            logits = None
+        else:
+            logits = self.head(h)
+            if self.final_softcap > 0:
+                logits = self.final_softcap * torch.tanh(logits / self.final_softcap)
         aux_info = {
             "moe_aux": float(total_aux.detach()) if isinstance(total_aux, torch.Tensor) else 0.0,
             "mt_loss": float(mt_loss.detach()) if isinstance(mt_loss, torch.Tensor) else 0.0,