Tunix v0.1.2: Expanded Model Support and Enhanced Flexibility

This release of Tunix introduces support for new models, enhances core functionalities for more flexible and efficient workflows, and includes several important fixes.

Highlights

• Expanded Model Support: We've added a configuration for qwen-8b and ported the Llama3 example to the Tunix CLI. Additionally, GRPO disaggregated llama3.1-70b is now supported through MaxText, including checkpoint saving.
• Enhanced Flexibility: Users can now specify a different data type for the rollout model and take advantage of more flexible PyTree support in the checkpoint manager. This release also introduces flexible collect modes and tokenization support, along with support for multiple EOS tokens in the vanilla sampler.

Other Changes

• Downgraded Jax version to 0.7.1 in prod mode due to performance regression, dev mode still supports Jax v0.7.2
• Fixes to the front page pip install command and GRPO examples.
• Improvements to the checkpoint manager and resharding library.
• Added a backward compatibility test for Orbax checkpoint restoration.
• Various code simplifications, refactoring, and documentation updates.

Full Changelog: v0.1.1...v0.1.2

What's Changed

• [Tunix] Allow specifying a different data type for the rollout model. by @copybara-service[bot] in #513
• Fix the front page pip install command. by @wang2yn84 in #515
• Remove prompt_template.py by @copybara-service[bot] in #514
• Tool code update by @copybara-service[bot] in #471
• simplify micro batching config by @copybara-service[bot] in #516
• Add explicit imports for specific TFDS datasets. by @copybara-service[bot] in #519
• Adding a qwen-8b config by @copybara-service[bot] in #522
• Environment code update by @copybara-service[bot] in #512
• [Tunix] Update the checkpoint manager with more flexible PyTree. by @copybara-service[bot] in #337
• Add BibTeX by @copybara-service[bot] in #528
• Ensure github workflow failure block on presubmit by @copybara-service[bot] in #487
• [Tunix] fix the grpo example which blocks copybara presubmit by @copybara-service[bot] in #531
• [Tunix] Add backward compatibility test for Orbax checkpoint restoration. by @copybara-service[bot] in #530
• [Tunix] Update reshard lib to respect logical axis rules. by @copybara-service[bot] in #518
• support multiple eos tokens in vanilla sampler by @copybara-service[bot] in #525
• Add sft smoke test. by @copybara-service[bot] in #533
• Port Precur.AI llama3 example to Tunix CLI by @copybara-service[bot] in #520
• Make GRPO disaggregated llama3.1 70b work with pathways including ckpt saving by @A9isha in #527
• Rename 'convert_messages_to_tokens_and_masks' to 'tokenize_and_generate_masks' by @copybara-service[bot] in #539
• Change docker version to jax0.7.1_rev1 by @copybara-service[bot] in #544
• Add flexible collect modes and tokenization support. by @copybara-service[bot] in #526
• fix mtnt import by @copybara-service[bot] in #546
• update function docstring for tokenize_and_generate_masks by @copybara-service[bot] in #547
• Move grpo shell scripts. by @copybara-service[bot] in #543
• minor simplification by @copybara-service[bot] in #541

Full Changelog: v0.1.1...v0.1.2