Merge pull request #328 from google:rwitten_add_env_vars

PiperOrigin-RevId: 598020582

Author: maxtex authors

MaxText/max_utils.py

@@ -153,7 +153,7 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
 
   return parallelism_vals
 
-def create_device_mesh(config, devices=None, logging=True):
+def create_device_mesh(config, devices=None):
   """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas """
   if devices is None:
     devices = jax.devices()
@@ -163,7 +163,6 @@ def create_device_mesh(config, devices=None, logging=True):
   except:
     num_slices = 1
   num_devices_per_slice = num_devices//num_slices
-  max_logging.log(f"Devices: {devices} (num_devices: {num_devices})")
 
   multi_slice_env = num_slices > 1
 
@@ -183,8 +182,7 @@ def create_device_mesh(config, devices=None, logging=True):
   else:
     mesh = mesh_utils.create_device_mesh(ici_parallelism, devices)
 
-  if logging:
-    max_logging.log(f"Decided on mesh: {mesh}")
+  max_logging.log(f"Num_devices: {num_devices}, shape {mesh.shape}")
 
   return mesh
 

MaxText/pyconfig.py

@@ -30,6 +30,10 @@
 
 from typing import Any, Union
 
+_MAX_PREFIX = "M_"
+def yaml_key_to_env_key(s: str) -> str:
+  return _MAX_PREFIX + s.upper()
+
 def string_to_bool(s: str) -> bool:
   if s.lower() == "true":
     return True
@@ -61,9 +65,17 @@ def _lists_to_tuples(l: list[Any]) -> Union[tuple[Any],list[Any]]:
 
 class _HyperParameters():
   # pylint: disable=missing-class-docstring
+  def _validate_env_variables(self, raw_data_from_yaml):
+    for environment_var in os.environ:
+      if environment_var[:len(_MAX_PREFIX)] == _MAX_PREFIX:
+        proposed_key = environment_var[len(_MAX_PREFIX):].lower()
+        if proposed_key not in raw_data_from_yaml:
+          raise ValueError(f"We received env {environment_var} but it doesn't match a key, so it is aassumed a mistake")
+
   def __init__(self, argv: list[str], **kwargs):
     with open(argv[1], "r", encoding="utf-8") as yaml_file:
       raw_data_from_yaml = yaml.safe_load(yaml_file)
+    self._validate_env_variables(raw_data_from_yaml)
     raw_data_from_cmd_line = self._load_kwargs(argv, **kwargs)
 
     for k in raw_data_from_cmd_line:
@@ -74,28 +86,43 @@ def __init__(self, argv: list[str], **kwargs):
 
     raw_keys = OrderedDict()
     for k in raw_data_from_yaml:
-      if k in raw_data_from_cmd_line and not isinstance(raw_data_from_cmd_line[k], type(raw_data_from_yaml[k])) and \
-                                         type(raw_data_from_yaml[k]) not in _yaml_types_to_parser:
+      if k in raw_data_from_cmd_line and yaml_key_to_env_key(k) in os.environ:
+        raise ValueError(f"You are passing overrides by both CLI and ENV for `{k}`. This isn't allowed.")
+
+      if not k in raw_data_from_cmd_line and not yaml_key_to_env_key(k) in os.environ:
+        raw_keys[k] = raw_data_from_yaml[k]
+        continue
+
+      if k in raw_data_from_cmd_line:
+        new_proposal = raw_data_from_cmd_line[k]
+      else:
+        new_proposal = os.environ.get(yaml_key_to_env_key(k))
+
+      if (not isinstance(new_proposal, type(raw_data_from_yaml[k]))) and \
+                                       (type(raw_data_from_yaml[k]) not in _yaml_types_to_parser):
         raise ValueError(
             f"For key '{k}', type {type(raw_data_from_yaml[k])} not in {_yaml_types_to_parser.keys()}, can't pass"
-            " at the command line"
+            " at the CLI or ENV"
         )
 
-      if k in raw_data_from_cmd_line and isinstance(raw_data_from_cmd_line[k], type(raw_data_from_yaml[k])):
-        raw_keys[k] = raw_data_from_cmd_line[k] # take the raw data, no type conversion
-      elif k in raw_data_from_cmd_line:
+      if isinstance(new_proposal, type(raw_data_from_yaml[k])):
+        raw_keys[k] = new_proposal # take the raw data, no type conversion
+      else:
         try:
           raw_keys[k] = _yaml_types_to_parser[type(raw_data_from_yaml[k])](
-              raw_data_from_cmd_line[k]
+              new_proposal
           )  # take the command line value, but type it like the config value.
         except ValueError as e:
-          raise ValueError(f"Couldn't parse value from command line '{raw_data_from_cmd_line[k]}' for key '{k}'") from e
-      else:
-        raw_keys[k] = raw_data_from_yaml[k]
+          raise ValueError(f"Couldn't parse value from CLI or ENV '{new_proposal}' for key '{k}'") from e
 
     _HyperParameters.update_model_vars(raw_keys)
     _HyperParameters.user_init(raw_keys)
+
     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]}")
 
   def _load_kwargs(self, argv: list[str], **kwargs):
     args_dict = dict(a.split("=") for a in argv[2:])

MaxText/train.py

@@ -345,7 +345,6 @@ def main(argv: Sequence[str]) -> None:
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS","") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
   pyconfig.initialize(argv)
-  print(f"Found {jax.device_count()} devices.")
   config = pyconfig.config
   validate_train_config(config)
   cc.initialize_cache(os.path.expanduser(config.jax_cache_dir))

end_to_end/llama_load_and_test.sh

@@ -1,9 +1,10 @@
 set -e
 idx=$(date +%Y-%m-%d-%H-%M)
+
 #TODO(internal bug -- migrate to XLML)
 #pip install torch
 #gsutil cp -r gs://maxtext-llama/llama2-7b/meta-ckpt /tmp/
 #python3 MaxText/convert_llama_ckpt.py --base-model-path /tmp/meta-ckpt --model-size 7b --maxtext-model-path gs://maxtext-llama/test/${idx}/decode-ckpt-maxtext/
 #python3 MaxText/decode.py MaxText/configs/base.yml run_name=runner_${idx} base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset load_from_other_directory=gs://maxtext-llama/test/${idx}/decode-ckpt-maxtext/ load_from_other_directory_step=0 per_device_batch_size=1 model_name='llama2-7b' assets_path=gs://maxtext-llama/llama2-7b ici_tensor_parallelism=4 steps=1 max_prefill_predict_length=4  max_target_length=16 async_checkpointing=false prompt="I love to" autoregressive_decode_assert="read. I love to write. I love to share." attention=dot_product
 
-python3 MaxText/decode.py MaxText/configs/base.yml run_name=runner_${idx} base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset load_from_other_directory=gs://maxtext-llama/test/2024-01-12-17-46/decode-ckpt-maxtext/ load_from_other_directory_step=0 per_device_batch_size=1 model_name='llama2-7b' assets_path=gs://maxtext-llama/llama2-7b ici_tensor_parallelism=4 steps=1 max_prefill_predict_length=4  max_target_length=16 async_checkpointing=false prompt="I love to" autoregressive_decode_assert="read. I love to write. I love to share." attention=dot_product
+python3 MaxText/decode.py MaxText/configs/base.yml run_name=runner_${idx} base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset load_from_other_directory=gs://maxtext-llama/test/2024-01-12-17-46/decode-ckpt-maxtext/ load_from_other_directory_step=0 per_device_batch_size=1 model_name='llama2-7b' assets_path=gs://maxtext-llama/llama2-7b ici_tensor_parallelism=4 steps=1 max_prefill_predict_length=4  max_target_length=16 async_checkpointing=false prompt="I love to" autoregressive_decode_assert="read. I love to write. I love to share." attention=dot_product