如何在GPU或TPU上用JAX对LLaMA、OpenLLaMA和XGen进行微调

LLaMA、OpenLLaM和XGen

LLaMA模型是由Meta公司在2023年2月发布的。这个生成性人工智能模型是一个开源的模型，提出了几种规模：7B参数、13B参数、33B参数和65B参数。

2021年6月，当GPT-J发布时，世界开始意识到开源生成式人工智能模型可以与OpenAI GPT-3认真竞争。现在有了LLaMA，标准显然又被提高了，这个模型似乎是OpenAI ChatGPT和GPT-4的一个非常好的开源替代品。

虽然LLaMA的许可证对商业不友好：这个模型不能用于商业目的......但好消息是，现在还有其他模式。

OpenLLaMA于2023年6月发布，是LLaMA的一个替代版本，由伯克利人工智能研究团队开发，给出了非常好的结果，可用于商业。截至目前，有2个版本可供选择：7B参数和13B参数。

XGen，由Salesforce在2023年6月发布，是另一个非常强大的基础模型，可以在商业应用中使用。截至目前，只有7B参数版本可用。值得注意的是，这个模型支持8k tokens上下文，而LLaMA和OpenLLaMA只支持2k tokens内容。

为什么要对你自己的模型进行微调？

上述模型是基础性模型，这意味着它们是以无监督的方式在大型文本语料库中训练出来的。

这些基础性的人工智能模型通常是一个很好的基础，但它们需要进行调整，以正确理解你想要的东西并返回良好的结果。实现这一目标的最简单方法是使用少量学习（也被称为 "提示工程"）。 欢迎在此阅读我们专门的几张照片学习指南。

少量的学习是很方便的，因为它可以在飞行中进行，而不需要创建一个新版本的生成性人工智能模型，但有时是不够的。

为了获得最先进的结果，你会想为自己的使用案例微调人工智能模型。微调意味着你将根据你自己的数据修改模型中的一些参数，然后得到你自己版本的模型。

微调比从头开始训练一个生成性人工智能模型要便宜得多，但它仍然需要计算能力，所以你需要先进的硬件来微调你自己的模型。最近的一些替代性微调技术需要更少的计算能力（见p-tuning、提示性微调、软微调、参数高效微调、适配器、LoRA、QLoRA......），但到目前为止，我们还没有设法用这些技术获得同样的质量水平，所以我们不打算在本教程中提及它们。

用JAX和EasyLM对TPU上的LLaMA进行微调

在本教程中，我们重点介绍了用伯克利人工智能研究团队发布的EasyLM库对LLaMA进行微调： https://github.com/young-geng/EasyLM. This library is based on JAX which makes the fine-tuning process fast and compatible with both GPUs and Google TPUs.

你也可以用同样的技术对OpenLLaMA或XGen进行微调。

在这里，我们在Google TPU V3-8上对LLaMA 7B进行了微调，但你完全可以在A100 GPU上做同样的事情（只需仔细阅读EasyLM文档中的 "安装 "部分，该部分略有不同）。当然，你也可以对更大版本的LLaMA进行微调（13B、33B、65B......），但你需要的东西远远超过TPU V3-8或单个A100 GPU。

我们来了!

首先，为你的用例创建一个文本生成数据集，以JSONL格式，使用 "text "作为每个例子的关键。这里是一个简单的情感分析数据集：

{"text":"[Content]: I love NLP Cloud, this company is awesome!\n[Sentiment]: Positive"}
{"text":"[Content]: Training LLMs is a complex but rewarding process.\n[Sentiment]: Neutral"}
{"text":"[Content]: My fine-tuning keeps crashing because of an OOM error! It just does not work at all!\n[Sentiment]: Negative"}

请注意几件重要的事情。首先，为了简单起见，这个数据集只包含3个例子，但在现实生活中，你将需要更多的例子。300个例子通常是一个好的开始。其次，当你将使用你的微调模型进行推理时，你将需要严格遵循相同的格式，使用"[内容]: "和"[情绪]: "的前缀。最后，"</s> "标记很重要，因为它意味着模型应该在这里停止生成。你可以在NLP Cloud文档中找到更多的数据集例子： 在此了解更多。

在谷歌云上用V2 Alpha软件版本创建一个TPU V3-8虚拟机：

SSH进入虚拟机并安装EasyLM：

git clone https://github.com/young-geng/EasyLM
cd EasyLM
bash ./scripts/tpu_vm_setup.sh

你现在可以下载并转换LLaMA的权重。第一个选择是向Meta索取官方权重： https://ai.facebook.com/blog/large-language-model-llama-meta-ai/. 然后用这个脚本将权重转换为EasyLM： https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_torch_to_easylm.py. 第二个选择是使用HuggingFace上的LLaMA权重： https://huggingface.co/decapoda-research/llama-7b-hf. 然后用这个脚本将权重转换为EasyLM： https://github.com/young-geng/EasyLM/blob/main/EasyLM/models/llama/convert_hf_to_easylm.py.

将你的数据集上传到虚拟机，使用HF LLaMA标记器计算它包含多少个标记：

pip install -U transformers
python -c "from transformers import LlamaTokenizer; tokenizer = LlamaTokenizer.from_pretrained('decapoda-research/llama-7b-hf'); f = open('/path/to/your/dataset', 'r'); print(len(tokenizer.encode(f.read())))"

如果你为1024个标记的上下文训练你的模型，你将需要把返回的标记数除以1024。

如果你为2048个令牌的上下文训练你的模型，你将需要把返回的令牌数除以2048。

这个数字将是每个epoch的步骤数。例如，如果你想训练5个epochs（这通常是一个很好的设置），你将需要用这个数字乘以5，然后把得到的值放在下面的--total_steps中。

这里有一个具体的例子：如果你的数据集包含100,000个标记，而你想要一个1024个标记的上下文和5个历时，你的总步数将是（100,000/1024）*5=488。

根据你的上下文长度，将--train_dataset.json_dataset.seq_length设置为1024或2048。请注意，为2048 tokens语境微调模型需要更多的内存，所以如果不是严格意义上的需要，我们建议你坚持使用1024 tokens语境。

现在你可以启动微调程序：

nohup python -u EasyLM/EasyLM/models/llama/llama_train.py \
--total_steps=your number of steps \
--save_model_freq=usually same as your number of steps \
--optimizer.adamw_optimizer.lr_warmup_steps=usually 10% of total steps \
--train_dataset.json_dataset.path='/path/to/your/dataset' \
--train_dataset.json_dataset.seq_length=1024 or 2048 \
--load_checkpoint='params::/path/to/converted/model' \
--tokenizer.vocab_file='/path/to/tokenizer' \
--logger.output_dir=/path/to/output  \
--mesh_dim='1,4,2' \
--load_llama_config='7b' \
--train_dataset.type='json' \
--train_dataset.text_processor.fields='text' \
--optimizer.type='adamw' \
--optimizer.accumulate_gradient_steps=1 \
--optimizer.adamw_optimizer.lr=0.002 \
--optimizer.adamw_optimizer.end_lr=0.002 \
--optimizer.adamw_optimizer.lr_decay_steps=100000000 \
--optimizer.adamw_optimizer.weight_decay=0.001 \
--optimizer.adamw_optimizer.multiply_by_parameter_scale=True \
--optimizer.adamw_optimizer.bf16_momentum=True &

一些解释：

--save_model_freq：你想在过程中保存模型的频率。如果你只对一个小数据集进行微调，你可以只在过程结束时保存，在这种情况下，这个值将等于--总步数。

--optimizer.adamw_optimizer.lr_warmup_steps：总步数的10%通常是一个好值。

--tokenizer.vocab_file：通往tokenizer.model文件的路径。例如，如果你使用HuggingFace上的decapoda存储库，这里是标记器的链接：https://huggingface.co/decapoda-research/llama-7b-hf/resolve/main/tokenizer.model

--logger.output_dir: 最终模型和日志的路径

其他参数可以不动。

一旦微调过程结束，你可以在你在--logger.output_dir中指定的路径下检索你的模型。

使用微调模型进行推理

你现在有了自己的微调模型，你当然想使用它了

第一个策略是使用EasyLM库进行推理。在这种情况下，你可以像这样启动推理服务器：

python EasyLM/EasyLM/models/llama/llama_serve.py  \
--mesh_dim='1,1,-1' \
--load_llama_config='7b' \
--load_checkpoint='params::/path/to/your/model' \
--tokenizer.vocab_file='/path/to/tokenizer'

然后简单地用cURL发送你的请求，像这样：

curl "http://0.0.0.0:5007/generate" \
-H "Content-Type: application/json" \
-X POST -d '{"prefix_text":["[Content]: EasyLM works really well!\n[Sentiment]:"]}'

第二个策略是将你的模型导出为HuggingFace格式，以便用另一个框架进行推理。下面是你如何导出它：

python EasyLM/EasyLM/models/llama/convert_easylm_to_hf.py \
--load_checkpoint='params::/path/to/output/model/streaming_params' \
--tokenizer_path='/path/to/tokenizer' \
--model_size='7b' \
--output_dir='/path/to/converted/model'

总结

2023年对于开源生成性人工智能模型来说是一个伟大的里程碑。从现在开始，每个人都可以使用伟大的模型，如LLaMA、OpenLLaMA、XGen、Orca、Falcon...

对这些模型进行微调是获得尖端结果的最佳方式，为您自己的使用案例量身定做，可以大大超过最好的专有AI模型，如ChatGPT（GPT-3.5）、GPT-4、Claude...

在本指南中，我向你展示了如何微调LLaMA、OpenLLaMA和XGen。 如果你有问题，请不要犹豫地与我联系，如果你想轻松地微调和部署先进的生成性人工智能模型，而没有任何技术上的复杂性、 have a look at the NLP Cloud dedicated documentation!

Mark
NLP云的机器学习工程师