Merge pull request #1596 from AI-Hypercomputer:jacobplatin/fix-microbenchmark-tokenizer-issue

PiperOrigin-RevId: 748349490

Author: maxtext authors

MaxText/inference_microbenchmark.py

@@ -407,7 +407,7 @@ def run_benchmarks(config):
 
   text = config.prompt
   metadata = engine.get_tokenizer()
-  vocab = token_utils.load_vocab(metadata.path, metadata.extra_ids)
+  tokenizer_model = engine.build_tokenizer(metadata)
   rng, rng_init_decode = jax.random.split(rng)
 
   generate_executable, params, decode_state_executable = engine.aot_compile(params, pass_rng_shape=True)
@@ -429,8 +429,8 @@ def run_benchmarks(config):
     rng_shape = jax.ShapeDtypeStruct([4], jax.numpy.dtype("uint32"))
 
     for prefill_length in prefill_lengths:
-      prefill_tokens[prefill_length], prefill_true_lengths[prefill_length] = token_utils.tokenize_and_pad(
-          text, vocab, is_bos=True, prefill_lengths=[prefill_length]
+      prefill_tokens[prefill_length], prefill_true_lengths[prefill_length] = tokenizer_model.encode(
+          text, is_bos=True, prefill_lengths=[prefill_length]
       )
 
       key_shape = jax.ShapeDtypeStruct([prefill_length], jax.numpy.dtype("int32"))

MaxText/tests/grpo_trainer_correctness_test.py

@@ -82,31 +82,30 @@ def prepare_maxtext_inputs(input_str, tokenizer_model):
   return input_ids, input_segmentation, input_position, completion_segmentation
 
 
-
 class GrpoTrainerTest(unittest.TestCase):
 
   def setUp(self):
     super().setUp()
     command = [
-      "gsutil",
-      "cp",
-      "-r",
-      "gs://maxtext-dataset/hf/llama3.1-tokenizer",
-      os.path.join(os.path.dirname(PKG_DIR), "assets", ""),
+        "gsutil",
+        "cp",
+        "-r",
+        "gs://maxtext-dataset/hf/llama3.1-tokenizer",
+        os.path.join(os.path.dirname(PKG_DIR), "assets", ""),
     ]
     exit_code = subprocess.call(command, cwd=os.path.dirname(PKG_DIR))
     if exit_code != 0:
       raise ValueError(f"{command} failed with exit code: {exit_code}")
     self.config = pyconfig.initialize(
         [None, "MaxText/experimental/rl/grpo_trainer_test.yml"],
         run_name="unit_test_grpo_trainer",
-        tokenizer_path=os.path.join(os.path.dirname(PKG_DIR), 'assets', 'llama3.1-tokenizer'),
+        tokenizer_path=os.path.join(os.path.dirname(PKG_DIR), "assets", "llama3.1-tokenizer"),
         enable_checkpointing=False,
     )
     self.config_inference = pyconfig.initialize(
         [None, "MaxText/experimental/rl/grpo_trainer_test.yml"],
         run_name="unit_test_grpo_trainer_inference",
-        tokenizer_path=os.path.join(os.path.dirname(PKG_DIR), 'assets', 'llama3.1-tokenizer'),
+        tokenizer_path=os.path.join(os.path.dirname(PKG_DIR), "assets", "llama3.1-tokenizer"),
         enable_checkpointing=False,
         ici_tensor_parallelism=4,
         per_device_batch_size=self.config.per_device_batch_size * self.config.num_generations,
@@ -115,11 +114,11 @@ def setUp(self):
     self.atol = 1e-08
     self.rng = jax.random.PRNGKey(self.config.init_weights_seed)
     self.tokenizer_model = transformers.AutoTokenizer.from_pretrained(
-      self.config.tokenizer_path,
-      add_bos_token=self.config.add_bos,
-      add_eos_token=self.config.add_eos,
-      legacy=False,
-      padding_side="left"
+        self.config.tokenizer_path,
+        add_bos_token=self.config.add_bos,
+        add_eos_token=self.config.add_eos,
+        legacy=False,
+        padding_side="left",
     )
     self.tokenizer_model.add_special_tokens({"pad_token": "<pad>"})
 
@@ -135,22 +134,23 @@ def test_grpo_trainer_correctness(self):
     )
     # Obtain per-token logits.
     maxtext_per_token_logps, _ = compute_log_probs(
-      maxtext_model,
-      state.params,
-      input_ids,
-      input_position,
-      input_segmentation,
-      completion_segmentation,
-      self.config,
-      is_train=False,
-      rngs=self.rng,
+        maxtext_model,
+        state.params,
+        input_ids,
+        input_position,
+        input_segmentation,
+        completion_segmentation,
+        self.config,
+        is_train=False,
+        rngs=self.rng,
     )
-    jax.debug.print("maxtext_per_token_logps={maxtext_per_token_logps}",maxtext_per_token_logps=maxtext_per_token_logps)
-    jax.debug.print("golden_per_token_logps={golden_per_token_logps}",golden_per_token_logps=golden_data["maxtext_per_token_logps_no_ckpt_loading"])
-    golden_maxtext_logits = np.array(golden_data["maxtext_per_token_logps_no_ckpt_loading"])
-    self.assertTrue(
-      jnp.all(np.array(golden_data["input_ids"]) == np.array(input_ids[0]))
+    jax.debug.print("maxtext_per_token_logps={maxtext_per_token_logps}", maxtext_per_token_logps=maxtext_per_token_logps)
+    jax.debug.print(
+        "golden_per_token_logps={golden_per_token_logps}",
+        golden_per_token_logps=golden_data["maxtext_per_token_logps_no_ckpt_loading"],
     )
+    golden_maxtext_logits = np.array(golden_data["maxtext_per_token_logps_no_ckpt_loading"])
+    self.assertTrue(jnp.all(np.array(golden_data["input_ids"]) == np.array(input_ids[0])))
     self.assertTrue(
         jax.numpy.allclose(
             maxtext_per_token_logps[0],

MaxText/train.py

@@ -663,7 +663,7 @@ def setup_train_loop(config):
   record_goodput(recorder, config, recorder.record_training_preparation_start_time if recorder else None)
   data_iterator, eval_data_iterator = create_data_iterator(config, mesh)
 
-  context_parallel_size = mesh.shape['context']
+  context_parallel_size = mesh.shape["context"]
   # Check if context parallelism is being used with sequence packing
   if context_parallel_size > 1 and config.packing and config.dataset_type != "synthetic":
     raise ValueError(