LLaMA 分布式训练实践

一、首先解释一下DeepSpeed：

        DeepSpeed 是由微软推出的一个开源深度学习优化库，旨在帮助研究人员和开发者在训练大规模深度学习模型时提高效率和性能，特别是在分布式训练和大规模模型微调等任务上。DeepSpeed 提供了多种技术和工具来优化计算资源的使用，从而在现有硬件（如 GPU 和 TPU）上加速训练速度、减少内存消耗，并降低训练成本。

DeepSpeed 的主要功能和特点：

1. Zero Redundancy Optimizer (ZeRO)： DeepSpeed 的 ZeRO 优化器是其最为知名的创新之一。它通过将模型的参数、梯度和优化器状态分布在多个 GPU 之间，极大地减少了显存占用，从而使得训练超大规模模型成为可能。例如，ZeRO 使得数百亿甚至上万亿参数的模型能够在常规硬件上训练。

2. 分布式训练： DeepSpeed 支持高度优化的分布式训练，能够在多个 GPU 和节点之间高效地划分和协调计算工作。它不仅支持数据并行、模型并行等常见的分布式训练模式，还针对大规模模型训练提供了专门的优化。

3. 混合精度训练： DeepSpeed 提供了混合精度训练的支持，可以在不牺牲模型性能的前提下，通过使用 16 位和 32 位浮点数混合运算来加速训练并减少显存的使用。

4. 高效的 GPU 资源管理： DeepSpeed 提供了各种优化策略来提高 GPU 的资源利用率，比如通过梯度累积和小批量训练，最大化 GPU 的内存和计算能力的使用。

5. 支持大规模模型的推理加速： 除了训练，DeepSpeed 还提供了对大模型推理的优化。通过其推理引擎，DeepSpeed 可以显著加速推理速度，并减少推理时的内存占用，适合用于生产部署。

6. 适配主流框架： DeepSpeed 主要与 PyTorch 框架集成良好，它兼容许多基于 PyTorch 的模型和工具，用户可以非常容易地将现有的 PyTorch 代码迁移到 DeepSpeed 环境中。

DeepSpeed 的优势

• 大规模模型训练：特别适合 GPT、BERT 等拥有大量参数的大型模型。借助 ZeRO 和分布式优化技术，它能够将训练这些超大规模模型所需的硬件资源显著减少。
• 高效利用内存和显存：通过先进的内存管理技术，如 ZeRO 和梯度检查点等，DeepSpeed 能够训练和推理更大的模型，同时减少内存占用。
• 成本效益：DeepSpeed 帮助用户通过更少的硬件和更短的时间训练大规模模型，从而减少硬件成本和能源消耗。
• 灵活性：用户可以选择不同的优化级别，适配各种硬件配置，从单 GPU 到大规模集群都能高效运行。

pip install deepspeed

二、 训练数据配置

模块/函数	来源	功能
DataLoader	PyTorch	用于从数据集中加载数据，以便模型进行训练或评估。
RandomSampler	PyTorch	随机采样数据集中的样本，适合在无特定顺序要求的训练场景中使用
SequentialSampler	PyTorch	顺序采样数据集中的样本，适合在有顺序要求的训练或评估场景中使用。
DistributedSampler	PyTorch	用于分布式训练，确保每个进程/节点获取数据集的不同子集，避免数据重复。
default_data_collator	Transformers	Transformers 库中的默认数据收集器，用于将多个样本整合成一个批量数据。
create_pretrain_dataset	自定义函数	用于创建预训练数据集，基于输入参数来生成适合模型训练的输入数据格式。

代码通过检查 args.local_rank 是否为 -1 来确定使用单机采样器还是分布式采样器，具体处理方式如下：

1. 单机模式 (args.local_rank == -1)： 如果 local_rank 等于 -1，表明模型运行在单机上，无需分布式处理。此时，代码会选择 PyTorch 的 RandomSampler 或 SequentialSampler：

   • RandomSampler 随机选择数据集中的样本进行训练，适合打乱数据的场景。
   • SequentialSampler 按照顺序采样数据，通常用于评估阶段，确保评估时数据顺序固定，保证结果的稳定性。

2. 分布式模式 (args.local_rank != -1)： 当 local_rank 不等于 -1，意味着模型运行在多机或多进程的分布式训练环境中。此时，代码会选择 DistributedSampler，该采样器能确保每个进程或节点接收数据集中的一个独立且不重复的子集。

   • DistributedSampler 的工作原理是将数据集划分为若干个子集，确保每个节点处理的数据互不重叠。这在大规模分布式训练中非常重要，因为它能显著提升计算效率，避免多次重复计算相同数据。
   • 这样，训练过程可以并行化进行，同时还能保证每个进程都能公平地获得数据子集，保持训练过程中的负载平衡。

---

ps：

• 在单机模式下：训练通常不涉及复杂的数据同步机制，数据加载器直接从本地数据集中按需采样。因此，使用随机或顺序采样器即可满足需求。随机采样有助于增强训练的泛化能力，而顺序采样则确保测试或验证时，输入数据的顺序一致。

• 在分布式模式下：由于多个 GPU 或节点同时处理不同部分的数据，需要使用分布式采样器进行数据分配。这样每个节点只处理数据集的一部分，保证了效率最大化，并避免多个进程处理相同的数据子集。分布式采样器还有助于减少通信开销，因而在大规模训练时非常重要。

• 采样器的重要性：采样器的选择不仅影响训练效率，还直接影响训练的稳定性与性能。在随机采样时，数据的打乱顺序可以帮助模型避免过拟合特定的样本顺序，而顺序采样有助于验证和测试时保持结果的一致性。此外，在分布式训练中，采样器的合理设计能显著减少节点间的数据冗余和同步延迟，提高总体训练性能。

  ---

  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
  from torch.utils.data.distributed import DistributedSampler
  from transformers import default_data_collator
  from utils.data.data_utils import create_pretrain_dataset
  
  # 准备数据集
  train_dataset, eval_dataset = create_pretrain_dataset(
      args.local_rank,
      args.data_path,
      args.data_split,
      args.data_output_path,
      args.seed,
      tokenizer,
      args.max_seq_len
  )
  
  # 选择采样器
  if args.local_rank == -1:
      train_sampler = RandomSampler(train_dataset)