decrease batch for hbm

stanford-crfm · Apr 24, 2024 · c94e433 · c94e433
1 parent 61d4e96
commit c94e433
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/config/llama2_1b_mark.yaml b/config/llama2_1b_mark.yaml
@@ -237,10 +237,11 @@ trainer:
     tags: ["pile", "llama", "web_comparison"]
 
   mp: p=f32,c=bfloat16
-  train_batch_size: 512
-  num_train_steps: 5000
+  train_batch_size: 256
+  num_train_steps: 10000
   steps_per_eval: 1000
-  per_device_eval_parallelism: 64
+  max_eval_batches: 2
+  per_device_eval_parallelism: -1
   tensor_parallel_axes: ["mlp", "heads"]
   fsdp_axis: "embed"
   batch_axis: "batch"