如何用 annotated_deep_learning_paper_implementations 的 Zero3 显存优化微调 GPT-NeoX 小模型-平芜编程栈

如何用 annotated_deep_learning_paper_implementations 的 Zero3 显存优化微调 GPT-NeoX 小模型

【免费下载链接】annotated_deep_learning_paper_implementations🧑‍🏫 60+ Implementations/tutorials of deep learning papers with side-by-side notes 📝; including transformers (original, xl, switch, feedback, vit, ...), optimizers (adam, adabelief, sophia, ...), gans(cyclegan, stylegan2, ...), 🎮 reinforcement learning (ppo, dqn), capsnet, distillation, ... 🧠项目地址: https://gitcode.com/gh_mirrors/an/annotated_deep_learning_paper_implementations

annotated_deep_learning_paper_implementations（labml-nn）提供了一套 Zero-DP 显存优化实现（代码里叫Zero3），可以把 GPT-NeoX 的参数、梯度和优化器状态切分到多张 GPU 上，从而在显存装不下整模型时仍能做多卡微调。本文的任务就是：在多张 CUDA GPU 上，用仓库里的 finetune_neox.py 脚本，借助 Zero-DP 优化器对 GPT-NeoX 做偏差参数（bias）微调，并用filter_layers把模型缩到只加载部分层的小模型。前提是你有可用的多张 NVIDIA GPU 和 NCCL 支持，并且本地已经拿到 GPT-NeoX 的 checkpoint 文件。

环境与前提

先安装库依赖。仓库 readme.md 给出的安装方式是：

pip install labml-nn

除了依赖，还要满足三个前提，缺任一项脚本都跑不起来：

多张 CUDA GPU + NCCL。脚本用torch.distributed.init_process_group('nccl', ...)组建进程组，每个 rank 绑定一张cuda:{rank}。脚本启动时会打印nccl可用性，需要先确认是True。
GPT-NeoX checkpoint 文件已就位。层权重从本地数据路径加载（见 checkpoint.py），加载用的是torch.load，不会自动下载。该模块提供了download()（其__main__会执行download()，默认n_layers=44）用于拉取这些文件。
训练文本会自动下载。text_dataset.py 首次运行会下载tiny_shakespeare语料，这块不用手动准备。

微调的对象是bias 参数。finetune.py 里FineTuneBiases把每个TransformerLayer整层设成requires_grad=False，再单独放开attention.output.bias、attention.qkv_lin.bias、ffn.dense_h_h4.bias三个 bias 让它们可训练。这正好对应 Zero3 文档里“只实现 Zero-DP、支持只训练一部分参数”的说明。

用 Zero-DP 包住 GPT-NeoX 的每一层

脚本的 finetune_neox.html 文档里，模型和 optimizer 各用一个@option覆盖：

@option(Configs.model, 'Zero3') def _model(c: Configs): from labml_nn.scaling.zero3 import Zero3Layer, Zero3Sequential _ = c.fine_tuner # 确保 fine tuner 先把可训练参数设好 modules = [] for m in monit.iterate('Zero3', c.layers): modules.append(Zero3Layer(m.to(c.device), c.rank, c.world_size, c.device, c.dtype)) model = Zero3Sequential(modules) return model @option(Configs.optimizer, 'Zero3Adam') def _optimizer(c: Configs): from labml_nn.optimizers.adam_fp16 import AdamFP16 return AdamFP16(c.model.get_trainable_chunk(), lr=c.learning_rate)

这里要做两件事：

逐层包装：c.layers是 GPT-NeoX 的各层，每层被Zero3Layer包住，再一起塞进Zero3Sequential。zero3 实现中，Zero3Layer在__init__里把该层的可训练/固定参数合并、按world_size切分并scatter到各设备，然后resize_(0)释放本地完整参数的显存；Zero3Sequential再为所有层挂上fetch_stream、backup_stream以及前/后层指针。
只把分片喂给优化器：get_trainable_chunk()返回的是当前 rank 上那一份可训练 shard，AdamFP16只对这一份分片维护状态，而不是完整参数。

这就是 Zero-DP 省显存的来源。index.html 给出的公式是：显存降到原模型的(2 + 2 + K)·Ψ / N_d，其中Ψ是参数量、N_d是分片数、K是每参数的优化器字节数；2 + 2指 16-bit 的参数与梯度，Adam 下K = 12（fp32 参数副本 + 两个 fp32 动量）。对应地，通信量是O(3Ψ)，而普通数据并行是O(2Ψ)。这是文档给出的期望关系，不是某次跑出来的固定数值，实际占用随层数与N_d变化。

缩小模型：只加载部分层（filter_layers）

默认Configs继承自 PipelineParallelTrainerConf，会加载全部 44 层。要跑“小模型”，用 model.py 里LayerGenerator的filter_layers参数——文档原文说明它是“用来测试只有更少层的更小模型版本”，并给了{0, 1}这种“只加载前两层”的示例（见 generate 文档）。

配置里c.layers由filter_layers透传进LayerGenerator(...).load()，所以可以直接在配置里指定一个子集，例如只加载前两个 transformer 层：

# 在 experiment.configs 的字典里追加这一项，把模型缩到只有指定层 'filter_layers': {0, 1}

不写filter_layers（或为None）时加载全部层；
写成set时，只加载集合里索引对应的层，层数越少显存越小。

注意：filter_layers只影响加载哪些已存在的层文件，这些层的 checkpoint 文件仍需先在本地（见上一节的download()）。

配置项与启动方式

脚本在 finetune_neox.py 里用experiment.configs灌入运行时参数：

experiment.configs(conf, { 'model': 'Zero3', 'optimizer': 'Zero3Adam', 'device': device, 'rank': rank, 'world_size': world_size, 'learning_rate': 3e-4, 'max_seq_len': 128, 'batch_size': 16, })

其中：

model/optimizer选中的是上面两个@option（'Zero3'、'Zero3Adam'）；
device、rank、world_size在main里按当前进程填好；
learning_rate、max_seq_len、batch_size是脚本选定的训练值（覆盖基类默认）；epochs沿用TrainerConf默认 16；
需要小模型时，在此字典里加'filter_layers': {0, 1}。

多卡启动由torch.multiprocessing.spawn负责，脚本对本机每张 GPU 起一个进程：

n_gpu = torch.cuda.device_count() # 每张 GPU 起一个进程；跨多台机器需要额外的 launcher torch.multiprocessing.spawn(main, args=(n_gpu,), nprocs=n_gpu, join=True)

于是world_size就等于本机 GPU 数，rank是spawn的下标（0 .. n_gpu-1）。每个进程在main里先初始化分布式、再绑定cuda:{rank}：

torch.distributed.init_process_group('nccl', timeout=datetime.timedelta(seconds=30), init_method=init_method, # 默认 'tcp://localhost:23456' rank=rank, world_size=world_size) device = torch.device(f'cuda:{rank}') torch.cuda.set_device(device)

文档明确提醒：本机一卡一进程用spawn即可，跨多台机器时需要另外的 launcher——默认脚本不覆盖多机场景。

判断运行是否正常

按顺序看这几个信号即可确认 Zero-DP 微调真的在多卡上跑起来：

启动配置日志。脚本入口会用inspect打印每张 GPU 名称、n_gpus、mpi、nccl。先确认nccl为True、n_gpus与你的卡数一致——否则 NCCL 进程组建不起来。
Distributed段完成。init_process_group('nccl')成功、随后逐层构建Zero3Layer模型（脚本特意在进训练循环前触发conf.model，为了让日志更清晰）。
每个 step 有 loss/acc 记录。训练循环走 trainer.py 的train_epoch()，tracker每步记录loss.train、acc.train；本脚本writers={'screen', 'labml'}，即屏幕和 labml 面板都能看到 loss 随 step 变化。
显存按分片占用。每张卡只保留自己那份 shard，文档给出的期望占用是(2 + 2 + K)·Ψ / N_d，卡数越多单卡占用越低。