Google开源Tunix:JAX生态的LLM微调方案来了

JAX生态这两年在LLM训练这块追赶得挺快。PyTorch虽然还是主流但JAX在并行计算、TPU加速和API组合性上确实有些独特的优势。Google今天放出了Tunix这个库,专门做LLM的后训练------微调、强化学习、知识蒸馏这些都能搞。

Tunix是什么

这是个构建在JAX之上的后训练库,和Flax NNX集成得比较紧密。主要解决三类问题:

  • 监督微调(Supervised Fine-Tuning)
  • 强化学习(Reinforcement Learning)
  • 知识蒸馏(Knowledge Distillation)

现在还在早期开发阶段,功能在持续迭代,支持的模型也在慢慢扩展。

核心功能

监督微调:既支持全参数微调,也支持LoRA和Q-LoRA这类参数高效的方法。内存和算力受限的时候,PEFT方案还是挺实用的。

强化学习:实现了几个主流算法:PPO(Proximal Policy Optimization)、GRPO(Group Relative Policy Optimization)、还有token级别的GSPO。另外还有DPO(Direct Preference Optimization)做偏好对齐,这个在RLHF场景用得比较多。

知识蒸馏:支持几种策略,包括基于logit的概率分布匹配、注意力机制的转移和投影、跨架构的特征池化与投影。这几种方法在不同场景下各有用处。

库的设计比较模块化,组件可以自由组合,想扩展自定义流程也不算麻烦。分布式训练支持数据并行(DP)、完全分片数据并行(FSDP)和张量并行(TP),对TPU做了专门优化。

安装

三种装法:

从PyPI装(推荐):

复制代码
 pip install "tunix[prod]"

或者直接从GitHub主分支:

复制代码
 pip install git+https://github.com/google/tunix

开发模式从源码装:

复制代码
 git clone https://github.com/google/tunix.git  
 cd tunix  
 pip install -e".[dev]"

TPU上用QLoRA微调Gemma

拿个英译法的任务来演示。用的是Google的Gemma 2B模型,跑在TPU v5e-8上。

环境准备

复制代码
 pip install -q kagglehub safetensors tensorflow tensorflow_datasets tensorboardX transformers grain datasets  
 pip install -q git+https://github.com/google/tunix  
 pip install -q git+https://github.com/google/qwix  
   
 # Flax需要升级到最新版
 pip uninstall -q -y flax  
 pip install -q git+https://github.com/google/flax.git

完整流程

第一步,从Kaggle拉预训练checkpoint:

复制代码
 import kagglehub  
   
 model_path = "google/gemma/flax/2b"  
 kaggle_ckpt_path = kagglehub.model_download(model_path)

初始化模型和tokenizer:

复制代码
 from flax import nnx  
from tunix.models.gemma import model as gemma_lib, params as params_lib  
from tunix.generate import tokenizer_adapter as tokenizer_lib  

base_model = gemma_lib.Transformer.from_params(  
    params_lib.load_and_format_params(kaggle_ckpt_path, "2b"),  
    version="2b"  
)  
 tokenizer = tokenizer_lib.Tokenizer(tokenizer_path=f"{kaggle_ckpt_path}/tokenizer.model")

挂上QLoRA adapter:

复制代码
 import qwix  

lora_provider = qwix.LoraProvider(  
    module_path=".*(q_einsum|kv_einsum|proj)",  
    rank=16,  
    alpha=2.0,  
    weight_qtype="nf4"  # enable QLoRA quantization
)  
 lora_model = qwix.apply_lora_to_model(base_model, lora_provider)

这里rank设成16,alpha是2.0,weight_qtype指定nf4量化格式。

加载训练数据:

复制代码
 from tunix.examples.data import translation_dataset  

train_ds, validation_ds = translation_dataset.create_datasets(  
    dataset_name="mtnt/en-fr",  
    global_batch_size=16,  
    max_target_length=256,  
    num_train_epochs=3,  
    tokenizer=tokenizer,  
 )

用的是mtnt的英法平行语料,batch size 16,目标序列最长256个token。

开始训练:

复制代码
 from tunix.sft import peft_trainer, utils  
import optax  

trainer=peft_trainer.PeftTrainer(  
    lora_model,  
    optimizer=optax.adamw(1e-3),  
    config=peft_trainer.TrainingConfig(max_steps=100)  
)  
 trainer.train(train_ds, validation_ds)

优化器用AdamW,学习率1e-3,跑100步看看效果。

推理测试:

训练完直接用adapter过的模型做生成。Tunix提供了Sampler工具:

复制代码
 from tunix.generate import sampler as sampler_lib  

# initialize sampler
sampler = sampler_lib.Sampler(  
    transformer=lora_model,  
    tokenizer=tokenizer,  
    cache_config=sampler_lib.CacheConfig(  
        cache_size=256,  
        num_layers=base_model.num_layers,  
        num_kv_heads=base_model.num_kv_heads,  
        head_dim=base_model.head_dim,  
    ),  
)  

# test prompts
input_batch = [  
    "Translate this into French:\nHello, my name is Morgane.\n",  
    "Translate this into French:\nThis dish is delicious!\n",  
    "Translate this into French:\nI am a student.\n",  
    "Translate this into French:\nHow's the weather today?\n",  
]  

# generate predictions
out_data = sampler(  
    input_strings=input_batch,  
    max_generation_steps=20,  
)  

# print results
for input_string, out_string in zip(input_batch, out_data.text):  
    print(f"----------------------")  
    print(f"Prompt:\n{input_string}")  
     print(f"Output:\n{out_string}")

如果用的是QLoRA,把lora_model换成qlora_model就行。生产环境可以考虑把adapter合并回基模型,推理延迟能降下来。

总结

100步训练之后,模型已经能生成一些翻译结果了,虽然质量还不够好。多训练一段时间,准确率会明显提升,而且内存开销和训练速度都保持在不错的水平。

Tunix现在还比较新,但已经能看出一些潜力。TPU优先的设计、模块化的API、LoRA/QLoRA支持、完整的分布式训练策略,这些对做LLM适配研究的人来说都挺有用。

后续应该会继续扩展支持的模型类型和训练算法,值得关注。

地址:https://avoid.overfit.cn/post/c434311d8a894922b6c52ea179cf8d97

作者:Abish Pius

相关推荐
运维行者_1 小时前
AWS云服务故障复盘——从故障中汲取的 IT 运维经验
大数据·linux·运维·服务器·人工智能·云计算·aws
Saniffer_SH1 小时前
搭载高性能GPU的英伟达Nvidia DGX Spark桌面性能小怪兽国内首台开箱视频!
人工智能·深度学习·神经网络·ubuntu·机器学习·语言模型·边缘计算
数字化脑洞实验室1 小时前
AI决策vs人工决策:效率的底层逻辑与选择边界
人工智能
可触的未来,发芽的智生2 小时前
追根索源:换不同的词嵌入(词向量生成方式不同,但词与词关系接近),会出现什么结果?
javascript·人工智能·python·神经网络·自然语言处理
递归不收敛2 小时前
三、检索增强生成(RAG)技术体系
人工智能·笔记·自然语言处理
特立独行的猫a2 小时前
ESP32使用笔记(基于ESP-IDF):小智AI的ESP32项目架构与启动流程全面解析
人工智能·架构·esp32·小智ai
Tiandaren2 小时前
自用提示词01 || Prompt Engineering || 学习路线大纲 || 作用:通过启发式的问题来带动学习
人工智能·pytorch·深度学习·nlp·prompt·1024程序员节
IT_陈寒2 小时前
React 19重磅前瞻:10个性能优化技巧让你少写30%的useEffect代码
前端·人工智能·后端
国科安芯2 小时前
AS32S601ZIT2抗辐照MCU在商业卫星飞轮系统中的可靠性分析
服务器·网络·人工智能·单片机·嵌入式硬件·fpga开发·1024程序员节