google-tunix 0.0.1__py3-none-any.whl

google_tunix-0.0.1.dist-info/METADATA

@@ -0,0 +1,155 @@
+Metadata-Version: 2.4
+Name: google-tunix
+Version: 0.0.1
+Summary: A lightweight JAX-native LLM post-training framework.
+Author-email: Tunix Developers <tunix-dev@google.com>
+License-Expression: Apache-2.0
+Project-URL: Source, https://github.com/google/tunix
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: datasets
+Requires-Dist: flax
+Requires-Dist: grain
+Requires-Dist: huggingface_hub
+Requires-Dist: jax
+Requires-Dist: jax[tpu]
+Requires-Dist: jaxtyping
+Requires-Dist: kagglehub
+Requires-Dist: omegaconf
+Requires-Dist: qwix
+Requires-Dist: sentencepiece
+Requires-Dist: tensorboardX
+Requires-Dist: tensorflow_datasets
+Requires-Dist: tqdm
+Requires-Dist: transformers
+Provides-Extra: docs
+Requires-Dist: sphinx>=8.2.3; extra == "docs"
+Requires-Dist: sphinx-book-theme>=1.1.4; extra == "docs"
+Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
+Requires-Dist: ipython>=8.8.0; extra == "docs"
+Requires-Dist: myst-nb>=1.3.0; extra == "docs"
+Requires-Dist: matplotlib>=3.10.0; extra == "docs"
+Requires-Dist: sphinx-gallery>=0.19.0; extra == "docs"
+Requires-Dist: sphinx-collections>=0.0.1; extra == "docs"
+Requires-Dist: sphinx_contributors; extra == "docs"
+Dynamic: license-file
+
+# Tunix: A JAX-native LLM Post-Training Library
+
+**Tunix(Tune-in-JAX)** is a JAX based library designed to streamline the
+post-training of Large Language Models. It provides efficient and scalable
+supports for:
+
+- **Supervised Fine-Tuning**
+- **Reinforcement Learning (RL)**
+- **Knowledge Distillation**
+
+Tunix leverages the power of JAX for accelerated computation and seamless
+integration with JAX-based modeling framework
+[Flax NNX](https://flax.readthedocs.io/en/latest/nnx_basics.html).
+
+**Current Status: Early Development**
+
+Tunix is in early development. We're actively working to expand its
+capabilities, usability and improve its performance. Stay tuned for upcoming
+updates and new features!
+
+## Key Features & Highlights
+
+Tunix is still under development, here's a glimpse of the current features:
+
+- **Supervised Fine-Tuning:**
+  - Full Weights Fine-Tuning
+  - Parameter-Efficient Fine-Tuning (PEFT) with LoRA/Q-LoRA Layers
+- **Reinforcement Learning (RL):**
+  - Proximal Policy Optimization (PPO)
+  - Group Relative Policy Optimization (GRPO)
+  - Token-level Group Sequence Policy Optimization (GSPO-token)
+- **Preference Fine-Tuning:**
+  - Preference alignments with Direct Preference Optimization (DPO)
+- **Knowledge Distillation:**
+  - Logit Strategy: A classic approach where the student learns to match the
+    teacher's output probability distribution.
+  - Attention Transfer & Projection Strategies: Methods to align the attention
+    mechanisms between the student and teacher models.
+  - Feature Pooling & Projection Strategies: General techniques for matching
+    intermediate feature representations, even between models of different
+    architectures.
+- **Modularity:**
+  - Components are designed to be reusable and composable
+  - Easy to customize and extend
+- **Efficiency:**
+  - Native support of common model sharding strategies such as DP, FSDP and TP
+  - Designed for distributed training on accelerators (TPU)
+
+## Upcoming
+
+- **Agentic RL Training:**
+  - Async Rollout
+  - Multi-turn & multi-step support
+  - Tool usage
+- **Advanced Algorithms:**
+  - Addtional state-of-the-art RL and distillation algorithms
+- **Scalability:**
+  - Multi-host distributed training
+  - Optimized rollout with vLLM
+- **User Guides:**
+  - More advanced RL recipe
+
+## Installation
+
+Tunix doesn't have a PyPI package yet. To use Tunix, you need to install from
+GitHub directly.
+
+```sh
+pip install git+https://github.com/google/tunix
+```
+
+## Getting Started
+
+To get started, we have a bunch of detailed examples and tutorials.
+
+- [PEFT Gemma with QLoRA](https://github.com/google/tunix/blob/main/examples/qlora_demo.ipynb)
+- [Training Gemma on grade school Math problems using GRPO](https://github.com/google/tunix/blob/main/examples/grpo_demo.ipynb)
+- [Logit Distillation using Gemma models](https://github.com/google/tunix/blob/main/examples/logit_distillation.ipynb)
+
+To setup Jupyter notebook on single host GCP TPU VM, please refer to the
+[setup script](https://github.com/google/tunix/blob/main/scripts/setup_notebook_tpu_single_host.sh).
+
+We plan to provide clear, concise documentation and more examples in the near
+future.
+
+## Contributing and Feedbacks
+
+We welcome contributions! As Tunix is in early development, the contribution
+process is still being formalized. A rough draft of the contribution process is
+present [here](https://github.com/google/tunix/blob/main/CONTRIBUTING.md). In
+the meantime, you can make feature requests, report issues and ask questions in
+our
+[Tunix GitHub discussion forum](https://github.com/google/tunix/discussions).
+
+## Collaborations and Partnership
+
+[GRL](https://github.com/lmgame-org/GRL/blob/tunix_integration_dev/README.md)
+(Game Reinforcement Learning), developed by
+[Hao AI Lab](https://hao-ai-lab.github.io/) from UCSD, is an open-source
+framework for post-training large language models through multi-turn RL on
+challenging games. In collaboration with Tunix, GRL integrates seamless TPU
+support—letting users quickly run scalable, reproducible RL experiments (like
+PPO rollouts on Qwen2.5-0.5B-Instruct) on TPU v4 meshes with
+[minimal setup](https://github.com/lmgame-org/GRL/blob/tunix_integration_dev/README.md#5-launch-the-quick-test-defaults-to-qwen2505b-supports-4-tpu-v4-with-mesh-22).
+This partnership empowers the community to push LLM capabilities further,
+combining Tunix’s optimized TPU runtime with GRL’s flexible game RL pipeline for
+cutting-edge research and easy reproducibility.
+
+## Stay Tuned!
+
+Thank you for your interest in Tunix. We're working hard to bring you a powerful
+and efficient library for LLM post-training. Please follow our progress and
+check back for updates!

google_tunix-0.0.1.dist-info/RECORD

@@ -0,0 +1,85 @@
+google_tunix-0.0.1.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+tunix/__init__.py,sha256=I_aCYRv7VOXgsrQ30fZVE36913XQBR-bRtMOY33fhsk,1295
+tunix/distillation/__init__.py,sha256=IZXRKcxXQlmirRKz3RMD_VkAjQIW5tG55vgIem498WM,869
+tunix/distillation/distillation_trainer.py,sha256=eQV3b500_TwFlyy1l4F9AzFn80NyNxIFW3RLpFCfAuc,4205
+tunix/distillation/feature_extraction/__init__.py,sha256=5ule3xq3xUjYqDVMzAfbE833Ibg3mUDpzAykHAnZuTY,1386
+tunix/distillation/feature_extraction/pooling.py,sha256=PRRDyQ3OmzUP0PnxAgWvW5hEvTZz8QmIlaxQi-lLLmQ,4661
+tunix/distillation/feature_extraction/projection.py,sha256=NietUThUfij3UcHW-oS-zkLHvtguBlDXDnSmQOiyf0k,4853
+tunix/distillation/feature_extraction/sowed_module.py,sha256=B66GzK7Me07iObLyA9ToOFj2J9Q0XngWyMdchoxTCiE,6210
+tunix/distillation/strategies/__init__.py,sha256=C-NHVX7IBAsvbkb37YFU4dJMYVuDYoERMcv7oRlqrC0,1200
+tunix/distillation/strategies/attention.py,sha256=7z9DYzkP-OSpJCDwkZuAKr8_m2oV02VQJE8mmi-i54Y,4272
+tunix/distillation/strategies/base_strategy.py,sha256=gIT-2DgTf-1hTMoEoOgUxMOHbb3nxosKVr3Rwx9mO6U,4625
+tunix/distillation/strategies/feature_pooling.py,sha256=itnHZ0QINnKBjg8-4c1Eewq6Gm3sYiEWAB2gSn2mkhU,6823
+tunix/distillation/strategies/feature_projection.py,sha256=0zzJOSM_76oB-TwWwYUBXUEvORGRUIX4aV6S_7L4DM4,6497
+tunix/distillation/strategies/logit.py,sha256=Z4x7JGY2XT43gOiSDd2bSkntUgWZdeNqog8su2gBhE8,4138
+tunix/examples/data/translation_dataset.py,sha256=2g60f9uN9JykecSO3WHKeZ8KfKzv2-BPUYOmUvkhdyo,9271
+tunix/generate/base_sampler.py,sha256=kL3IXt_yik1WPGTbxoXeZQ3Pw8Es3QRNRhkWINkrDWE,2054
+tunix/generate/beam_search.py,sha256=se9zBmpBOoEXENgaJPgcPMhQoN77jYbJdk2ZK8_M-Ps,9722
+tunix/generate/sampler.py,sha256=x52acyTzOcck6HpSKyXov_caFRuph8Wb2__v6cKsuWA,27007
+tunix/generate/tokenizer_adapter.py,sha256=zsEsjOZvg32gb9SfWUVQT8yibCOtG4XSWV6zP_-cfVA,4326
+tunix/generate/utils.py,sha256=2R3ILV9TLAtHZZt67SkKiyT7DBHzsk6o753iDH59nbI,18774
+tunix/generate/vllm_sampler.py,sha256=IL1_burNb72j46GSmheGwqHTPyist3IrYWX8q2BEGJ0,10499
+tunix/models/safetensors_loader.py,sha256=gyF6230EW7qaBJVtQZLTy5tbxaMrf4xUh27xrGl7Ips,5026
+tunix/models/gemma/gemma.py,sha256=4ks6sLHxrQyEIA7bJV2oPKEQNYe5mAwD9u2PJ4NexjQ,28652
+tunix/models/gemma/params.py,sha256=lHQa2cpMOKk6v1gnIKKe-IgihI0U1UYlJOo3uEAxJ8U,2442
+tunix/models/gemma/params_safetensors.py,sha256=kN3YfKSNjKgZ7lSauVCjPdcSfrEjrd8huwQTttd2V9g,6905
+tunix/models/gemma/sampler.py,sha256=w2xlax3N_WLEUb7fWebT5ZAwotU19nhwIwiIs0LS5W4,21736
+tunix/models/gemma3/model.py,sha256=5BrfnL7nlzMcpNZHOFOmSpJq5TP3iNNssrclGOBC-_A,24243
+tunix/models/gemma3/params.py,sha256=YCKf68VZONa0Vopm9e1-Kearf8AVoCte09DXE-IFinI,5501
+tunix/models/gemma3/params_safetensors.py,sha256=cl0-Bfbs4GvGC6oGXaWm1xfI4Op2fDDx5JYetEVmAWU,7442
+tunix/models/llama3/model.py,sha256=HKt1AIoPBx1DOqdLyn1JZatZW1uQKuOgeYegRm7dkdc,23319
+tunix/models/llama3/params.py,sha256=7FJDCj56nRptkSqUbdT-T_I4y2dEaZsC3o7WF4myuHY,3374
+tunix/models/qwen2/model.py,sha256=YyvHdNZiHLaQ3gGb1Dx_NiA1C5k0KOKKRtElwphLs2o,22266
+tunix/models/qwen2/params.py,sha256=i_-W76RpqvN_bc4BoZ7P3ApZ_1Gd3HufEVYyASdPBQs,3483
+tunix/models/qwen3/model.py,sha256=3GCwScMYAeHl_Zl_bxFGp5HPlwYp9dyal63_6qeMnKo,19697
+tunix/models/qwen3/params.py,sha256=XQ-nVYQ9UwJ1kDGEK0nA4OlyFgmtDQ3yRlJnP7LKAmU,4656
+tunix/oss/utils.py,sha256=JfoIYpHkB4qIM1JgoR8SeIuO5TzJ14z-8m5o2CsRPZY,859
+tunix/rl/common.py,sha256=33OFG8NbXPzRypbyd-TpjCgWsEu1_ZTaYjs013RSw-k,10151
+tunix/rl/reshard.py,sha256=uAe3r3h0W4NHsxNIHteUZwC74YgHFPWyYXeBv4ObHJM,5992
+tunix/rl/rl_cluster.py,sha256=pnHZDwMHoI12OhK9KiGNlrMLOXyuIwpajuIHURncSxo,28079
+tunix/rl/rl_learner.py,sha256=kkBXI_enhYyPegC7t9ShANrlyD3DaMhRMQek1sVpJiI,22379
+tunix/rl/trainer.py,sha256=ZjwPpTnI90YXKJE49fWS_dSys16GU-EOpWDJpMMcRFY,3050
+tunix/rl/utils.py,sha256=Km9WzQmuL8xHqxwnmL96-XFFZ3xgaV8rpKvaHpcr-uE,9232
+tunix/rl/experimental/agentic/agents/agent_types.py,sha256=trqLQsMZYXdQyzTMh4bkdwowyJgb1UEwXRje5wsH6ps,3207
+tunix/rl/experimental/agentic/agents/base_agent.py,sha256=hzPUb_xFxhBl7YnYh42s5I0LrFjV3HzyktK4sk_6SXY,8376
+tunix/rl/experimental/agentic/agents/tool_agent.py,sha256=Ym7s7YWlpBPtmrRdutM7jGKHUxzOzJlmb9lTKpwNuPc,10089
+tunix/rl/experimental/agentic/environments/base_environment.py,sha256=w-MS2coPItpKFPY0at9MlriUOX696jdIxJ0WE-N5t0A,5538
+tunix/rl/experimental/agentic/environments/tool_environment.py,sha256=KWKZ3tqh4hMLRaFK662XgZIVSX_0B9d1iMuNmCalJ_M,8412
+tunix/rl/experimental/agentic/parser/tool_parser/gemini_parser.py,sha256=aczEg3xgBZuaCDoRN_vv0JWbGAR1A1eMCoxSQqXJ8gA,601
+tunix/rl/experimental/agentic/parser/tool_parser/qwen_parser.py,sha256=Z5LOCzQ94LMzn5S5olocXcfsMM2AJXpK_68Q7jl-Xc8,3306
+tunix/rl/experimental/agentic/parser/tool_parser/tool_parser_base.py,sha256=7ht1eiuieuhEUY2-Y8tgCo5jqcnfvZRV0R3MuDcXVOw,2605
+tunix/rl/experimental/agentic/parser/tool_parser/tool_parser_registry.py,sha256=3KX5uvhrbSVxDCDt2iTLy177qXzLeNNQC5-xLFOp84s,684
+tunix/rl/experimental/agentic/prompts/prompt_template.py,sha256=JTFXmy3axxnk9Vd3gw7zW7UItaGsFEnzrTLVrygO5xA,1414
+tunix/rl/experimental/agentic/rewards/reward.py,sha256=3srEyf7QkUyJj6QX6cWRXTAnRCVk90MjIXemWPh15u4,4354
+tunix/rl/experimental/agentic/rewards/reward_types.py,sha256=eL5q5n0C7ArLMKx-CXRA8Euz6170Bqwd8trZEx7yvaA,1798
+tunix/rl/experimental/agentic/tools/base_tool.py,sha256=5dJt-C1BTrplZZA_nKxDyKjlUi6sEtyzVXZBaHPIvRY,7238
+tunix/rl/experimental/agentic/tools/calculator_tool.py,sha256=euTcCMkeTvN1GiDTmC34bSxF0QnKd6ioQ2FjWAlOcU8,3483
+tunix/rl/experimental/agentic/tools/tool_manager.py,sha256=8HS3kQZ6gvJ3S6kjiTysbuBiDS8samDSmEdV63ffKD8,6586
+tunix/rl/experimental/agentic/trajectory/trajectory_collect_engine.py,sha256=njJvTU6X2ery6z6p7erGBm9Hn_Nc7925o-O5cEEJhe4,8494
+tunix/rl/grpo/grpo_helpers.py,sha256=A4Qr7N4KbfZuM5jJ_NXqAPLyOQjV5WNL3fKlC_U6F8s,1256
+tunix/rl/grpo/grpo_learner.py,sha256=xHS3DvCnyrtdLIbfedyM-krHWCNDxpN6jNa4qpw02I8,16001
+tunix/rl/inference/inference_worker.py,sha256=XoEXf0uYNfGaRk-Ao4SlL5XyDsEreiyt0TOgOCFOVAM,2831
+tunix/rl/ppo/ppo_helpers.py,sha256=IgGw5T4POG90AdQHxxZ3HzM-_SUckSew9hsGUOMGa6w,4199
+tunix/rl/ppo/ppo_learner.py,sha256=V7L3GBhrpon7iOwiH12M3gnivM9O9lt2Ohi-4U9Zda8,19063
+tunix/rl/queue/data_queue.py,sha256=nEH3cGsfEWHBgcFaxP_aDLN_ZqpNb1FpSV7-OhWGUis,1797
+tunix/rl/rollout/base_rollout.py,sha256=UNGtkqqkvKZPluNpqUpFuuQFVlIQ43ecwdidatmklkY,3476
+tunix/rl/rollout/vanilla_rollout.py,sha256=6yoVRNMSYM1WHqN1xFgFtRVSMxnR4N6KGlqrvXt_PZ0,3589
+tunix/rl/rollout/vllm_rollout.py,sha256=I5rAzNz2Fp8TJEJfPw4qVWkMqJEqUagyMJlGVc-5LI4,3666
+tunix/sft/checkpoint_manager.py,sha256=iRvP2Wd4edCXL7R9VSRIfkequmh1q13g3SwXHMKT7rg,4354
+tunix/sft/config.py,sha256=VNrWt6CUkCk-SJaf3-393tyXlImOkWlBP9NumjwJGeQ,13477
+tunix/sft/hooks.py,sha256=WNkaJm4gRPB1KWz4PV8C0pn_lsXcmfeD0cSSKBbeh68,2232
+tunix/sft/inflight_throttler.py,sha256=IIJdtJRwsWaApxXbTVCFUq4GqG4RS57TgOfqVTN2sJs,2124
+tunix/sft/metrics_logger.py,sha256=Luo8IoW6ID6yH7ZTigpDI1IatCBQbwd8xZYDLwgEYgo,5845
+tunix/sft/peft_main.py,sha256=HEiB-iKraJ6TDAQK5lH24O8-R4nYsSWrkbO1fV-kAmc,13466
+tunix/sft/peft_trainer.py,sha256=n9ZpFYC11_d_XAi034-SYdfqMalzXOdwTZZuCKeNBNw,26447
+tunix/sft/profiler.py,sha256=nWZmf9j_sxqli6aERoRhiJUFLw2OhhFWEIkOWFEv5Kw,3329
+tunix/sft/progress_bar.py,sha256=3zkRWgo2peOc5s-54lvoCtKk3ULd4eQHAEERUdjGyuk,2465
+tunix/sft/sharding_utils.py,sha256=-n3Ba7eUMYFuqEhFAqrk42_ZujLkLeatk6ZBoSb0Zlk,1544
+tunix/sft/system_metrics_calculator.py,sha256=ziPoe1BgTdZCZ8DMDcw4sgu8vNfBBnnnVp-JIvf2J48,2506
+tunix/sft/utils.py,sha256=MtZCezCXUdxnJDuY52Sd2MZ-QWnoYoLKuCrhIA2-TH8,1178
+tunix/sft/dpo/dpo_trainer.py,sha256=ovXpjq7wjt353f3FCAuFxKZmGPYIZCbnTEfNv-fqhVM,15270
+tunix/tests/test_common.py,sha256=1jZkMdIR3y0bKTaUygj1eWxyqOubjumLZl9b3T-Q5dc,6133
+google_tunix-0.0.1.dist-info/METADATA,sha256=aRalgublhANE7VD2cMYR-F74svrzED_PvLSYW_UjLZk,6157
+google_tunix-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+google_tunix-0.0.1.dist-info/top_level.txt,sha256=VOWD8y-ahrChIDAFYazDhmKAfJUyKY7C0Ovgpky4Z8Y,6
+google_tunix-0.0.1.dist-info/RECORD,,

google_tunix-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

google_tunix-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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
+
+       http://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.

google_tunix-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+tunix

tunix/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2024 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
+#
+#      http://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.
+
+"""Tunix API."""
+
+# pylint: disable=g-multiple-import, g-importing-member
+
+from tunix.distillation.distillation_trainer import DistillationTrainer, TrainingConfig as DistillationTrainingConfig
+from tunix.generate.sampler import CacheConfig, Sampler
+from tunix.rl.grpo.grpo_learner import GrpoConfig, GrpoLearner, RewardFn
+from tunix.rl.rl_cluster import ClusterConfig, RLCluster, RLTrainingConfig, Role
+from tunix.rl.rollout.base_rollout import RolloutConfig
+from tunix.sft.dpo.dpo_trainer import DpoTrainer, DpoTrainingConfig
+from tunix.sft.metrics_logger import MetricsLogger, MetricsLoggerOptions
+from tunix.sft.peft_trainer import PeftTrainer, TrainingConfig
+
+# pylint: enable=g-multiple-import, g-importing-member

tunix/distillation/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2025 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
+#
+#     http://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.
+"""Module containing the distillation trainers."""
+
+# pylint: disable=g-importing-member
+
+from tunix.distillation.distillation_trainer import DistillationTrainer
+from tunix.distillation.distillation_trainer import TrainingConfig
+from tunix.distillation.distillation_trainer import TrainingInput

tunix/distillation/distillation_trainer.py

@@ -0,0 +1,136 @@
+# Copyright 2025 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
+#
+#     http://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.
+
+"""Distillation trainer."""
+import dataclasses
+from typing import Any, Callable, Tuple
+
+import flax
+from flax import nnx
+from jax.typing import ArrayLike  # pylint: disable=g-importing-member
+import optax
+from tunix.distillation import strategies
+from tunix.sft import metrics_logger
+from tunix.sft import peft_trainer
+from typing_extensions import override
+
+
+@dataclasses.dataclass(slots=True, kw_only=True)
+class TrainingConfig(peft_trainer.TrainingConfig):
+  """Distillation training config."""
+
+
+@flax.struct.dataclass(frozen=True)
+class TrainingInput(peft_trainer.TrainingInput):
+  """Distillation training input."""
+
+  teacher_output: Any = None
+
+
+class DistillationTrainer(peft_trainer.PeftTrainer):
+  """Distillation trainer."""
+
+  def __init__(
+      self,
+      student_model: nnx.Module,
+      teacher_model: nnx.Module,
+      strategy: strategies.BaseStrategy,
+      optimizer: optax.GradientTransformation,
+      training_config: TrainingConfig,
+  ):
+    """Initializes the DistillationTrainer.
+
+    Args:
+        student_model: The student model to train.
+        teacher_model: The teacher model to use for distillation.
+        strategy: The distillation strategy to use.
+        optimizer: The optimizer to use for training.
+        training_config: The training config.
+    """
+    student_model, teacher_model = strategy.pre_process_models(
+        student_model, teacher_model
+    )
+    super().__init__(student_model, optimizer, training_config)
+    self.strategy = strategy
+    self.teacher_model = teacher_model
+    self.loss_fn = self.get_train_loss
+    self.eval_loss_fn = self.get_eval_loss
+    self.gen_model_input_fn = lambda x: {
+        "inputs": {"input_tokens": x.input_tokens, "input_mask": x.input_mask},
+        "teacher_output": (
+            x.teacher_output if hasattr(x, "teacher_output") else None
+        ),
+    }
+
+  @override
+  def with_gen_model_input_fn(
+      self, gen_model_input_fn: Callable[[Any], dict[str, ArrayLike]]
+  ) -> "DistillationTrainer":
+    self.gen_model_input_fn = lambda x: {
+        "inputs": gen_model_input_fn(x),
+        "teacher_output": (
+            x.teacher_output if hasattr(x, "teacher_output") else None
+        ),
+    }
+    return self
+
+  @override
+  def with_loss_fn(
+      self,
+      loss_fn: Callable[..., ArrayLike | Tuple[ArrayLike, Any]],
+      has_aux: bool = False,
+  ) -> "DistillationTrainer":
+    raise NotImplementedError(
+        "with_loss_fn is not supported for distillation. Use the strategy to"
+        " define the loss."
+    )
+
+  @override
+  def _prepare_inputs(self, input_data: TrainingInput) -> TrainingInput:
+    inputs = self.gen_model_input_fn(input_data)["inputs"]
+    if self._mode == metrics_logger.Mode.EVAL:
+      teacher_output = None
+    else:
+      teacher_output = self.strategy.get_teacher_outputs(
+          self.teacher_model, inputs
+      )
+
+    return TrainingInput(
+        input_tokens=input_data.input_tokens,
+        input_mask=input_data.input_mask,
+        teacher_output=teacher_output,
+    )
+
+  def get_train_loss(
+      self,
+      model: nnx.Module,
+      teacher_output: Any,
+      inputs: dict[str, ArrayLike],
+  ) -> ArrayLike:
+    return self.strategy.get_train_loss(model, teacher_output, inputs)
+
+  def get_eval_loss(
+      self,
+      model: nnx.Module,
+      teacher_output: Any,
+      inputs: dict[str, ArrayLike],
+  ) -> ArrayLike:
+    del teacher_output  # Not computed in eval.
+    return self.strategy.get_eval_loss(model, inputs)
+
+  def close(self):
+    super().close()
+    self.model, self.teacher_model = self.strategy.post_process_models(
+        self.model, self.teacher_model
+    )

tunix/distillation/feature_extraction/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2025 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
+#
+#     http://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.
+"""Feature extraction utilities for distillation."""
+
+# pylint: disable=g-importing-member
+
+from tunix.distillation.feature_extraction.pooling import avg_pool_array_to_target_shape
+from tunix.distillation.feature_extraction.projection import ModelWithFeatureProjection
+from tunix.distillation.feature_extraction.projection import remove_feature_projection_from_models
+from tunix.distillation.feature_extraction.projection import setup_models_with_feature_projection
+from tunix.distillation.feature_extraction.sowed_module import pop_sowed_intermediate_outputs
+from tunix.distillation.feature_extraction.sowed_module import SowedModule
+from tunix.distillation.feature_extraction.sowed_module import unwrap_sowed_modules
+from tunix.distillation.feature_extraction.sowed_module import wrap_model_with_sowed_modules