larger tiles and better speed, full loss curve checked

facebookresearch · Dec 27, 2021 · 8b9befa · 8b9befa
1 parent b36e8ca
commit 8b9befa
Show file tree

Hide file tree

Showing 27 changed files with 8 additions and 10 deletions.
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float16_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float16_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float16_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float16_Act:_gelu.png
diff --git a/...lots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float16_Act:_squared_relu.png b/...lots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float16_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float32_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float32_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float32_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float32_Act:_gelu.png
diff --git a/...lots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float32_Act:_squared_relu.png b/...lots/fused_dropout/Dropout_Bias_False_FW+BW_torch.float32_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float16_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float16_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float16_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float16_Act:_gelu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float16_Act:_squared_relu.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float16_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float32_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float32_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float32_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float32_Act:_gelu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float32_Act:_squared_relu.png b/docs/plots/fused_dropout/Dropout_Bias_False_FW_torch.float32_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float16_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float16_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float16_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float16_Act:_gelu.png
diff --git a/...plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float16_Act:_squared_relu.png b/...plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float16_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float32_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float32_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float32_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float32_Act:_gelu.png
diff --git a/...plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float32_Act:_squared_relu.png b/...plots/fused_dropout/Dropout_Bias_True_FW+BW_torch.float32_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float16_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float16_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float16_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float16_Act:_gelu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float16_Act:_squared_relu.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float16_Act:_squared_relu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float32_Act:_None.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float32_Act:_None.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float32_Act:_gelu.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float32_Act:_gelu.png
diff --git a/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float32_Act:_squared_relu.png b/docs/plots/fused_dropout/Dropout_Bias_True_FW_torch.float32_Act:_squared_relu.png
diff --git a/examples/microGPT.py b/examples/microGPT.py
@@ -73,7 +73,7 @@ def __init__(
                 "feedforward_config": {
                     "name": "FusedMLP",  # Use MLP if Triton is not available
                     "dropout": self.hparams.mlp_pdrop,
-                    "activation": "squared_relu",
+                    "activation": "gelu",
                     "hidden_layer_multiplier": self.hparams.hidden_layer_multiplier,
                 },
             }
@@ -273,7 +273,7 @@ def top_k_logits(logits, k):
     # Adjust batch depending on the available memory on your machine.
     # You can also use reversible layers to save memory
     REF_BATCH = 512
-    BATCH = 128
+    BATCH = 256
 
     WORKERS = 4
     EPOCHS = 1

diff --git a/xformers/triton/dropout.py b/xformers/triton/dropout.py
@@ -19,13 +19,10 @@
     get_triton_activation_kernel,
 )
 from xformers.triton.k_dropout import k_dropout_bw, k_dropout_fw
-from xformers.triton.sum_strided import sum_2d_dim_0
 
-# NOTE: GROUP_M and BLOCK_N need to be kept low (<16x64)
-# for the random numbers to be good enough
-GROUP_M = 8
+GROUP_M = 32  # 32
 BLOCK_M = GROUP_M // 4
-BLOCK_N = 128
+BLOCK_N = 128  # 128
 
 
 # Helper to handle the SPMD launch grid and error cases
@@ -145,7 +142,7 @@ def grid(meta):
         return (
             grad_in.reshape_as(grad_out),
             None,
-            sum_2d_dim_0(grad_bias) if ctx.trainable_bias else None,
+            torch.sum(grad_bias, dim=0) if ctx.trainable_bias else None,
             None,
             None,
             None,

diff --git a/xformers/triton/sum_strided.py b/xformers/triton/sum_strided.py
@@ -36,10 +36,10 @@ def sum_2d_dim_0(x: torch.Tensor):
     )
 
     BLOCK_M = min(triton.next_power_of_2(M), 2048)
-    BLOCK_N = 48
+    BLOCK_N = 32
     if BLOCK_M > 256:
         BLOCK_N = 16
-    if BLOCK_M >= 1024:
+    if BLOCK_M > 1024:
         BLOCK_N = 8
 
     def grid(meta):
@@ -53,6 +53,7 @@ def grid(meta):
         x.dtype == torch.float16,
         BLOCK_M=BLOCK_M,
         BLOCK_N=BLOCK_N,
+        num_stages=4,
     )
     # fmt: on