add addtional config to import libtpu flags

Author: NuojCheng

src/maxtext/configs/base.yml

@@ -879,6 +879,7 @@ rope_attention_scaling: False # Scale the rotary embedding output
 compiled_trainstep_file: "" # Name of saved serialized compiled train_step, e.g. compiled_train_v5e-256.pickle
 compile_topology: '' # Target hardware version, e.g. 'v5e-256'
 compile_topology_num_slices: -1 # Number of target slices, set to a positive integer.
+compile_libtpu_flags: "" # LIBTPU_INIT_ARGS for compilation only
 
 decode_sampling_strategy: "greedy" # decode_sampling_strategy should be one of greedy, weighted, nucleus, topk, or composite(top_k -> top_p -> weighted temperature)
 decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p

src/maxtext/configs/types.py

@@ -840,6 +840,7 @@ class LayoutAndSharding(BaseModel):
       description="Allowed percentage of non-sharded parameters.",
   )
   shard_optimizer_over_data: bool = Field(False, description="Enable ZeRO-1 optimizer sharding over the data axis.")
+  compile_libtpu_flags: str = Field("", description="LIBTPU_INIT_ARGS for compilation only.")
   internal_compile: bool = Field(False, description="Use internal_compile to bypass open-source topology mappings.")
   internal_compile_num_devices: int = Field(-1, description="Number of devices when using internal_compile.")
 
@@ -2174,6 +2175,9 @@ def validate_and_set_hlo_dump_defaults():
 
     # pylint: enable=access-member-before-definition
 
+    # Add LIBTPU FLAGS
+    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + self.compile_libtpu_flags
+
     # Validate and initiate hlo dump related configs
     validate_and_set_hlo_dump_defaults()
 

src/maxtext/trainers/pre_train/train_compile.py

@@ -219,7 +219,7 @@ def is_oom(argv: Sequence[str]) -> bool:
 def main(argv: Sequence[str]) -> None:
   jax.config.update("jax_default_prng_impl", "unsafe_rbg")
   os.environ["LIBTPU_INIT_ARGS"] = (
-      os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+      os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true "
   )
   print("Starting train_compile.py...", flush=True)