Merge pull request #3977 from CIeNET-International:maxtext/user/dora/v5e_config_oom

PiperOrigin-RevId: 921524627

Author: Google-ML-Automation

benchmarks/maxtext_v5e_model_configs.py

@@ -40,6 +40,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -63,6 +64,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -86,6 +88,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -112,6 +115,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -133,6 +137,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -155,6 +160,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -167,7 +173,7 @@
         model_type="llama2-13b",
         tuning_params={
             "ici_fsdp_parallelism": -1,
-            "per_device_batch_size": 8,
+            "per_device_batch_size": 6,
             "remat_policy": "qkv_proj_offloaded",
             "max_target_length": 2048,
             "use_iota_embed": True,
@@ -177,6 +183,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -198,6 +205,7 @@
             "reuse_example_batch": 1,
             "enable_checkpointing": False,
             "profiler": "xplane",
+            "use_qwix_quantization": True,
         },
         xla_flags=(xla_flags_library.DATA_PARALLEL_OVERLAP + xla_flags_library.CF_FOR_ALL_GATHER),
     ),
@@ -236,6 +244,7 @@
             "profiler": "xplane",
             "skip_first_n_steps_for_profiler": 10,
             "profiler_steps": 5,
+            "use_qwix_quantization": True,
         },
         xla_flags=(
             xla_flags_library.DENSE_VMEM_LIMIT_FLAG