Merge pull request #1099 from AI-Hypercomputer:mattdavidow-jdi-telemetry

PiperOrigin-RevId: 707276379

Author: maxtext authors

MaxText/configs/base.yml

@@ -342,6 +342,11 @@ grain_worker_count: 1
 steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps
 log_period: 100 # Flushes Tensorboard
 
+jax_distributed_initialization_timeout: 300 # This is the default timeout in https://github.com/jax-ml/jax/blob/main/jax/_src/distributed.py
+# Note there are two separate initializations - the jax coordination service (aka jax.distributed.initialize) and the backend (e.g. PjRT), the timeout above refers
+# only to the jax coordination service.
+jax_debug_log_modules: "" # Set this to "jax" to enable jax verbose logging such as for the jax coordination service initialization. 
+
 # We take inspiration from Llama2's learning rate (LR) schedule, see https://arxiv.org/pdf/2307.09288.pdf section 2.2
 # Learning rate schedule has either two or three parts:
 # 1) Linear warmup from 0 to [learning_rate] over steps 0 to [learning_rate_schedule_steps * warmup_steps_fraction]
@@ -477,6 +482,8 @@ prometheus_port: 0
 enable_jax_profiler: False
 jax_profiler_port: 9999
 
+log_config: True # Prints the config (after defaults have been set by pyconfig logic)
+
 # Checkpoint Structured logging
 enable_checkpoint_cloud_logger: False
 

MaxText/max_utils.py

@@ -224,11 +224,11 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     return
   if is_gpu_backend(raw_keys):
     max_logging.log("Attempting to initialize the jax distributed system for GPU backend...")
-    initialize_jax_for_gpu()
+    initialize_jax_for_gpu(raw_keys)
     max_logging.log("Jax distributed system initialized on GPU!")
   elif is_cpu_backend(raw_keys):
     max_logging.log("Attempting to initialize the jax distributed system for CPU backend...")
-    initialize_jax_for_cpu()
+    initialize_jax_for_cpu(raw_keys)
     max_logging.log("Jax distributed system initialized on CPUs!")
   elif (
       raw_keys["enable_checkpointing"]
@@ -238,13 +238,13 @@ def maybe_initialize_jax_distributed_system(raw_keys):
   ) or raw_keys["hardware"] == "gpu_multiprocess":
     max_logging.log("Attempting to initialize the jax distributed system...")
     if not raw_keys["enable_emergency_checkpoint"]:
-      jax.distributed.initialize()
+      jax.distributed.initialize(initialization_timeout=raw_keys["jax_distributed_initialization_timeout"])
     else:
       initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys)
     max_logging.log("Jax distributed system initialized!")
 
 
-def initialize_jax_for_gpu():
+def initialize_jax_for_gpu(raw_keys):
   """Jax distributed initialize for GPUs."""
   if os.environ.get("JAX_COORDINATOR_IP") is not None:
     coordinator_ip = str(os.getenv("JAX_COORDINATOR_IP"))
@@ -253,11 +253,12 @@ def initialize_jax_for_gpu():
         coordinator_address=f"{coordinator_ip}:{coordinator_port}",
         num_processes=int(os.getenv("NNODES")),
         process_id=int(os.getenv("NODE_RANK")),
+        initialization_timeout=raw_keys["jax_distributed_initialization_timeout"],
     )
     max_logging.log(f"JAX global devices: {jax.devices()}")
 
 
-def initialize_jax_for_cpu():
+def initialize_jax_for_cpu(raw_keys):
   """Jax distributed initialize for CPUs. Includes retries until the coordinator is ready."""
   coordinator_ip_address = get_coordinator_ip_address()
   coordinator_address = coordinator_ip_address + ":1234"  # JAX coordinator port used in XPK
@@ -272,6 +273,7 @@ def initialize_jax_for_cpu():
       coordinator_address=coordinator_address,
       process_id=pid,
       num_processes=int(os.environ.get("JAX_PROCESS_COUNT")),
+      initialization_timeout=raw_keys["jax_distributed_initialization_timeout"],
   )
 
 
@@ -288,7 +290,11 @@ def initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys):
         f"Using {process_id} as the process_id and {coordinator_address} as the"
         " coordinator_address to initialize JAX distributed runtime..."
     )
-    jax.distributed.initialize(coordinator_address=coordinator_address, process_id=int(process_id))
+    jax.distributed.initialize(
+        coordinator_address=coordinator_address,
+        process_id=int(process_id),
+        initialization_timeout=raw_keys["jax_distributed_initialization_timeout"],
+    )
     if raw_keys["use_replicator_service"]:
       REPLICATOR_FILE = "replicator.yaml"
       TEMP_FILE = REPLICATOR_FILE + ".tmp"
@@ -324,7 +330,7 @@ def initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys):
         "Initializing JAX distributed runtime without args when emergency checkpointing is"
         " enabled. This should not happen and your workload may have unexpected behavior."
     )
-    jax.distributed.initialize()
+    jax.distributed.initialize(initialization_timeout=raw_keys["jax_distributed_initialization_timeout"])
 
   ocp.multihost.initialize_runtime_to_distributed_ids()
   ocp.multihost.initialize_distributed_to_device_ids()

MaxText/pyconfig.py

@@ -345,6 +345,8 @@ def __init__(self, argv: list[str], **kwargs):
     validate_no_keys_overwritten_twice(keys_from_env_and_command_line, keys_from_model)
 
     # We initialize the jax distributed system here because it must be done before device backend is initialized.
+    if raw_keys["jax_debug_log_modules"]:
+      jax.config.update("jax_debug_log_modules", raw_keys["jax_debug_log_modules"])
     max_utils.maybe_initialize_jax_distributed_system(raw_keys)
 
     if raw_keys["jax_cache_dir"]:
@@ -367,8 +369,10 @@ def __init__(self, argv: list[str], **kwargs):
     self.keys = raw_keys
     keys = [k for k in raw_keys]  # pylint: disable=unnecessary-comprehension
     keys.sort()
-    for k in keys:
-      max_logging.log(f"Config param {k}: {raw_keys[k]}")
+
+    if raw_keys["log_config"]:
+      for k in keys:
+        max_logging.log(f"Config param {k}: {raw_keys[k]}")
 
   @staticmethod
   def user_init(raw_keys):