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

model/scripts/viu_moe.py

@@ -316,9 +316,9 @@ class NextGenMoELayer(nn.Module):
             rows, cols = torch.where(flat_top_i == ei)
             if rows.numel() == 0:
                 continue
-            w = flat_top_p[rows, cols].unsqueeze(-1)
+            w = flat_top_p[rows, cols].unsqueeze(-1).to(flat_x.dtype)
             y_sel = expert(flat_x[rows])
-            flat_out.index_add_(0, rows, w * y_sel)
+            flat_out.index_add_(0, rows, (w * y_sel).to(flat_out.dtype))
 
         out = flat_out.view(B, T, D)
         out = self.drop(out)
@@ -346,7 +346,7 @@ class ManifoldHyperConnection(nn.Module):
         self.gate = nn.Parameter(torch.zeros(2))
 
     def forward(self, x, fx):
-        w = F.softmax(self.gate, dim=0) * 2.0
+        w = (F.softmax(self.gate, dim=0) * 2.0).to(x.dtype)
         return w[0] * x + w[1] * fx