Tunix v0.1.7

What's Changed

• Instrument agentic loop with perf v2 & perfetto improvements by @copybara-service[bot] in #1159
• add support for GCS endpoint for perf trace writing by @copybara-service[bot] in #1221
• [Tunix] Expose many training hyperparameters via argparse. by @copybara-service[bot] in #1243
• fix sglang warning by @copybara-service[bot] in #1242
• fix: return type of _train_step() in PeftTrainer. by @copybara-service[bot] in #1241
• remove extra from training prompt ids by @copybara-service[bot] in #1244
• Uiuc vlm pr compressed fixed by @immuntasir in #511
• [Tunix] Implement batch logging and add signal handlers for graceful shutdown for trajectory logger. by @copybara-service[bot] in #1245
• measure batch solve_all & solve_none by @copybara-service[bot] in #1247
• [Tunix] Update math eval script with vLLM sampler configuration to use vllm_jax backend and engine_kwargs. Adjust seeding for vLLm. by @copybara-service[bot] in #1246
• [Tunix] Defaults to best mesh configurations. by @copybara-service[bot] in #1250
• chore: start v0.1.7 development. by @copybara-service[bot] in #1248
• use timeline snapshot for export to avoid race condition by @copybara-service[bot] in #1240
• obtain old_logps in async rollout by @copybara-service[bot] in #1252
• [Tunix] Add requirements.txt file with vllm and tpu-inference versions by @copybara-service[bot] in #1253
• Add vllm rollout config for deepswe. by @copybara-service[bot] in #1222
• Speed-up unscanning with jax unstack. by @NicoGrande in #1256
• [Tunix] Remove eos token from user config. use default tokenizer.eos_id() directly. Besides for deepseek-distilled model we should have used <|end_of_sentence|> instead. by @copybara-service[bot] in #1267
• fix copybara by @copybara-service[bot] in #1266
• Support tunable parameter for submit train job in xpk by @copybara-service[bot] in #1269
• [Tunix] Set max_logprobs to 1 in vLLM rollout. by @copybara-service[bot] in #1270
• use background thread for writing trace by @copybara-service[bot] in #1249
• fix oldprobs. Remove patching rollout config. Add validation check at agentic learner init. by @copybara-service[bot] in #1271
• [Tunix] Fix aggregation functions for reward metrics. by @copybara-service[bot] in #1277
• [Tunix][Oncall] Update expected output in vLLM sampler test. by @copybara-service[bot] in #1281
• mixed precision for Qwen3 by @copybara-service[bot] in #1285
• enable flash attn in Tunix model by @copybara-service[bot] in #1273
• [Tunix] Update resharding to use latest pathwaysutils tests by @copybara-service[bot] in #1286
• make lane assignment helper fn and fix naming by @copybara-service[bot] in #1272
• Add more metrics for agentic rl. by @wang2yn84 in #1284
• [JAX] Suppress type errors found by pytype after correcting definition of jax.typing.ArrayLike. by @copybara-service[bot] in #1254
• [Tunix] Add jax.block_until_ready to test common utils. by @copybara-service[bot] in #1290
• [Tunix] Clamp sequence importance ratio in agentic GRPO learner. by @copybara-service[bot] in #1296
• [Tunix] Update expected output in vllm_sampler_test. by @copybara-service[bot] in #1300
• pass role to device assignment to writer by @copybara-service[bot] in #1274
• add flash attn to qwen3 model by @copybara-service[bot] in #1287
• Disable progress bar if stderr is not TTY and not running in a notebook. by @copybara-service[bot] in #1292
• remove step time metrics by @copybara-service[bot] in #1289
• make async timeline failure log errors by @copybara-service[bot] in #1301
• Role base track grouping by @copybara-service[bot] in #1261
• [Tunix] Use self.config.return_logprobs to determine if logprobs should be returned. by @copybara-service[bot] in #1304
• [Tunix] Add special handling for math answer grading. by @copybara-service[bot] in #1297
• nit by @copybara-service[bot] in #1293
• remove dummy reward fns for agentic learner. by @copybara-service[bot] in #1275
• Add support for dynamically setting the number of steps for GRPO. by @niting in #1257
• Add env time tracking by @copybara-service[bot] in #1302
• Add megablox MoE support by @copybara-service[bot] in #1306
• [Tunix] Hide internal helper functions for 1p only in reshard.py. by @copybara-service[bot] in #1291
• [Tunix] Add backend_kwargs to MetricsLoggerOptions. by @copybara-service[bot] in #1305
• add logging of example rewards with snipped output, controlled under debug=True flag from GRPOConfig by @copybara-service[bot] in #1298
• Code update by @copybara-service[bot] in #1314
• instrument environment interactions for perf metrics by @copybara-service[bot] in #1303
• update metrics docs with new perf v2 api by @copybara-service[bot] in #1312
• revert changes to reward_manager_test.py from b6e98ae by @andytwigg in #1319
• update news in readme by @copybara-service[bot] in #1320
• fix qwen2 mixed precision & notebook by @copybara-service[bot] in #1315
• [Tunix] Add support for loading models from Orbax checkpoints for evaluation of individual trained ckpts. by @copybara-service[bot] in #1318
• [Tunix] Always log kl regardless if it's enabled. by @copybara-service[bot] in #1313
• return logps in vanilla sampler by @copybara-service[bot] in #1328
• Code update by @copybara-service[bot] in #1327
• reward fn is only set at env reset. by @copybara-service[bot] in #1329
• adding s-noghabi to github codeowners by @copybara-service[bot] in #1308
• [Tunix] Update qwix dependency and TPU test image. by @copybara-service[bot] in #1145
• enable sequence parallelism by @copybara-service[bot] in #1332
• Add overlong filter to optionally filter overlong trajectory. by @copybara-service[bot] in #1283
• Add reward time and env time. by @copybara-service[bot] in #1310
• [Tunix] Update Pathways worker temporary storage and remote IDE port. by @copybara-service[bot] in #1343
• fix deepscaler sort failure by @copybara-service[bot] in #1342
• fix name of trace_dir in base_config by @copybara-service[bot] in #1346
• add support for repeating kv head tensors during weight sync. by @NicoGrande in #1251
• make cluster grouping names include multiple roles in perfetto by @copybara-service[bot] in #1347
• Add num_batches to base_config.yaml. by @copybara-service[bot] in #1349
• Adding force_compute_kl knob and fixing exception for kl computation by @copybara-service[bot] in #1344
• [Tunix] Update expected output in vLLM sampler test. by @copybara-service[bot] in #1357
• [Tunix] Add validation for vLLM server mode in AgenticRLLearner. by @copybara-service[bot] in #1354
• [Tunix] Add configurable loss aggregation mode and KL loss mode to DeepScaler. by @copybara-service[bot] in #1355
• instrument tunix main with perf v2 by @copybara-service[bot] in #1345
• [Tunix]Adds a command-line argument for maximum prompt length and configures the metrics logger to log training hyperparameters to Weights & Biases (wandb). by @copybara-service[bot] in #1356
• [Tunix] Add dual-clip on pg_loss in agentic GRPO learner and compute kl_loss before applying kl penalty. by @copybara-service[bot] in #1334
• [Tunix] Refactor loss aggregation to handle empty sequences and improve test coverage. by @copybara-service[bot] in #1358
• [Tunix] Update vLLM/tpu-inference commits. by @copybara-service[bot] in #1359
• [Tunix] Fix the bug of partial write being caught by a read operation. by @copybara-service[bot] in #1363
• make REMOTE_PW_PORT configurable by @copybara-service[bot] in #1360
• enable flash attention option in cli. by @copybara-service[bot] in #1339
• Remove "steps_per_sec" from tqdm train metrics. by @copybara-service[bot] in #1362
• Tunix gemma4 by @copybara-service[bot] in #1365
• gemma4 26b moe by @copybara-service[bot] in #1367
• use State instead of GraphState by @copybara-service[bot] in #1369
• update readme by @copybara-service[bot] in #1368
• [Tunix] Add retry logic to Hugging Face model existence checks. by @copybara-service[bot] in #1374
• mixed precision support in gemma4 by @copybara-service[bot] in #1372
• [Tunix] Add Pathways Cloud support to grpo_main. by @copybara-service[bot] in #1371
• Add huggingface dataset support. by @wang2yn84 in #1375
• [Tunix] Remove rollout_vllm_swap_space_size_gb from Tunix. by @copybara-service[bot] in #1383
• support flash attn in gemma4 by @copybara-service[bot] in #1376
• add knobs for perf metrics in grpo_demo_llama3_qwen2.py by @copybara-service[bot] in #1348
• move filtering degenerate groups to process results by @copybara-service[bot] in #1381
• Avoid implicit conversion of jax arrays to dtypes. by @copybara-service[bot] in #1382
• [Tunix] Add DeepScaleR CLI recipe by @wang2yn84 in #1373
• Move RLOO advantage estimator from GRPOLearner to AgenticGRPOLearner. by @copybara-service[bot] in #1340
• Remove the step_time parameter from the SFT TrainingHooks.on_train_step_end function. by @copybara-service[bot] in #1395
• Convert deepswe to cli by @wang2yn84 in #1392
• [Tunix] Add documentation for Trajectory Logging in Tunix. by @copybara-service[bot] in #1405
• improve rl_cluster init corner case handling and test coverage by @copybara-service[bot] in #1311
• Check if GCS path and raise error in v1 perf tracer by @copybara-service[bot] in #1389
• Include vLLM and tpu inference to Tunix build. by @wang2yn84 in #1411
• add multi-node setup doc by @copybara-service[bot] in #1412
• Bug fix for trajectory engine by @copybara-service[bot] in #1394
• Fix broken base_config for RL training. by @niting in #1413
• Add DeepSWE evaluation script with ActionGuard. by @copybara-service[bot] in #1406
• mock rollout engine by @copybara-service[bot] in #1391
• Allow various length distributions in mock rollout by @copybara-service[bot] in #1398
• update docs with mock rollout engine by @copybara-service[bot] in #1400
• [tunix/sft] Avoid eager re-sharding of globally distributed arrays by @copybara-service[bot] in #1325
• Fix agentic cli dataloading by @wang2yn84 in #1416
• Fix multiple dataset issue. by @wang2yn84 in #1419
• check if limit reached more proactivately by @copybara-service[bot] in #1388
• Internal by @copybara-service[bot] in #1390
• add test for export failing by @copybara-service[bot] in #1421
• [Tunix] Bump up the vllm/tpu-inference commit by @copybara-service[bot] in #1424
• Adding support for fused moe kernels in sampler. by @NicoGrande in #1385
• Save timelines per step by @copybara-service[bot] in #1423
• use current and commited step in perf tracer by @copybara-service[bot] in #1425
• [Tunix] Update Dockerfile dependencies and installation order. by @copybara-service[bot] in #1428
• add timeline sequentializing to perf tracer by @copybara-service[bot] in #1422
• Update training dataset to R2E-Gym-Subset. by @copybara-service[bot] in #1426
• Prevent adding environment tokens for terminal steps in trajectory collection. by @copybara-service[bot] in #1427
• Fix deleting destination buffers during reshard. by @NicoGrande in #1430
• Add sequence packing by @copybara-service[bot] in #1337
• [Tunix] Update logging for Pathways checkpointing based on persistence API usage. by @copybara-service[bot] in #1434
• Disable sglang_jax with LoRA test in nightly regression by @sizhit2 in #1435
• Fix env logging metrics tab by @sizhit2 in #1437
• add model registry by @copybara-service[bot] in #1446
• algo fix by @copybara-service[bot] in #1450
• fix error message for max steps by @copybara-service[bot] in #1453
• set vllm logprobs_mode by @copybara-service[bot] in #1451
• Fix RL cluster mesh logic for null objects + Update model name in Gemma2 config for SFT and RL and add appropriate data_source. by @copybara-service[bot] in #1444
• update nightly regression test by @copybara-service[bot] in #1455
• add gemma 4 to tests by @copybara-service[bot] in #1447
• Integrate MaxText Model Source in AutoModel by @sizhit2 in #1442
• minor speedup by @copybara-service[bot] in #1460
• Disable failed vLLM test in TPU nightly regression. by @copybara-service[bot] in #1459
• [Tunix] Fix the nightly regression test path. by @copybara-service[bot] in #1463
• [Tunix] Add the repository root to sys.path, so from tunix.cli import grpo_main resolves against the local tunix package in the workspace by @copybara-service[bot] in #1464
• [Tunix] Remove pytree=False from Llama3 model. by @copybara-service[bot] in #1466
• populate deepswe train branch to main by @copybara-service[bot] in #1420
• Remove @nnx.jit from Gemma4 model call. by @copybara-service[bot] in #1468
• Adding support for TP for fused MoE kernel. by @NicoGrande in #1448
• Fix Arbitrary Code Execution via eval() in Reward Calculation. by @copybara-service[bot] in #1471
• fix qwen parser thinking mode by @copybara-service[bot] in #1475
• [Tunix] Add the missing train_fraction to cli scripts. by @copybara-service[bot] in #1477
• [Tunix] Fix the lora filter for llama and qwen. by @copybara-service[bot] in #1478
• Update nightly run env by @copybara-service[bot] in #1470
• bug fix for agentic reward calculation. by @copybara-service[bot] in #1472
• Fixed a test failure. by @copybara-service[bot] in #1473
• Fix Command Injection in TPU single host script. by @copybara-service[bot] in #1480
• [Tunix] Unpin jax and flax versions for sglnag tests in CI. by @copybara-service[bot] in #1483
• Disable more cli tests in nightly regression and add TODO by @copybara-service[bot] in #1492
• add control token sanitization to prevent prompt injection. by @copybara-service[bot] in #1489
• set download path for all cli scripts by @copybara-service[bot] in #1491
• Remove context_ratio from agentic GRPO config. by @copybara-service[bot] in #1495
• Dynamic batch should be write back to config if it's inferred. by @wang2yn84 in #1479
• Fix flash attention issue on decode path. by @wang2yn84 in #1481
• Add weights mappings and sync from Tunix to VLLM-tpu-jax for Gemma4 models. by @copybara-service[bot] in #1486
• Add "seq-mean-token-sum" loss aggregation mode. by @copybara-service[bot] in #1476
• Fix mesh in cli scripts by @copybara-service[bot] in #1485
• [Tunix] Remove trailing whitespace in cpu-tests.yml. by @copybara-service[bot] in #1500
• Add degenerate_group_masking option to DeepSWE training script. by @copybara-service[bot] in #1488
• make BaseTool immutable. by @copybara-service[bot] in #1497
• Relax the constraint on batch dim for computing ref_logps by @copybara-service[bot] in #1431
• Add sharding to the embedder parameters in Gemma4 model. by @copybara-service[bot] in #1498
• Allow configurable micro-batching for compute_logps in agentic RL. by @copybara-service[bot] in #1436
• Bump qwix version to 0.1.6 by @copybara-service[bot] in #1504
• Add option to recompute old policy logprobs on the trainer. by @copybara-service[bot] in #1487
• Specify shell: bash for TPU nightly RL script by @copybara-service[bot] in #1507
• yaml based configs for llama cli scripts by @copybara-service[bot] in #1452
• [Tunix] Update Dockerfile to fix qwix version conflict. by @copybara-service[bot] in #1509
• Standardize Tunix model naming to align with HuggingFace conventions. by @copybara-service[bot] in #1433
• yaml based configs for gemma3 cli scripts by @copybara-service[bot] in #1508
• Stabilize Tunix SFT pipelines: resolve NaN divergence and persistent state corruption by @copybara-service[bot] in #1511
• Fix non agentic colocate mode in cli. by @wang2yn84 in #1457
• Add chunked resharding and target buffer deletion options for tunix models' weight sync. by @copybara-service[bot] in #1499
• [Tunix] minor bug fix -- respect completion_mask when _compute_ref_per_token_logps. by @copybara-service[bot] in #1501
• add orpo_trainer test to build file by @copybara-service[bot] in #1515
• Add run_latest CI job for Tunix nightly regression by @copybara-service[bot] in #1510
• Remove stale context_ratio from GRPO configurations. by @copybara-service[bot] in #1519
• fix PEFT trainer early exit by @copybara-service[bot] in #1418
• fix naming in scripts by @copybara-service[bot] in #1521
• Fix the nightly regression to report error status correctly. by @wang2yn84 in #1522
• tighten atol rtol on common_test by @copybara-service[bot] in #1524
• Make it easier to see individual test results in TPU nightly by @copybara-service[bot] in #1525
• Support remat for gemma4 by @copybara-service[bot] in #1531
• [Tunix] Replace HTML-based image centering with standard g3doc Markdown image syntax. by @copybara-service[bot] in #1503
• [Tunix] replace info logging to debug logging on state. by @copybara-service[bot] in #1527
• Adjust top_p sampling to compute top-p of top-k logits, rather than top-k of top-p logits, matching vllm implementation by @copybara-service[bot] in #1528
• support gemma4 in automodel by @copybara-service[bot] in #1533
• Agentic GRPO: TIS correction, eval dedup, flash-attn segment_ids by @colincai-mc in #1523
• [Tunix] Update TPU nightly regression workflow to build vLLM/tpu-inference from source. by @copybara-service[bot] in #1532
• Gemma4 sampling optimizations. by @copybara-service[bot] in #1415
• [Tunix] Fix weight mapping for gemma4 E2B model. by @copybara-service[bot] in #1540
• Simplify logps compute by @copybara-service[bot] in #1548
• [Tunix]Ensure anchor policy state is on device before resharding as pathways reshard doesn't support resharding between different memory kinds. by @copybara-service[bot] in #1549
• Enhance performance and avoid OOM for data parallelism over a large number of devices O(100) by @copybara-service[bot] in #1550
• [Tunix] minor fix on anchor trainer state. by @copybara-service[bot] in #1557
• [Tunix] Default degenerate_group_masking to false. by @copybara-service[bot] in #1559
• Add orbax-checkpoint to tunix dependencies. by @copybara-service[bot] in #1462
• Introduce checkpoint_options to manage options for V1 Orbax Checkpointing. by @copybara-service[bot] in #1408
• Migrate Tunix's usage of Orbax V0 Checkpoint Manager to V1 Checkpointer. by @copybara-service[bot] in #1409
• Update Checkpointing Options construction through CLI configuration to move from Orbax V0 CheckpointManagerOptions to Tunix's CheckpointOptions. by @copybara-service[bot] in #1410
• [Tunix] Ensure TITO for prompt tokens. Now the trajectory prompt tokens is updated to what's used for rollout sampler. by @copybara-service[bot] in #1561
• Fix GSM8K reward function logic and string metrics logging. by @copybara-service[bot] in #1558
• Add Gemma 4 E2B config and support for local HuggingFace tokenizers. by @copybara-service[bot] in #1562
• Add optional output_clamp_value to compute_kl_divergence by @py4 in #1560
• Update Checkpointing Options construction through CLI configuration to move from Orbax V0 CheckpointManagerOptions to Tunix's CheckpointOptions. by @copybara-service[bot] in #1563
• Migrate Tunix's usage of Orbax V0 Checkpoint Manager to V1 Checkpointer. by @copybara-service[bot] in #1564
• Fixes the tfds prompt not sanitizated issue for agentic path. by @wang2yn84 in #1556
• Topology aware mesh creation by @wang2yn84 in #1465
• Fix Tunix package test by @copybara-service[bot] in #1568
• Fix run_prod CI: cap safetensors<0.8 to fix .ctypes breakage by @copybara-service[bot] in #1572
• Add Gemma 4 Orbax checkpoint loading support to Tunix. by @copybara-service[bot] in #1553
• pin safetensors in project toml by @copybara-service[bot] in #1573
• Undo transformers<=4.57.1 pin by @ringohoffman in #1566

New Contributors

• @niting made their first contribution in #1257
• @colincai-mc made their first contribution in #1523
• @py4 made their first contribution in #1560
• @ringohoffman made their first contribution in #1566

Full Changelog: v0.1.6...v0.1.7