Add llama-70b gpu config.

PiperOrigin-RevId: 634267313
google · May 16, 2024 · fa92a21 · fa92a21
1 parent 3e65eb6
commit fa92a21
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/MaxText/configs/llama2_70b_gpu.yml b/MaxText/configs/llama2_70b_gpu.yml
@@ -0,0 +1,18 @@
+base_config: "base.yml"
+
+run_name: "gpu_train_test"
+# Args coming from the NVIDIA spreadsheet http:https://shortn/_AhULYn1mX4.
+hardware: "gpu"
+steps: 30
+model_name: "llama2-70b"
+enable_checkpointing: False
+attention: "cudnn_flash_te"
+remat_policy: "full"
+use_iota_embed: True
+scan_layers: True
+dataset_type: "synthetic"
+async_checkpointing: False
+logits_dot_in_fp32: False
+
+per_device_batch_size: 6
+max_target_length: 4096