Colab+Drive微调大模型

假如说，你最近发现了 Hugging Face 和大量开源模型，如 BERT、Llama、BART 以及 Mistral AI、Facebook、Salesforce 和其他公司的大量生成语言模型。现在，你想尝试微调一些大型语言模型，用于你的业余项目。事情开始很顺利，但后来你发现它们在计算上非常耗时，而你手边没有 GPU 处理器。

Google Colab 慷慨地为你提供了一种访问免费计算的方法，以便你可以解决这个问题。缺点是，你需要在基于浏览器的临时环境中完成所有操作。更糟糕的是，整个过程都有时间限制，所以似乎无论您做什么，当内核最终关闭并且环境被摧毁时，你都会失去宝贵的微调模型和所有结果。

不要害怕。有一种方法可以解决这个问题：使用 Google Drive 保存任何中间结果或模型参数。这将允许你在稍后阶段继续实验，或者在其他地方使用经过训练的模型进行推理。

为此，你需要一个 Google 帐户，该帐户具有足够的 Google Drive 空间来存储您的训练数据和模型检查点。我假设你已经在 Google Drive 中创建了一个名为 data 的文件夹，其中包含你的数据集。然后创建了另一个名为 checkpoints 的文件夹，它是空的。

然后在你的 Google Colab Notebook 中使用以下命令安装你的 Drive：

from google.colab import drive
drive.mount('/content/drive')  

现在，你可以在新单元格中使用以下两个命令列出数据和检查点目录的内容：

!ls /content/drive/MyDrive/data
!ls /content/drive/MyDrive/checkpoint

如果这些命令有效，那么你现在可以访问笔记本中的这些目录。如果命令不起作用，那么你可能错过了授权步骤。上面的 drive.mount 命令应该会生成一个弹出窗口，要求你点击并授权访问。你可能错过了弹出窗口，或者没有选择所有必需的访问权限。尝试重新运行单元格并检查。

获得访问权限后，你就可以编写脚本，以便将模型和结果序列化到 Google Drive 目录中，从而在会话中保留下来。理想情况下，你可以编写训练作业，以便任何运行时间过长的脚本都可以从上一个会话中加载部分训练的模型，并从该点继续训练。

实现这一点的一个简单方法是创建一个供训练脚本使用的保存和加载函数。在初始化新模型之前，训练过程应始终检查是否存在部分训练的模型。以下是一个示例保存函数：

def save_checkpoint(epoch, model, optimizer, scheduler, loss, model_name, overwrite=True):
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'loss': loss
    }
    direc = get_checkpoint_dir(model_name)
    if overwrite:
        file_path = direc + '/checkpoint.pth'
    else:
        file_path = direc + '/epoch_'+str(epoch) + '_checkpoint.pth'
    if not os.path.isdir(direc):
       try:
          os.mkdir(direc)
       except:
          print("Error: directory does not exist and cannot be created")
          file_path = direc +'_epoch_'+str(epoch) + '_checkpoint.pth'
    torch.save(checkpoint, file_path)
    print(f"Checkpoint saved at epoch {epoch}")

在此实例中，我们将模型状态以及一些元数据（epoch 和损失）保存在字典结构中。我们包含一个选项，用于覆盖单个检查点文件，或为每个 epoch 创建一个新文件。我们使用 torch.save 函数，但原则上你可以使用其他序列化方法。关键思想是你的程序打开文件并确定现有文件使用了多少个训练 epoch。这允许程序决定是继续训练还是继续前进。

同样，在加载函数中，我们传入对我们希望使用的模型的引用。如果已经有一个序列化模型，我们将参数加载到我们的模型中并返回它训练的 epoch 数。这个 epoch 值将决定需要多少个额外的 epoch。如果没有模型，那么我们会得到零个 epoch 的默认值，并且我们知道模型仍然具有初始化时的参数。

def load_checkpoint(model_name, model, optimizer, scheduler):
    direc = get_checkpoint_dir(model_name)
    if os.path.exists(direc):
        file_path = get_path_with_max_epochs(direc)
        checkpoint = torch.load(file_path, map_location=torch.device('cpu'))
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        epoch = checkpoint['epoch']
        loss = checkpoint['loss']
        print(f"Checkpoint loaded from {epoch} epoch")
        return epoch, loss
    else:
        print(f"No checkpoint found, starting from epoch 1.")
        return 0, None

这两个函数需要在训练循环中调用，并且需要确保使用 epochs 值的返回值来更新训练迭代中的 epochs 值。结果是，你现在有一个可以在内核死亡时重新启动的训练过程，并且它将从中断的地方继续进行。

该核心训练循环可能如下所示：

EPOCHS = 10
for exp in experiments: 
    model, optimizer, scheduler = initialise_model_components(exp)
    train_loader, val_loader = generate_data_loaders(exp)
    start_epoch, prev_loss = load_checkpoint(exp, model, optimizer, scheduler)
    for epoch in range(start_epoch, EPOCHS):
        print(f'Epoch {epoch + 1}/{EPOCHS}')
        # ALL YOUR TRAINING CODE HERE
        save_checkpoint(epoch + 1, model, optimizer, scheduler, train_loss, exp)

注意：在此示例中，我正在尝试训练多个不同的模型设置（在名为实验的列表中），可能使用不同的训练数据集。支持函数 initialise_model_components 和 generate_data_loaders 负责确保我为每个实验获得正确的模型和数据。

上述核心训练循环允许我们重复使用训练和序列化这些模型的整体代码结构，确保每个模型都达到所需的训练次数。如果我们重新启动该过程，它将再次遍历实验列表，但它将放弃任何已经达到最大次数的实验。

原文链接：Training Language Models on Google Colab

汇智网翻译整理，转载请标明出处