From 733c8a1bd1d9db71002172f71d4c8f9177218201 Mon Sep 17 00:00:00 2001
From: YannAhlgrim <yann@ahlgrim.net>
Date: Wed, 8 Jul 2026 11:39:00 +0200
Subject: [PATCH] remove gradient accumulation

---
 configs/grids/lp_grid.yaml               | 4 ++--
 configs/supervised_vitg16_224_in22k.yaml | 4 ++--
 configs/supervised_vith14_224.yaml       | 4 ++--
 configs/supervised_vith14_224_in1k.yaml  | 4 ++--
 configs/supervised_vith14_224_in22k.yaml | 4 ++--
 configs/supervised_vith16_448.yaml       | 4 ++--
 configs/supervised_vith16_448_in1k.yaml  | 4 ++--
 7 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/configs/grids/lp_grid.yaml b/configs/grids/lp_grid.yaml
index 70ec4ee..bb1e624 100644
--- a/configs/grids/lp_grid.yaml
+++ b/configs/grids/lp_grid.yaml
@@ -9,8 +9,8 @@ grid:
   optimization.weight_decay: [5.0e-4, 0]
   optimization.momentum: [0.9]
   optimization.lr_schedule: [cosine]
-  optimization.gradient_accumulation_steps: [8, 16, 32, 64]
-  data.batch_size: [16]
+  optimization.gradient_accumulation_steps: [1]
+  data.batch_size: [128, 256, 512, 1024]
   meta.representation_type: [last_avgpool, last4_avgpool_concat]
   meta.head_type: [linear, bn_linear]
 
diff --git a/configs/supervised_vitg16_224_in22k.yaml b/configs/supervised_vitg16_224_in22k.yaml
index fc58d0f..d7a5ef4 100644
--- a/configs/supervised_vitg16_224_in22k.yaml
+++ b/configs/supervised_vitg16_224_in22k.yaml
@@ -12,7 +12,7 @@ meta:
   use_gradient_checkpointing: true
 
 data:
-  batch_size: 16
+  batch_size: 256
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
@@ -33,7 +33,7 @@ optimization:
   lr: 0.001
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  gradient_accumulation_steps: 16
+  gradient_accumulation_steps: 1
   final_lr: 0.0
   early_stopping:
     enabled: true
diff --git a/configs/supervised_vith14_224.yaml b/configs/supervised_vith14_224.yaml
index b0d312f..d3b9f07 100644
--- a/configs/supervised_vith14_224.yaml
+++ b/configs/supervised_vith14_224.yaml
@@ -10,7 +10,7 @@ meta:
   head_type: bn_linear
 
 data:
-  batch_size: 16
+  batch_size: 512
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
@@ -31,7 +31,7 @@ optimization:
   lr: 0.001
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  gradient_accumulation_steps: 32
+  gradient_accumulation_steps: 1
   final_lr: 0.0
   early_stopping:
     enabled: true
diff --git a/configs/supervised_vith14_224_in1k.yaml b/configs/supervised_vith14_224_in1k.yaml
index b1de2af..6dcc577 100644
--- a/configs/supervised_vith14_224_in1k.yaml
+++ b/configs/supervised_vith14_224_in1k.yaml
@@ -11,7 +11,7 @@ meta:
   head_type: linear
 
 data:
-  batch_size: 16
+  batch_size: 1024
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
@@ -32,7 +32,7 @@ optimization:
   lr: 0.01
   weight_decay: 0
   lr_schedule: cosine
-  gradient_accumulation_steps: 64
+  gradient_accumulation_steps: 1
   final_lr: 0.0
   early_stopping:
     enabled: true
diff --git a/configs/supervised_vith14_224_in22k.yaml b/configs/supervised_vith14_224_in22k.yaml
index d9a59e7..89756e2 100644
--- a/configs/supervised_vith14_224_in22k.yaml
+++ b/configs/supervised_vith14_224_in22k.yaml
@@ -11,7 +11,7 @@ meta:
   head_type: bn_linear
 
 data:
-  batch_size: 16
+  batch_size: 512
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
@@ -32,7 +32,7 @@ optimization:
   lr: 0.001
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  gradient_accumulation_steps: 32
+  gradient_accumulation_steps: 1
   final_lr: 0.0
   early_stopping:
     enabled: true
diff --git a/configs/supervised_vith16_448.yaml b/configs/supervised_vith16_448.yaml
index 3e4b95f..e7ce62f 100644
--- a/configs/supervised_vith16_448.yaml
+++ b/configs/supervised_vith16_448.yaml
@@ -10,7 +10,7 @@ meta:
   head_type: bn_linear
 
 data:
-  batch_size: 16
+  batch_size: 1024
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
@@ -31,7 +31,7 @@ optimization:
   lr: 0.01
   weight_decay: 5.0e-4
   lr_schedule: cosine
-  gradient_accumulation_steps: 64
+  gradient_accumulation_steps: 1
   final_lr: 0.0
   early_stopping:
     enabled: true
diff --git a/configs/supervised_vith16_448_in1k.yaml b/configs/supervised_vith16_448_in1k.yaml
index 5eae316..fc8f82b 100644
--- a/configs/supervised_vith16_448_in1k.yaml
+++ b/configs/supervised_vith16_448_in1k.yaml
@@ -12,7 +12,7 @@ meta:
   use_gradient_checkpointing: true
 
 data:
-  batch_size: 64
+  batch_size: 512
   root_path: ./wilds_data
   num_workers: 10
   pin_mem: true
@@ -33,7 +33,7 @@ optimization:
   lr: 0.001
   weight_decay: 0
   lr_schedule: cosine
-  gradient_accumulation_steps: 8
+  gradient_accumulation_steps: 1
   final_lr: 0.0
   early_stopping:
     enabled: true