decrease mbz

allenai · Nov 19, 2024 · a24e58c · a24e58c
1 parent 072893a
commit a24e58c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/scripts/train/nGPT-1B.py b/src/scripts/train/nGPT-1B.py
@@ -39,7 +39,7 @@ def build_trainer_config(common: CommonComponents) -> TrainerConfig:
     return (
         TrainerConfig(
             save_folder=common.save_folder,
-            rank_microbatch_size=8 * 4096,
+            rank_microbatch_size=4 * 4096,  # TODO: can we increase this?
             save_overwrite=True,
             metrics_collect_interval=10,
             cancel_check_interval=1,