Update 128B config on v5e to use qkv_proj_offloaded remat_policy

google · Apr 24, 2024 · 62b5012 · 62b5012
1 parent 718d9e7
commit 62b5012
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/configs/v5e/128b.sh b/MaxText/configs/v5e/128b.sh
@@ -43,7 +43,7 @@ export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xl
 
 python3 MaxText/$EXECUTABLE MaxText/configs/base.yml\
  steps=30 per_device_batch_size=1 enable_checkpointing=false\
- enable_profiler=false remat_policy=minimal_offloaded global_parameter_scale=128\
+ enable_profiler=false remat_policy=qkv_proj_offloaded global_parameter_scale=128\
  ici_fsdp_parallelism=16 ici_tensor_parallelism=16\
  max_target_length=2048 base_output_directory=gs:https://runner-maxtext-logs\
  use_iota_embed=true reuse_example_batch=1\