当前位置：首页 > news >正文

探索多种方案下 LLM 的预训练性能

news 来源：原创 2025/5/4 3:17:35

Base 版

此处采用最基础 Torch
的 DDP

算法，作为基础的对比版本，运行的脚本如下：

torchrun --nnodes 1 --nproc_per_node 8 pretrain_hf.py
–model_config_path …/config/config.json
–tokenizer_name_or_path …/ckpt/Llama-2-13b-hf
–per_device_train_batch_size 8
–do_train
–seed 1234
–fp16
–num_train_epochs 1
–lr_scheduler_type cosine
–learning_rate 2e-5
–warmup_ratio 0.05
–weight_decay 0.01
–logging_strategy steps
–logging_steps 10
–save_strategy steps
–save_total_limit 1
–save_steps 100
–gradient_accumulation_steps 8
–model_max_length 2048
–output_dir ‘./hf_logs’
–overwrite_output_dir
–gradient_checkpointing
–ddp_find_unused_parameters False

FSDP 版

FSDP是从DeepSpeed ZeRO
以及FairScale的FSDP中获取灵感的一种完全数据分片并行的方法。

HuggingFace 的 transformers 已经原生支持了 FSDP，可以在原生代码的运行命令上增加2行配置来使用：

--fsdp "full_shard auto_wrap" \
--fsdp_config 'fsdp.json' \

HuggingFace 的 transformers 已经原生支持了 DeepSpeed，其用法也非常简单，只需要添加一个配置文件和一个参数即可：

以下是zero-2的配置文件

{
“fp16”: {
“enabled”: “auto”,
“loss_scale”: 0,
“loss_scale_window”: 100,
“initial_scale_power”: 16,
“hysteresis”: 2,
“min_loss_scale”: 1e-10
},

"zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 1e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 1e8,
    "contiguous_gradients": true
},

"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": false

}

使用该配置文件如下：

–deepspeed ‘ds_zero2_no_offload.json’ \

运行结果如下所示，与原版 DDP 相比性能提升约 15.4% ：

***** train metrics *****
epoch = 0.98
train_loss = 7.6718
train_runtime = 0:16:35.35
train_samples_per_second = 31.297
train_steps_per_second = 0.04

如果使用 zero3 算法，结果如下，zero2 略有下降：

***** train metrics *****
epoch = 0.98
train_loss = 9.5239
train_runtime = 0:16:55.82
train_samples_per_second = 30.667
train_steps_per_second = 0.01

如果使用 zero3 算法同时使用 offload，性能略有下降：

***** train metrics *****
epoch = 0.98
train_loss = 9.3981
train_runtime = 0:17:21.20
train_samples_per_second = 29.919
train_steps_per_second = 0.01

FlashAttention
if training_args.flash_attn:
from utils.flash_attn_patch import replace_llama_attn_with_flash_attn
replace_llama_attn_with_flash_attn()

只需要在运行文件中增加一句使能即可：

--flash_attn \

运行结果如下所示，与原版 DDP 相比性能提升约 52.7% ：

***** train metrics *****
epoch = 0.98
train_loss = 8.0186
train_runtime = 0:12:32.22
train_samples_per_second = 41.413
train_steps_per_second = 0.04

FlashAttention + FSDP 性能如下，略有下降：

***** train metrics *****
epoch = 0.98
train_loss = 9.9216
train_runtime = 0:12:55.09
train_samples_per_second = 40.191
train_steps_per_second = 0.039

FlashAttention + DeepSpeed 性能如下，略有提升：

***** train metrics *****
epoch = 0.98
train_loss = 8.0758
train_runtime = 0:12:17.55
train_samples_per_second = 42.237
train_steps_per_second = 0.041

ColossalAI 的优化算法实现 llama 模型的预训练。

首先尝试 Gemini 算法，这是 ColossalAI 提出的异构内存空间管理器。它通过在 CPU 和 GPU 中容纳模型数据，并仅在必要时将数据移动到当前设备，可以同时利用 GPU 内存、CPU 内存（由 CPU DRAM 或 NVMe SSD内存组成）来突破单GPU内存墙的限制。

目前 DeepSpeed采用的 Zero-offload 在CPU和GPU内存之间静态划分模型数据，并且它们的内存布局对于不同的训练配置是恒定的。如下图左边所示，当 GPU 内存不足以满足其相应的模型数据要求时，即使当时CPU上仍有可用内存，系统也会崩溃。而 ColossalAI 可以通过将一部分模型数据换出到CPU上来完成训练，这就是Gemini。它管理CPU和GPU二者内存空间。它的内存管理器由两部分组成，分别是MemStatsCollector(MSC)和StatefulTensorMgr(STM)，可以让张量在训练过程中动态分布在CPU-GPU的存储空间内。
Zero-Offload和Gemini的内存管理方案比较

Gemini 还利用了深度学习网络训练过程的迭代特性，将迭代分为warmup和non-warmup两个阶段，开始时的一个或若干迭代步属于预热阶段，其余的迭代步属于正式阶段。在warmup阶段为MSC收集信息，而在non-warmup阶段STM入去MSC收集的信息来移动tensor，以达到最小化CPU-GPU数据移动volume的目的。
Gemini在不同训练阶段的运行流程

下面是该算法的用法：

colossalai run --nproc_per_node 8 pretrain_colossalai.py
–config …/config/config.json
–plugin gemini_cuda
–batch_size 28
–lr 2e-5
–weigth_decay 0.01
–warmup_steps 2
–grad_checkpoint
–max_length 2048
–mixed_precision fp16
–flash_attention

从日志可以看到该算法对GPU及CPU的内存分配情况：

Booster init max CUDA memory: 1523.89 MB
Booster init max CPU memory: 6737.21 MB
Max CUDA memory usage: 36608.19 MB

此时可观察到 GPU 及 CPU 的使用情况，可以清楚看到该算法在二者间动态使用的过程
GPU 的使用情况
CPU 的使用情况

性能如下所示，与原版 DDP 相比性能提升约 50.7%，其中主要的性能提升来源于 FlashAttention：

train_runtime = 0:12:42.53
train_samples_per_second = 40.882

如果不使能 FlashAttention，性能如下，与原版 DDP 相比性能提升约 8.1 %：

train_runtime = 0:17:42.36
train_samples_per_second = 29.333

Sophia Optimizer

Sophia 优化器使用随机估计作为 Hessian 矩阵对角线的 pre-conditioner，并采用剪切（clipping）机制来控制最坏情况下的参数大小更新。在像 GPT-2 这样的预训练语言模型上，Sophia 与 Adam 相比，在减少了 50% step 数量的情况下实现了相同的验证预训练损失，这相当于总计算量减少了 50%，wall-clock 时间减少了 50%。