在其他任何事情之前,您可能需要申请访问TPU研究云(TRC)。再加上谷歌云免费试用,你可以免费在这里做任何事情。一旦你进入TRC,你需要创建一个项目,然后用新项目的名称填写发送给你的表格。使用脚本create_finetune_tfrecords.py将数据准备为tfrecords;我可能会做一个单独的指南。您可能需要做的另一件事是fork mesh-transformer-jax代码,以便更容易添加和修改配置文件。
- 安装Google Cloud SDK,我们稍后需要它。
- 如果你还没有完成一个项目并通过TRC激活TPU访问(或者如果你计划自掏腰包),免费申请。
- TPU使用谷歌云存储桶进行存储,现在就创建一个吧。确保它位于TPU VM所在的区域;TRC的电子邮件将告诉您可以在哪个地区使用免费TPU。
- 为了微调模型,您需要完全预训练的权重。下载地址
现在云上有了一个存储桶,电脑上有了权重,您需要分两步将权重上传到存储桶:
- 解压缩并提取GPT-J-6B/step_383500.tar.zstd,这样您就剩下包含碎片检查点的未压缩文件夹。
- 打开Google Cloud SDK并运行以下命令,根据需要替换路径名:gsutil-m cp -R LOCAL_path_TO/step_383500 gs://YOUR-BUCKET。如果成功,控制台将显示正在上载的文件。注:我从加州上传到荷兰,花了大约12个小时;希望你的地理位置比我好!我最初也犯了上传仍然打包的.tar的错误。不要这样做,TPU虚拟机没有足够的本地存储空间供您解压缩。为了避免重新上传,我不得不在Colab中解压缩。
你也会想上传你的数据的tfrecords,说实在的,没有人会想通过网络界面上传近70GB的权重。
请注意,稍后可以通过在VM的文本编辑器中编辑基本repo来完成准备索引和配置文件的步骤6和7。相反,对您自己的repo分支进行以下更改更有效:
- 在数据文件夹中,创建一个新的文件foo.train.index,将foo替换为您希望引用数据集的任何内容。对于存储桶中要训练的每个tfrecord,将路径添加为索引中的一行。创建foo.val.index,并对验证数据集执行同样的操作(如果有)。有关示例,请参见现有文件。
- 复制配置文件6B_roto_256.json,将其重命名为适合您项目的名称。打开它并进行以下编辑:
- tpu_size:从256更改为8
- bucket:换成你的bucket
- model_dir:更改到要保存checkpoints的目录
- train_set和val_set:从最后一步更改为索引文件
- eval_haness_tasks:如果不打算使用eval线束,则可以删除
- val_every&ckpt_every&keep_every:用法应该直观。但不要将foo_every值设置为0,否则会出现被零除的错误。如果没有val_set,只需将val_every设置为高于total_steps的值
- val_batches:这应该等于val数据集中的序列数。您可以在create_finetune_tfrecords.py生成的.tfsrecords文件的末尾找到此数字
- name:更改为模型的名称
- warmup_steps、lr、val_batches等:请参阅本指南末尾的学习率注释部分
- 将更改推送到GitHub存储库。
- 遵循本指南,直到连接到您的云TPU VM。
此时,您应该可以远程访问TPU VM!
- 在新的VM终端中,键入git clone https://github.com/kingoflolz/mesh-transformer-jax(或者,最好是在推送配置文件和索引文件之后,您自己的fork)。
- 使用cd mesh-transformer-jax移动到新目录,并运行pip install-r requirements.txt。由于requirements.txt文件没有固定微调所需的确切jax版本,因此运行pip安装jax==0.2.12,您就完成了。
- 最后,运行python3 device_train.py --config=YOUR_config.json --tune model path=gs://YOUR-BUCKET/step_383500/。如果一切设置正确,这将开始微调过程。首先,必须将模型加载到内存中;当加载控制台上显示的网络时,需要大约10-15分钟才能执行下一步,设置WandB进行日志记录。选项3允许您在不使用WandB时跳过该选项。步骤1检查点将保存下来,真正的训练将开始。如果您有一个小数据集,这将很快完成;TPU虚拟机可以每秒约5000个令牌进行训练。
- 完成了!现在不要忘记你需要采取的清理步骤,比如关闭TPU虚拟机或删除存储桶中不需要的数据,这样你以后就不会收到谷歌的任何意外收费。
更多
本指南被标记为“基础知识”,我们迄今为止没有涉及的任何内容都超出了范围,但请查看存储库的其余部分!尝试python3 device_sample.py --config=configs/YOUR_config.json获取基本采样接口。使用slim_model.py准备一个更容易部署的新权重的slim版本,以便进行推断、实验。
运行HuggingFace
要使用pytorch在HuggingFace的转换器库中使用模型,需要将权重转换为它可以识别的格式。这可以使用to_hf_weights.py完成。建议您在尝试将权重移动到pytorch/transform格式之前使用slim_model.py。使用python to_hf_weights.py --help查看用法详细信息。
- 注:截至2021 9月1日,GPT-J已并入变压器主分支,但尚未投产。运行pip安装git+https://github.com/huggingface/transformers#transformers安装当前主分支。
炼丹参数
首先要确定的是训练时期的长度。gradient_accumulation_steps是批次大小,默认为16,nostalgebraist建议32。tfrecord文件的文件名中应该有一个数字,指示数据集中有多少序列。将该数字除以批次大小,结果就是一个时期中有多少步。现在我们可以写时间表了。
lr建议介于1e-5和5e-5之间,endlr设置为lr的1/5或1/10。weight_decay可以保持0.1。totalsteps至少应该是一个历元,如果您有一个验证集来确定您的训练损失,则可能更长。warmup_steps应该是总数的5-10%,最后anneal_steps应该为totalsteps-warmup_steps。(在warmup_steps+anneal_steps之后,lr设置为end_lr,然后保持训练直到total_steps,但通常应该在退火完成后停止)
举个例子:我有一个小数据集,它标记成1147个序列作为.tfrecord。除以gradient_aaccumulation_steps设置为16,四舍五入以确保我使用所有数据,等于每个历元72步。我将lr设置为5e-5,end_lr设置为其中的五分之一,1e-5;这可能太多了,它处于推荐范围的高端。我将为一个epoch将total_steps设置为72,因为我没有验证集。然后我将anneal_steps设置为65,warm_steps设置设置为7。这么简单,但你可能需要自己处理细节。