feat: Improve permission tests and squash feature branch

igorts-git · igorts-git · commit c6a7ef65184d · 2026-03-09T14:58:35.000-07:00
- Reworks tests for read-only directories to use /sys instead of temporary directories with chmod. This fixes CI failures when tests are run as root.
- Squashes commits related to int4 support, checkpoint directory checks, and RL parsing unit tests.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -2,16 +2,16 @@
 
 # Model bring-up
 src/MaxText/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande
-src/MaxText/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @suexu1025 @jesselu-google
+src/MaxText/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @suexu1025 @jesselu-google @NuojCheng
 src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande
-src/MaxText/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @suexu1025 @jesselu-google
+src/MaxText/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @suexu1025 @jesselu-google @NuojCheng
 
 # Features
 src/maxtext/experimental/rl @A9isha @khatwanimohit @xuefgu @gagika @richjames0 @shralex @NicoGrande
 src/MaxText/input_pipeline @aireenmei @SurbhiJainUSC @richjames0 @shralex @NicoGrande
 src/MaxText/kernels/megablox @RissyRan @michelle-yooh @gagika @richjames0 @shralex @suexu1025 @jesselu-google
 src/MaxText/kernels/ragged_attention.py @patemotter @vipannalla @richjames0 @shralex
-src/MaxText/layers/pipeline.py @gobbleturk @richjames0 @shralex
+src/MaxText/layers/pipeline.py @gobbleturk @richjames0 @shralex @NuojCheng
 src/MaxText/layers/moe.py @RissyRan @michelle-yooh @gagika @richjames0 @shralex @suexu1025 @jesselu-google
 src/MaxText/layers/multi_token_prediction.py @parambole @RissyRan @gagika @richjames0 @shralex
 src/MaxText/elastic_train.py @lukebaumann @shauryagup @richjames0 @shralex
diff --git a/dependencies/requirements/base_requirements/requirements.txt b/dependencies/requirements/base_requirements/requirements.txt
@@ -17,6 +17,7 @@ jax
 jaxlib
 jaxtyping
 jsonlines
+math-verify
 ml-collections
 ml-goodput-measurement
 numpy
diff --git a/dependencies/requirements/generated_requirements/tpu-requirements.txt b/dependencies/requirements/generated_requirements/tpu-requirements.txt
@@ -120,6 +120,7 @@ lxml>=6.0.2
 markdown-it-py>=4.0.0
 markdown>=3.10
 markupsafe>=3.0.3
+math-verify>=0.9.0
 matplotlib>=3.10.7
 mccabe>=0.7.0
 mdurl>=0.1.2
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
@@ -28,6 +28,7 @@
 from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator
 from maxtext.utils import exceptions
 from maxtext.utils import max_logging
+from maxtext.utils import gcs_utils
 import numpy as np
 import orbax.checkpoint as ocp
 from orbax.checkpoint import v1 as ocp_v1
@@ -245,8 +246,7 @@ def create_orbax_checkpoint_manager(
     item_handlers["iter"] = GrainCheckpointHandler()
 
   # local storage checkpoint needs parent directory created
-  p = epath.Path(checkpoint_dir)
-  p.mkdir(exist_ok=True, parents=True)
+  p = gcs_utils.mkdir_and_check_permissions(checkpoint_dir)
   if enable_continuous_checkpointing:
     save_decision_policy = save_decision_policy_lib.ContinuousCheckpointingPolicy()
     preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep)
@@ -300,19 +300,19 @@ def create_orbax_emergency_checkpoint_manager(
   flags.FLAGS.experimental_orbax_use_distributed_process_id = True
   max_logging.log("Creating emergency checkpoint manager...")
 
+  persistent_p = gcs_utils.mkdir_and_check_permissions(persistent_checkpoint_dir)
+
   # Only create directories if running on GPUs as the previous
   # directory structure might be assumed by TPUs
   if global_mesh.devices.flatten()[0].platform == "gpu":
     # pylint: disable=protected-access
     local_checkpoint_dir = f"{local_checkpoint_dir}/{jax._src.distributed.global_state.process_id}"
     local_p = epath.Path(local_checkpoint_dir)
-    persistent_p = epath.Path(persistent_checkpoint_dir)
     local_p.mkdir(exist_ok=True, parents=True)
-    persistent_p.mkdir(exist_ok=True, parents=True)
 
   manager = EmergencyCheckpointManager(
       local_checkpoint_dir,
-      epath.Path(persistent_checkpoint_dir),
+      persistent_p,
       global_mesh=global_mesh,
       abstract_state=abstract_state,
       options=emergency_checkpoint_manager.CheckpointManagerOptions(
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
@@ -81,6 +81,7 @@ class QuantizationType(str, Enum):
   """Supported quantization schemes."""
 
   NONE = ""
+  INT4 = "int4"
   INT8 = "int8"
   INTMP = "intmp"
   FP8 = "fp8"
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
@@ -1318,7 +1318,7 @@ def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index):
         )
 
         # Sum up the partial outputs across the expert shards.
-        output = jnp.reshape(output, (-1, sequence_length, self.config.emb_dim))
+        output = jnp.reshape(output, (-1, sequence_length, self.config.emb_dim // self.get_tensor_parallelism_size()))
         output = jax.lax.psum_scatter(output, self._expert_parallelism_name, scatter_dimension=0, tiled=True)
 
       else:
diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py
@@ -655,6 +655,15 @@ def get_fp8_full_qwix_rule(config: Config):
 
 def get_quantization_rule(config: Config):
   match config.quantization:
+    case "int4":
+      return qwix.QtRule(
+          module_path="decoder/.*layers.*",
+          weight_qtype=jnp.int4,
+          act_qtype=jnp.int4,
+          bwd_qtype=jnp.int4,
+          bwd_weight_grad_tile_size=1 / config.quantization_local_shard_count,
+          op_names=("dot_general",),
+      )
     case "int8":
       return qwix.QtRule(
           module_path="decoder/.*layers.*",
@@ -702,6 +711,8 @@ def get_qt_provider(config):
   match config.quantization:
     case "int8":
       return qwix.QtProvider([get_quantization_rule(config)])
+    case "int4":
+      return qwix.QtProvider([get_quantization_rule(config)])
     case "fp8":
       return qwix.QtProvider([get_quantization_rule(config)])
     case "fp8_full":
diff --git a/src/maxtext/trainers/post_train/rl/evaluate_rl.py b/src/maxtext/trainers/post_train/rl/evaluate_rl.py
@@ -100,6 +100,7 @@ def score_responses(tmvp_config, question, responses, answer):
       Tuple of (is_correct, is_partially_correct, has_correct_format)
   """
   match_format = utils_rl.get_match_format_regex(tmvp_config)
+  answer_fallback = utils_rl.get_answer_fallback_regex(tmvp_config)
 
   if tmvp_config.debug.rl:
     max_logging.log("========================================")
@@ -113,10 +114,19 @@ def score_responses(tmvp_config, question, responses, answer):
   has_correct_format = False
 
   for response in responses:
-    # Extract numerical response
-    extracted_response = guess.group(1) if (guess := match_format.search(response)) is not None else "-1000000"
+    # Extract answer: prefer the full format match; fall back to the last
+    # <answer>...</answer> tag if full format match is not found, so result
+    # scoring is decoupled from format.
+    full_match = match_format.search(response)
+    if full_match is not None:
+      extracted_response = full_match.group(1)
+    else:
+      # Find the *last* occurrence of the answer tag (most likely the final answer).
+      fallback_matches = answer_fallback.findall(response)
+      extracted_response = fallback_matches[-1].strip() if fallback_matches else "-1000000"
     if tmvp_config.debug.rl:
-      max_logging.log(f"Evaluation extracted_response: {extracted_response}")
+      used = "full format" if full_match is not None else "answer-tag fallback"
+      max_logging.log(f"Evaluation extracted_response ({used}): {extracted_response}")
 
     # Check exact correctness
     try:
@@ -146,8 +156,8 @@ def score_responses(tmvp_config, question, responses, answer):
         max_logging.log(f"Evaluation Exception: {e}")
         max_logging.log("SKIPPED")
 
-    # Check format correctness
-    if match_format.search(response) is not None:
+    # Check format correctness (requires the full <reasoning>...</reasoning><answer>...</answer> structure)
+    if full_match is not None:
       has_correct_format = True
 
     # Early exit if all criteria are met
diff --git a/src/maxtext/trainers/post_train/rl/utils_rl.py b/src/maxtext/trainers/post_train/rl/utils_rl.py
@@ -118,6 +118,19 @@ def get_match_format_regex(tmvp_config):
   return match_format
 
 
+def get_answer_fallback_regex(tmvp_config):
+  """Returns a compiled regex that finds the *last* answer tag in a completion.
+
+  Used as a fallback when the full <reasoning>...</reasoning><answer>...</answer>
+  format is incomplete (e.g. missing the closing reasoning tag).  The result
+  reward can still be computed independently from the format reward.
+  """
+  return re.compile(
+      rf"{re.escape(tmvp_config.solution_start_token)}(.+?){re.escape(tmvp_config.solution_end_token)}",
+      flags=re.MULTILINE | re.DOTALL,
+  )
+
+
 def match_format_exactly(prompts, completions, tmvp_config, **kargs):
   """
   Give the model a reward of tmvp_config.reward_exact_format_match points if the format matches exactly.
diff --git a/src/maxtext/utils/gcs_utils.py b/src/maxtext/utils/gcs_utils.py
@@ -18,6 +18,8 @@
 import os
 import socket
 from pathlib import Path
+from etils import epath
+import uuid
 
 import yaml
 
@@ -242,3 +244,45 @@ def write_dict_to_gcs_json(data_dict, file_path):
     blob.upload_from_string(json_string, content_type="application/json")
   except (ValueError, TypeError, RecursionError) as e:
     print(f"Failed to write json file at {file_path} with error: {str(e)}")
+
+
+def mkdir_and_check_permissions(p: str | epath.Path) -> epath.Path:
+  """Creates a directory if it doesn't exist and verifies write permissions.
+
+  This function prevents the program from hanging when an output directory is inaccessible. The standard
+  `epath.Path.mkdir` can hang or fail silently when pointed at a path in a non-existent or inaccessible GCS bucket.
+
+  For example, the following code can hang indefinitely:
+
+    from etils import epath
+    p = epath.Path("gs://no_such_bucket/path/to/output")
+    p.mkdir(exist_ok=True, parents=True)
+  """
+  if isinstance(p, str):
+    p = epath.Path(p)
+
+  if p.as_posix().startswith("gs://"):
+    bucket_name = p.parts[2]
+    try:
+      storage_client = storage.Client()
+      storage_client.get_bucket(bucket_name)
+    except Exception as e:
+      raise FileNotFoundError(f"GCS bucket 'gs://{bucket_name}' not found or accessible.") from e
+  p.mkdir(exist_ok=True, parents=True)
+  if not p.exists():
+    raise PermissionError(f"Failed to create the directory '{p}'. Please check that you have write access.")
+
+  # Verify write permissions by creating and deleting a temporary file.
+  # This handles the case where the directory exists but is not writable.
+  temp_file_path = p / f".write_test_{uuid.uuid4()}"
+  try:
+    temp_file_path.write_text("test")
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    raise PermissionError(f"Directory '{p}' exists, but is not writable. Please check permissions.") from e
+  finally:
+    try:
+      temp_file_path.unlink()  # Delete the temp file.
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass  # Suppress errors during cleanup to not hide the original error.
+
+  return p
diff --git a/tests/unit/gcs_utils_test.py b/tests/unit/gcs_utils_test.py
@@ -0,0 +1,127 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for GCS utility functions."""
+
+import unittest
+from unittest import mock
+import os
+import tempfile
+
+# Module to be tested
+from maxtext.utils import gcs_utils
+
+
+class GcsUtilsTest(unittest.TestCase):
+  """Unit tests for GCS utility functions."""
+
+  def test_add_trailing_slash(self):
+    """Tests the simple add_trailing_slash utility."""
+    self.assertEqual(gcs_utils.add_trailing_slash("a/b"), "a/b/")
+    self.assertEqual(gcs_utils.add_trailing_slash("a/b/"), "a/b/")
+
+  def test_mkdir_non_existing_dir(self):
+    """Tests that a non-existing directory is created and is empty."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+      new_dir_path = os.path.join(temp_dir, "new_dir")
+      self.assertFalse(os.path.exists(new_dir_path))
+
+      # Act
+      gcs_utils.mkdir_and_check_permissions(new_dir_path)
+
+      # Assert
+      self.assertTrue(os.path.isdir(new_dir_path))
+      self.assertEqual(os.listdir(new_dir_path), [])
+
+  def test_mkdir_existing_non_empty_dir(self):
+    """Tests that an existing, non-empty directory's contents are unmodified."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+      existing_dir_path = os.path.join(temp_dir, "existing_dir")
+      os.makedirs(existing_dir_path)
+      dummy_file_path = os.path.join(existing_dir_path, "dummy.txt")
+      with open(dummy_file_path, "w", encoding="utf-8") as f:
+        f.write("test")
+
+      # Act
+      gcs_utils.mkdir_and_check_permissions(existing_dir_path)
+
+      # Assert
+      self.assertTrue(os.path.isdir(existing_dir_path))
+      self.assertEqual(os.listdir(existing_dir_path), ["dummy.txt"])
+
+  def test_mkdir_existing_read_only_dir(self):
+    """Tests that a PermissionError is raised for a read-only directory."""
+    # /sys is a universally read-only directory, even for root
+    read_only_dir_path = "/sys"
+    with self.assertRaises(PermissionError):
+      gcs_utils.mkdir_and_check_permissions(read_only_dir_path)
+
+  def test_mkdir_read_only_parent_dir(self):
+    """Tests that a PermissionError is raised when the parent is read-only."""
+    # /sys is a universally read-only directory, even for root
+    new_dir_path = "/sys/new_dir"
+    with self.assertRaises(PermissionError):
+      gcs_utils.mkdir_and_check_permissions(new_dir_path)
+
+  @mock.patch("maxtext.utils.gcs_utils.storage.Client")
+  def test_mkdir_gcs_no_such_bucket(self, mock_storage_client):
+    """Tests that an exception is raised for a non-existent GCS bucket."""
+    mock_client_instance = mock_storage_client.return_value
+    mock_client_instance.get_bucket.side_effect = Exception("Bucket not found!")
+    gcs_path = "gs://no_such_bucket"
+
+    with self.assertRaises(FileNotFoundError):
+      gcs_utils.mkdir_and_check_permissions(gcs_path)
+    mock_client_instance.get_bucket.assert_called_with("no_such_bucket")
+
+  @mock.patch("maxtext.utils.gcs_utils.storage.Client")
+  def test_mkdir_gcs_no_such_bucket_with_path(self, mock_storage_client):
+    """Tests an exception for a non-existent bucket with a subdirectory."""
+    mock_client_instance = mock_storage_client.return_value
+    mock_client_instance.get_bucket.side_effect = Exception("Bucket not found!")
+    gcs_path = "gs://no_such_bucket/some/dir"
+
+    with self.assertRaises(FileNotFoundError):
+      gcs_utils.mkdir_and_check_permissions(gcs_path)
+    mock_client_instance.get_bucket.assert_called_with("no_such_bucket")
+
+  @mock.patch("maxtext.utils.gcs_utils.uuid.uuid4", return_value="test-uuid")
+  @mock.patch("maxtext.utils.gcs_utils.epath.Path")
+  @mock.patch("maxtext.utils.gcs_utils.storage.Client")
+  def test_mkdir_gcs_valid_bucket(self, mock_storage_client, mock_epath, mock_uuid):
+    """Tests that a valid GCS path is handled correctly without errors."""
+    # Arrange: Mock the GCS client to simulate a valid bucket
+    mock_client_instance = mock_storage_client.return_value
+
+    # Arrange: Mock epath to prevent real GCS calls
+    mock_path_instance = mock.MagicMock()
+    mock_path_instance.as_posix.return_value = "gs://valid_bucket/some/dir"
+    mock_path_instance.parts = ["gs:", "", "valid_bucket", "some", "dir"]
+    mock_path_instance.exists.return_value = True
+
+    mock_temp_file_instance = mock.MagicMock()
+    mock_path_instance.__truediv__.return_value = mock_temp_file_instance
+
+    mock_epath.return_value = mock_path_instance
+    gcs_path = "gs://valid_bucket/some/dir"
+
+    # Act
+    gcs_utils.mkdir_and_check_permissions(gcs_path)
+
+    # Assert
+    mock_client_instance.get_bucket.assert_called_with("valid_bucket")
+    mock_path_instance.mkdir.assert_called_with(exist_ok=True, parents=True)
+    mock_path_instance.exists.assert_called_once()
+    mock_temp_file_instance.write_text.assert_called_once_with("test")
+    mock_temp_file_instance.unlink.assert_called_once()
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
diff --git a/tests/unit/rl_utils_test.py b/tests/unit/rl_utils_test.py

Original file line number	Diff line number	Diff line change
`@@ -1318,7 +1318,7 @@ def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index):`
`1318`	`1318`	`)`
`1319`	`1319`
`1320`	`1320`	`# Sum up the partial outputs across the expert shards.`
`1321`		`- output = jnp.reshape(output, (-1, sequence_length, self.config.emb_dim))`
	`1321`	`+ output = jnp.reshape(output, (-1, sequence_length, self.config.emb_dim // self.get_tensor_parallelism_size()))`
`1322`	`1322`	`output = jax.lax.psum_scatter(output, self._expert_parallelism_name, scatter_dimension=0, tiled=True)`
`1323`	`1323`
`1324`	`1324`	`else:`