Merge pull request #3570 from AI-Hypercomputer:chengnuojin-compile-flags

PiperOrigin-RevId: 895595405

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -442,6 +442,7 @@ hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu', 'gpu_multiprocess'
 # internal_compile allows bypassing open-source topology name mappings when using internal topologies directly via get_topology_desc.
 internal_compile: False
 internal_compile_num_devices: -1 # You must specify the number of devices when using internal_compile.
+compile_xla_flags: "" # Compiler options e.g. compile_xla_flags="--xla_tpu_num_sparse_cores_for_gather_offloading=1 --xla_tpu_scoped_vmem_limit_kib=65536" 
 
 # Parallelism
 shard_mode: "auto" # can be either auto or explicit

src/maxtext/configs/types.py

@@ -842,6 +842,7 @@ class LayoutAndSharding(BaseModel):
   shard_optimizer_over_data: bool = Field(False, description="Enable ZeRO-1 optimizer sharding over the data axis.")
   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.")
+  compile_xla_flags: str = Field("", description="Compiler options for compilation only.")
 
 
 class DcnParallelism(BaseModel):

src/maxtext/trainers/pre_train/train_compile.py

@@ -144,7 +144,9 @@ def jit_and_compile(
     )
     maxtext_utils.maybe_dump_jaxpr(config, jitted, func_input_args)
     lowered = jitted.lower(*func_input_args, **func_input_kwargs)
-  compiled = lowered.compile()
+  # Import libtpu flags as compiler options. Defaults to empty dict if string is empty.
+  compiler_options = max_utils.parse_libtpu_flags_to_dict(config.compile_xla_flags)
+  compiled = lowered.compile(compiler_options=compiler_options)
   return compiled
 
 

src/maxtext/utils/max_utils.py

@@ -21,6 +21,7 @@
 from functools import partial
 import os
 import socket
+import re
 import subprocess
 import time
 from typing import Any
@@ -52,6 +53,43 @@
 # pylint: disable=too-many-positional-arguments
 
 
+def parse_libtpu_flags_to_dict(flags_str: str) -> dict:
+  """
+  Parses a string of XLA flags into a dictionary of compilation options.
+  This function is only for compilation usage.
+  """
+  if not flags_str or not flags_str.strip():
+    return {}
+
+  # Clean the string by removing line-continuation backslashes
+  cleaned_str = flags_str.replace("\\", " ")
+
+  # Split by any whitespace (handles single spaces, multiple spaces, newlines)
+  tokens = cleaned_str.split()
+
+  options_dict = {}
+
+  # Regex to strictly match '--key=value' for an isolated token
+  # Key assumes alphanumeric + underscores. Value is anything after the '='.
+  token_pattern = re.compile(r"^--([a-zA-Z0-9_]+)=(.+)$")
+
+  for token in tokens:
+    match = token_pattern.match(token)
+    if not match:
+      # Throw an error immediately if any token fails the strict format
+      raise ValueError(f"Invalid flag format detected: '{token}'. Expected format: '--key=value'")
+
+    key, value = match.groups()
+
+    # Optional: Catch duplicate flags
+    if key in options_dict:
+      raise ValueError(f"Duplicate flag detected: '--{key}'")
+
+    options_dict[key] = value
+
+  return options_dict
+
+
 def with_memory_kind(t, memory_kind):
   return jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind=memory_kind), t)