Merge pull request #2335 from AI-Hypercomputer:tunix_fix

PiperOrigin-RevId: 806064971

Author: Google-ML-Automation

pyproject.toml

@@ -236,7 +236,7 @@ dependencies = [
     "tqdm>=4.67.1",
     "transformers>=4.56.1",
     "treescope>=0.1.10",
-    "tunix @ https://github.com/google/tunix/archive/4c5561be36d8a2f1f0858c2685554ca4e1a65fd2.zip",
+    "tunix @ https://github.com/google/tunix/archive/d770659621eb16ef6588268e26fa687fa068df20.zip",
     "typing-extensions>=4.14.1",
     "typing-inspection>=0.4.1",
     "tzdata>=2025.2",

requirements.txt

@@ -36,7 +36,7 @@ tensorflow-text
 tensorflow
 tiktoken
 transformers
-tunix @ https://github.com/google/tunix/archive/4c5561be36d8a2f1f0858c2685554ca4e1a65fd2.zip
+tunix @ https://github.com/google/tunix/archive/d770659621eb16ef6588268e26fa687fa068df20.zip
 google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/daedc21c393f23449fb54ddc4f75fca34348ea9c.zip
 mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
 qwix @ https://github.com/google/qwix/archive/f2fd7b9114ff8d09e5b0131a453351578502da8a.zip

src/MaxText/examples/grpo_llama3_demo.py

@@ -474,7 +474,6 @@ def get_ref_maxtext_model(config):
 if DEBUG:
   print("Model initialized successfully")
   print(f"Model mesh shape: {mesh_policy.shape}")
-  print(f"Model config: {model_config_policy}")
 
   # Sanity check that weights are loaded correctly
   _maxtext_state_flatten = nnx.state(llama3_1_8b_policy).flat_state()
@@ -955,7 +954,7 @@ def evaluate(
   # verify if vllm sampler works
   output = rl_cluster.rollout.generate(
       ["The capital of France is"],
-      rollout_config=RolloutConfig(n=1, max_tokens_to_generate=64, temperature=0.1),
+      rollout_config=RolloutConfig(max_tokens_to_generate=64, temperature=0.1),
   )
 
   print(f"Output: {output}")