Merge pull request #3131 from AI-Hypercomputer:igorts-dev

PiperOrigin-RevId: 871361489

Author: Google-ML-Automation

src/MaxText/__init__.py

@@ -27,6 +27,12 @@
 
 from collections.abc import Sequence
 
+import os
+# In order to have any effect on the C++ logging this has to be set before we import anything from jax.
+# When jax is imported, its `__init__.py` calls `cloud_tpu_init()`, which also initializes the C++ logger.
+os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "0")
+del os
+
 from jax.sharding import Mesh
 
 from MaxText import pyconfig

src/MaxText/train.py

@@ -538,7 +538,6 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any, Any]
   # TF allocates extraneous GPU memory when using TFDS data
   # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
   tf.config.set_visible_devices([], "GPU")
-  os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
     os.environ["LIBTPU_INIT_ARGS"] = (
         os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"