fix config to use shard grad op

PrimeIntellect-ai · Sep 26, 2024 · ba392de · ba392de
1 parent 294c36d
commit ba392de
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/configs/150M/3090.toml b/configs/150M/3090.toml
@@ -3,7 +3,7 @@ project = "debug_150m_zero_band"
 
 [train]
 micro_bs = 16 # change this base on the gpu
-sharding_strategy = "NO_SHARD"
+sharding_strategy = "SHARD_GRAD_OP"
 
 [optim]
 batch_size = 512

diff --git a/configs/150M/A40.toml b/configs/150M/A40.toml
@@ -3,7 +3,7 @@ project = "debug_150m_zero_band"
 
 [train]
 micro_bs = 32 # change this base on the gpu
-sharding_strategy = "NO_SHARD"
+sharding_strategy = "SHARD_GRAD_OP"
 
 [optim]
 batch_size = 512

diff --git a/configs/150M/H100.toml b/configs/150M/H100.toml
@@ -3,7 +3,7 @@ project = "debug_150m_zero_band"
 
 [train]
 micro_bs = 64 # change this base on the gpu
-sharding_strategy = "NO_SHARD"
+sharding_strategy = "SHARD_GRAD_OP"
 
 [optim]
 batch_size = 512