Merge pull request #619 from google:raymondzou-128b-offload-fix

PiperOrigin-RevId: 627768373
google · Apr 24, 2024 · 206e84c · 206e84c
2 parents 718d9e7 + 62b5012
commit 206e84c
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\