我为什么无法在本地训练 Gemma 3n（以及为什么我转而使用 Vertex AI）
我一直在开发一个发音辅导应用，它需要一个能够同时理解文本和音频的多模态AI。该应用目前使用Google Gemini的原生音频模型，但我希望针对发音辅导微调一个专用模型。Google的Gemma 3n似乎非常合适，因为它通过通用语音模型（Universal Speech Model）编码器支持原生音频。我有一台配备48GB内存的MacBook Pro M4 Pro。这肯定足以在本地训练一个50亿参数的模型了吧？
剧透：事实并非如此。以下是我从失败中学到的经验。
我为什么尝试本地训练
其吸引力是显而易见的：
- 零成本：在RunPod或Vast.ai等服务上，云GPU训练每次尝试的费用为2-5美元。
- 隐私性：我的训练数据保留在本地。
- 快速迭代：无需上传数据集或配置云环境。
- 动手学习：我想试验PyTorch训练脚本，了解内存优化技术，并通过反复试验来学习。在本地，我可以在几秒钟内修改代码并重试，而不是部署到云端并等待脚本运行。
我有988个训练样本，涵盖发音错误、口语表达分析和辅导对话。这似乎是进行隔夜本地训练的完美候选。
训练模型到底意味着什么？
快速澄清一下：当人们说“训练”时，他们可能指的是两种截然不同的事情。
从头开始训练是Google和OpenAI所做的工作。你从随机数字开始，教神经网络从零开始学习一切。这需要数十亿的文本样本、数百个GPU和数月的计算时间。这既昂贵又缓慢。
微调是采用一个已经训练好的模型，并教它执行你的特定任务。困难的工作已经完成。模型已经了解语言和推理。你只是在向它展示你的特定格式和领域。这就像雇佣一个已经知道如何编程的人，然后教他你的代码库，而不是从头教一个人编程一样。
对于我的用例，我需要的是微调。模型已经理解了英语和发音。我只需要它以我的格式提供反馈。
如何使用Hugging Face进行微调
Hugging Face使整个过程比预期的要简单得多。他们托管了数千个预训练模型，并提供了用于微调它们的Python库。你不需要博士学位就能做到这一点。
基本工作流程如下：
from transformers import AutoProcessor, AutoModelForImageTextToText
# 下载模型和处理器
processor = AutoProcessor.from_pretrained("google/gemma-3n-E2B-it")
model = AutoModelForImageTextToText.from_pretrained("google/gemma-3n-E2B-it")
两行代码。这会将一个54亿参数的AI模型下载到你的笔记本电脑上。processor将你的数据转换为正确的格式，而model就是神经网络本身。
然后你使用Trainer类来实际运行训练：
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="outputs/gemma3n",
num_train_epochs=3,
per_device_train_batch_size=1,
learning_rate=3e-5,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=my_dataset,
)
trainer.train() # 开始训练！
这很棒，因为你跳过了所有底层工作，比如编写梯度下降循环和管理检查点。缺点是什么？你不太清楚幕后发生了什么，尤其是每样东西占用了多少内存。这正是我遇到的问题。
内存问题
训练大型语言模型需要在内存中同时容纳几样东西：
- 模型权重：对于Gemma 3n E2B（54亿参数），以float16格式约为11GB。
- 梯度：另外约11GB（与权重大小相同）。
- 优化器状态：Adam为每个参数保留2个状态，增加了约22GB。
- 激活值：中间计算结果。
- 批次数据：正在处理的训练样本。
总计：较小的E2B模型需要50-60GB，E4B模型需要70-80GB。我的48GB内存根本不够用。
我的尝试
尝试1：使用E4B进行完整微调
在优化器初始化期间出现内存不足（Out of Memory, OOM）。模型加载成功（16GB），但优化器无法分配其状态。
尝试2：更小的批次大小
将batch_size从2减小到1只是将OOM错误推迟到了反向传播阶段。
尝试3：使用E4B进行LoRA
LoRA（低秩适应）只训练小的适配器层，而不是整个模型，将可训练参数从78亿减少到4000万（0.5%）。但冻结的基础模型仍然需要装入内存：仅E4B就需要59GB。
尝试4：使用E2B进行LoRA
这次几乎成功了。使用LoRA的较小E2B模型占用了40-45GB，留下了3-8GB的余量。训练开始并运行了几个小时，然后在处理具有较长序列的特定批次时遇到内存错误。
最终失败的原因
即使采取了积极的优化措施（LoRA、batch_size=1、无梯度检查点），48GB对于一致的训练来说仍然不够。内存使用量因样本而异：
- 短样本：38-42GB（正常）
- 长样本：46-50GB（崩溃）
我或许可以通过过滤掉较长的样本或将max_length减小到1024个token来使其工作，但到那时，我就是在为了适应硬件限制而牺牲模型的性能。
转向Vertex AI
经过三天的失败尝试后，我将转向Google Cloud的Vertex AI进行训练：
- A100 40GB GPU：对于使用LoRA的E4B模型来说，内存绰绰有余。
- 托管基础设施：无需调试内存问题，不会崩溃。
- 成本：一次完整的训练运行（3个epoch）约3-5美元。
- 时间：2-3小时，而我在本地尝试了6-9小时。
我将撰写一篇单独的博客文章介绍Vertex AI的训练过程，内容包括：
- 使用Gemma 3n设置自定义训练作业。
- 将数据集上传到Cloud Storage。
- 使用TensorBoard监控训练。
- 将训练好的模型部署到Vertex AI端点。
什么时候可以进行本地训练？
在Apple Silicon上进行本地训练是可行的，前提是：
- 你的模型更小：使用LoRA的1-3B参数模型可以在48GB内存中舒适地运行。
- 你有更多内存：配备128GB内存的M4 Max可以轻松处理Gemma 3n E2B。
- 你使用量化：4位或8位量化可以进一步减少内存占用，但会增加复杂性。
关键经验教训
- 内存是瓶颈：48GB的统一内存听起来很多，直到你尝试训练50亿参数以上的模型。仅优化器状态就可能使你的内存需求增加一倍或三倍。
- LoRA有帮助但并非万能：它将可训练参数减少了99%，但在前向和反向传播过程中，冻结的基础模型权重仍然需要装入内存。
- 云训练具有成本效益：与花费数天时间调试OOM错误和失败的实验相比，一次训练运行花费3-5美元是合理的。
- 了解你的硬件限制：Apple Silicon的统一内存非常适合推理和较小模型的训练（1-3B参数），但要可靠地训练5B以上的模型，需要专用的GPU VRAM。