单机多卡分布式训练策略——MirroredStrategy

前言

分布式训练是一种用于在多个设备或机器上同时训练深度学习模型的技术，它有助于减少训练时间，允许使用更多数据更快训练大模型。分布式训练重点关注数据并行性，本次试验使用的是单机多卡的分布式训练策略，也就是 MirroredStrategy 。通常单台机器上的有 1-8 个 GPU ， 这也是研究人员和小公司最常见的配置。

MirroredStrategy 简介

tf.distribute.MirroredStrategy 的步骤如下：

• 训练开始前，该策略在 N 个 GPU 上各复制一份完整的模型备份；
• 每次训练传入一个 batch 的数据，将数据分成 N 份，分别传入 N 个计算设备（即数据并行）；
• N 个计算设备在自己的内存中分别计算自己所获得数据的模型梯度；
• 使用分布式计算的 All-reduce 操作，在所有 GPU 间高效交换梯度数据并进行求和，使得最终每个设备都有了所有设备的梯度之和；
• 使用梯度求和的结果更新各个 GPU 中的模型权重；
• 因为该策略是同步的，所以只有当所有设备均更新模型后，才进入下一轮训练。

虚拟出 4 个 2G 的GPU

1. 这里先查找主机系统中可用的物理 GPU 设备，因为我这里只有一块 4090 ，所以结果肯定是包含只有一个物理 GPU 的列表，并将它们存储在 physical_devices 列表中。
2. 将我们唯一的物理 GPU 设备 physical_devices[0] 划分成了四个虚拟 GPU 设备，每个虚拟 GPU 的内存限制被设置为 2048MB 。这样是为了模拟一个单机多卡的分布式环境，方便我们试验 MirroredStrategy 策略。

physical_devices = tf.config.list_physical_devices('GPU')
tf.config.set_logical_device_configuration(
    physical_devices[0],
    [
        tf.config.LogicalDeviceConfiguration(memory_limit=2048),
        tf.config.LogicalDeviceConfiguration(memory_limit=2048),
        tf.config.LogicalDeviceConfiguration(memory_limit=2048),
        tf.config.LogicalDeviceConfiguration(memory_limit=2048),
    ]
)

logical_devices = tf.config.list_logical_devices('GPU')
print(f'从一张物理 GPU 中虚拟出 {len(logical_devices)} 个逻辑 GPU')

结果打印：

从一张物理 GPU 中虚拟出 4 个逻辑 GPU

数据准备

这里主要是准备用于训练神经网络的文本数据集，并对数据进行一些预处理，具体如下：

1. 设置每个训练批次的大小为 128，训练的总轮数为 5 。 从指定的 URL 下载 wiki 文本数据集，并将其解压缩到本地。
2. train_ds、val_ds 和 test_ds：这三个变量分别用于表示训练、验证和测试数据集。对每个数据集进行了类似的处理步骤：
   • 通过过滤器函数，剔除长度小于 100 个字符的文本行，以排除短文本。
   • 对数据进行随机洗牌，以打乱样本的顺序，有助于模型的训练。
   • 将数据批次大小设置为 BATCH_SIZE 。
   • 将数据缓存，以提高数据加载的效率。
   • 使用预取策略，允许在模型训练时异步加载下一个批次的数据，以减少训练时的等待取数据的时间。

BATCH_SIZE = 128
EPOCHS = 5
keras.utils.get_file(origin="https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip", extract=True, )
wiki_dir = os.path.expanduser("~/.keras/datasets/wikitext-2/")
train_ds = (tf.data.TextLineDataset(wiki_dir + 'wiki.train.tokens').filter(lambda x: tf.strings.length(x) > 100).shuffle(buffer_size=500).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE))
val_ds = (tf.data.TextLineDataset(wiki_dir + 'wiki.valid.tokens').filter(lambda x: tf.strings.length(x) > 100).shuffle(buffer_size=500).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE))
test_ds = (tf.data.TextLineDataset(wiki_dir + 'wiki.test.tokens').filter(lambda x: tf.strings.length(x) > 100).shuffle(buffer_size=500).batch(BATCH_SIZE).cache().prefetch(tf.data.AUTOTUNE))

分布式训练

这里介绍训练一个 BERT 掩码语言模型（Masked Language Model，MLM），并使用分布式训练策略 tf.distribute.MirroredStrategy 进行训练。具体如下：

1. TRAINING_STEP 定义了记录了每个 epoch 总的训练步骤。lr_schedule定义了一个学习率衰减策略，学习率在训练过程中从初始值（0.0001）线性地衰减到结束值（0.0），衰减的步数由 TRAINING_STEP 决定，也就是每个 epoch 过后学习率进行衰减下调。
2. callbacks：定义一些回调函数，用于在训练过程中执行特定的操作。包括了早停（EarlyStopping）和记录训练日志（TensorBoard）。
3. 创建一个 MirroredStrategy ，用于多 GPU 分布式训练。 在 MirroredStrategy 的作用域内定义模型和训练过程。这意味着模型同时将在多个 GPU 上进行训练。在作用域中创建一个 BERT 掩码语言模型，其中包括预训练的 BERT 模型，并将其最后一层池化层设置为不可训练。
4. 编译模型，定义了损失函数为稀疏分类交叉熵、优化器为 Adam 和评估指标为稀疏分类准确率。

TRAINING_STEP = sum(1 for _ in train_ds.as_numpy_iterator()) * EPOCHS
lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(initial_learning_rate=0.0001, decay_steps=TRAINING_STEP, end_learning_rate=0.)
callbacks = [
    tf.keras.callbacks.EarlyStopping(patience=2),
    tf.keras.callbacks.TensorBoard('./logs')]

strategy = tf.distribute.MirroredStrategy()
print(f"可用的分布式训练的 GPU 设备有 {strategy.num_replicas_in_sync} 个")
with strategy.scope():
    model_d = keras_nlp.models.BertMaskedLM.from_preset("bert_tiny_en_uncased")
    model_d.get_layer("bert_backbone").get_layer("pooled_dense").trainable = False
    model_d.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                    optimizer=tf.keras.optimizers.Adam(lr_schedule),
                    weighted_metrics=tf.keras.metrics.SparseCategoricalAccuracy())
    model_d.fit(train_ds, validation_data=val_ds, epochs=EPOCHS, callbacks=callbacks, verbose=1)
model_d.evaluate(test_ds, verbose=1)

结果打印，理论上训练时间是与显卡数量成反比，卡越多训练时间越快，但是在小数据集中效果不是很明显，因为多块显卡之间的通信、复制、同步信息都会耗时，在实际训练中还要保证既要跑满每张显卡又不能 OOM ：

可用的分布式训练的 GPU 设备有 4 个
Epoch 1/5
120/120 [==============================] - 48s 270ms/step - loss: 1.9297 - sparse_categorical_accuracy: 0.0579 - val_loss: 1.7024 - val_sparse_categorical_accuracy: 0.1913
Epoch 2/5
120/120 [==============================] - 29s 241ms/step - loss: 1.6517 - sparse_categorical_accuracy: 0.1644 - val_loss: 1.4522 - val_sparse_categorical_accuracy: 0.2798
Epoch 3/5
120/120 [==============================] - 29s 240ms/step - loss: 1.5088 - sparse_categorical_accuracy: 0.2163 - val_loss: 1.3278 - val_sparse_categorical_accuracy: 0.3198
Epoch 4/5
120/120 [==============================] - 29s 240ms/step - loss: 1.4406 - sparse_categorical_accuracy: 0.2370 - val_loss: 1.2749 - val_sparse_categorical_accuracy: 0.3361
Epoch 5/5
120/120 [==============================] - 29s 241ms/step - loss: 1.4113 - sparse_categorical_accuracy: 0.2448 - val_loss: 1.2603 - val_sparse_categorical_accuracy: 0.3402
15/15 [==============================] - 4s 112ms/step - loss: 1.2633 - sparse_categorical_accuracy: 0.3531

可以看出损失在下降，准确率在提升，可以使用更大的 epoch 继续进行训练。使用 tensorboard 查看训练过程 loos 的变化过程如下：

如何系统的去学习大模型LLM ？

作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。

但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的 AI大模型资料 包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。 那以下这些PDF籍就是非常不错的学习资源。

四、AI大模型商业化落地方案

阶段1：AI大模型时代的基础理解

• 目标：了解AI大模型的基本概念、发展历程和核心原理。
• 内容：
  • L1.1 人工智能简述与大模型起源
  • L1.2 大模型与通用人工智能
  • L1.3 GPT模型的发展历程
  • L1.4 模型工程
  • L1.4.1 知识大模型
  • L1.4.2 生产大模型
  • L1.4.3 模型工程方法论
  • L1.4.4 模型工程实践
  • L1.5 GPT应用案例

阶段2：AI大模型API应用开发工程

• 目标：掌握AI大模型API的使用和开发，以及相关的编程技能。
• 内容：
  • L2.1 API接口
  • L2.1.1 OpenAI API接口
  • L2.1.2 Python接口接入
  • L2.1.3 BOT工具类框架
  • L2.1.4 代码示例
  • L2.2 Prompt框架
  • L2.2.1 什么是Prompt
  • L2.2.2 Prompt框架应用现状
  • L2.2.3 基于GPTAS的Prompt框架
  • L2.2.4 Prompt框架与Thought
  • L2.2.5 Prompt框架与提示词
  • L2.3 流水线工程
  • L2.3.1 流水线工程的概念
  • L2.3.2 流水线工程的优点
  • L2.3.3 流水线工程的应用
  • L2.4 总结与展望

阶段3：AI大模型应用架构实践

• 目标：深入理解AI大模型的应用架构，并能够进行私有化部署。
• 内容：
  • L3.1 Agent模型框架
  • L3.1.1 Agent模型框架的设计理念
  • L3.1.2 Agent模型框架的核心组件
  • L3.1.3 Agent模型框架的实现细节
  • L3.2 MetaGPT
  • L3.2.1 MetaGPT的基本概念
  • L3.2.2 MetaGPT的工作原理
  • L3.2.3 MetaGPT的应用场景
  • L3.3 ChatGLM
  • L3.3.1 ChatGLM的特点
  • L3.3.2 ChatGLM的开发环境
  • L3.3.3 ChatGLM的使用示例
  • L3.4 LLAMA
  • L3.4.1 LLAMA的特点
  • L3.4.2 LLAMA的开发环境
  • L3.4.3 LLAMA的使用示例
  • L3.5 其他大模型介绍

阶段4：AI大模型私有化部署

• 目标：掌握多种AI大模型的私有化部署，包括多模态和特定领域模型。
• 内容：
  • L4.1 模型私有化部署概述
  • L4.2 模型私有化部署的关键技术
  • L4.3 模型私有化部署的实施步骤
  • L4.4 模型私有化部署的应用场景

学习计划：

• 阶段1：1-2个月，建立AI大模型的基础知识体系。
• 阶段2：2-3个月，专注于API应用开发能力的提升。
• 阶段3：3-4个月，深入实践AI大模型的应用架构和私有化部署。
• 阶段4：4-5个月，专注于高级模型的应用和部署。

这份完整版的大模型 LLM 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

😝有需要的小伙伴，可以Vx扫描下方二维码免费领取🆓