1. 项目概述ReLoRA——一种高效的大模型预训练新范式如果你正在训练或微调一个参数规模达到数十亿甚至百亿级别的大型语言模型那么“显存爆炸”和“训练成本高昂”这两个词对你来说一定不陌生。传统的全参数微调Full Fine-tuning需要为每个任务都保存一份完整的模型副本这在大模型时代几乎成了不可承受之重。参数高效微调PEFT技术尤其是LoRALow-Rank Adaptation通过冻结预训练模型的主干参数只训练注入的低秩适配器极大地降低了显存和存储开销。然而LoRA在从头开始预训练一个全新的大模型时其低秩更新的表达能力是否足够一直是个疑问。ReLoRAReinitialized Low-Rank Adaptation正是为了解决这个问题而生。它不是一个全新的微调方法而是一种全新的预训练策略。其核心思想非常巧妙在预训练过程中周期性地将累积的LoRA适配器权重合并回主干网络然后重新初始化新的LoRA适配器继续训练。你可以把它想象成“搭积木”我们不是一次性用所有积木搭一个巨大的结构全参数训练也不是只用几块小积木在表面修修补补标准LoRA微调而是先用几块积木搭一个稳固的底座Warm-up阶段的全参数训练然后不断地用小批量的新积木LoRA更新往上添加每搭完一批就把它们和底座牢牢粘合在一起合并再换一批新积木继续。这样我们最终能用有限的“同时使用的积木”显存搭建出一个庞大而复杂的结构高秩模型。我最初看到这个思路时觉得它既大胆又优雅。它挑战了“预训练必须进行高秩更新”的固有认知通过低秩更新的循环累积理论上可以逼近甚至达到全参数训练的效果同时将训练过程中的峰值显存占用降低一个数量级。这对于我们这些资源有限的团队或个人研究者来说意味着可以用更少的GPU去探索更大规模的模型或者用相同的资源进行更长时间的实验。接下来我将结合Guitaricet团队的开源代码和论文深入拆解ReLoRA的每一个技术细节、实操步骤以及我踩过的那些坑。2. ReLoRA核心原理与设计思路拆解要理解ReLoRA为什么有效以及如何正确使用它我们需要先跳出代码从第一性原理出发看看它到底在解决什么问题。2.1 问题根源低秩更新的表达能力瓶颈标准的LoRA在微调预训练好的大模型时表现优异是因为预训练模型本身已经具备了强大的、高秩的语言表示能力。LoRA所做的是在这个已经非常丰富的表示空间里进行一些低秩的“方向性调整”以适应下游任务。这就像在一幅已经画好的精美素描上用细笔进行一些局部的着色和修饰。但是在从头开始预训练一个模型时情况完全不同。模型初始时是一张白纸随机初始化我们需要通过训练为其注入大量的、复杂的信息构建出高秩的表示空间。如果全程只使用低秩的LoRA适配器进行更新就好比试图只用一支细笔从头画一幅巨幅油画——笔触太细表达力有限训练效率会非常低下很可能无法收敛到一个好的解。2.2 ReLoRA的破局之道循环累积与重启ReLoRA的解决方案可以概括为“低秩更新高秩累积周期性重启”。它包含三个关键阶段全参数预热Warm-up这是整个流程的基石。在训练的最开始我们进行数千步的标准全参数训练。这一步至关重要它的目标不是让模型学会什么具体任务而是让模型从一个随机初始化的“白纸”状态进化到一个具备基础语言建模能力的“粗糙草图”状态。这个阶段的训练为模型注入了初始的、相对高秩的梯度信息建立了一个稳定的优化起点。没有这个阶段直接上LoRA就像在流沙上盖房子。LoRA循环训练阶段预热结束后我们冻结主干网络的所有参数并为其附加一组新的、随机初始化的LoRA适配器。随后我们进入一个循环低秩更新在接下来的N个训练步例如5000步里我们只训练LoRA适配器的参数。由于主干网络被冻结训练时的显存占用和计算量大幅降低。合并与重启达到预设的步数--relora参数后我们将当前LoRA适配器的权重矩阵A和B乘以缩放因子alpha/r然后加到对应的主干网络权重上。紧接着丢弃当前的LoRA适配器及其优化器状态重新初始化一组全新的LoRA适配器。学习率重启伴随着适配器的重启学习率调度器也会被重置通常包含一个短暂的热身--restart_warmup_steps然后进入新的周期。这有助于模型在每次合并后快速适应新的参数空间。循环往复重复步骤2直到达到总训练步数。为什么这样设计打破低秩瓶颈每次合并操作都将本轮低秩更新“固化”到了高秩的主干网络中。经过多个周期的累积主干网络实际上接收了来自多个不同低秩子空间的更新这些更新的总和可以表达出单个低秩更新无法捕捉的复杂模式从而逼近全参数训练的效果。优化稳定性定期重置优化器尤其是Adam中的动量和方差估计可以防止优化器状态因长期累积而变得陈旧或产生偏差这对于这种周期性剧烈变化的参数空间尤为重要。内存效率训练过程中的活跃参数始终是“主干参数 当前LoRA参数”而LoRA参数通常只占主干参数的0.1%~1%。这带来了巨大的显存节省允许使用更大的批次大小Microbatch Size或在不减少批次大小的情况下训练更大的模型。2.3 与相关技术的对比为了更清晰地定位ReLoRA我们可以将其与几种常见技术做个对比技术更新范围训练开销适用场景关键特点全参数训练所有参数极高模型预训练、大规模数据微调表达能力最强资源消耗巨大标准LoRA微调低秩适配器极低下游任务微调参数高效但假设主干模型已预训练好(IA)^3/Adapter少量注入参数低下游任务微调另一种PEFT方式同样依赖预训练主干ReLoRA循环的低秩适配器低大模型从头预训练核心创新将低秩更新用于预训练通过循环合并实现高秩表达梯度检查点所有参数计算时间换空间所有场景减少显存但增加计算时间与ReLoRA正交可结合使用注意ReLoRA并非要取代下游任务微调时的标准LoRA。它的主战场是资源受限情况下的模型预训练。一旦你用ReLoRA预训练好了一个基础模型在用它做下游任务时你依然可以愉快地使用标准的、更轻量的LoRA进行微调。3. 环境搭建与数据预处理实操指南纸上得来终觉浅绝知此事要躬行。理论再完美也需要在代码上跑通。Guitaricet提供的代码库清晰且易于上手但其中仍有一些细节需要特别注意。3.1 环境配置避坑第一站项目要求Python 3.10和PyTorch 2.0主要是为了支持参数注解风格和Flash Attention 2。我的建议是直接使用Conda创建一个干净的环境。# 创建并激活环境 conda create -n relora python3.10 -y conda activate relora # 安装PyTorch请根据你的CUDA版本到官网获取最新命令 # 例如对于CUDA 11.8 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 克隆项目 git clone https://github.com/Guitaricet/peft_pretraining.git cd peft_pretraining # 安装项目依赖 pip install -e .接下来是Flash Attention。这是提升Transformer训练速度的关键尤其是对于长序列。但这里有个小坑官方requirements.txt里没有包含flash-attn因为它的安装脚本对PyTorch等前置依赖的版本非常敏感。因此我们需要在安装完PyTorch和项目基础依赖后再单独安装它。# 安装Flash Attention 2 pip install flash-attn --no-build-isolation如果安装过程中报错通常是因为CUDA环境或PyTorch版本不匹配。请务必确保你的CUDA版本、PyTorch版本和flash-attn版本兼容。可以查阅flash-attn的官方GitHub仓库获取详细的安装指南。实操心得在服务器上我遇到过因为ninja编译工具缺失导致安装失败的情况。可以先尝试pip install ninja或者使用预编译的wheel文件。如果实在无法安装代码也提供了回退方案不使用Flash Attention但训练速度会慢很多。3.2 数据预处理一切训练的基础项目使用Hugging Face的datasets库来加载数据并提供了pretokenize.py脚本进行预处理。预处理的目的很明确将原始文本数据转换为模型可以直接消费的、预先分词并打包好的token id序列并保存成内存映射文件.mmap这样在训练时可以实现高效的流式读取避免IO瓶颈。以预处理C4数据集为例python pretokenize.py \ --save_dir ./preprocessed_data \ --tokenizer t5-base \ # 注意这里默认用T5 tokenizer如果你训练LLaMA需要更改 --dataset c4 \ --dataset_config en \ --text_field text \ --sequence_length 512关键参数解析与避坑点--tokenizer这是第一个大坑。代码示例中使用了t5-base但如果你要预训练一个LLaMA架构的模型你必须使用对应的LLaMA tokenizer。例如你可以使用meta-llama/Llama-2-7b-hf需要授权或社区提供的开源版本。tokenizer不匹配会导致训练完全无效。正确做法在configs/目录下的模型配置文件如llama_250m.json中通常会指定tokenizer_name。确保pretokenize.py使用的tokenizer与之一致。--sequence_length需要与训练时的--max_length参数保持一致。这里设置为512意味着每条样本都会被截断或填充到512个token。--dataset支持Hugging Face datasets库中任何数据集。如果你想用自己的文本数据最简单的方法是将其整理成JSON Lines格式每行一个JSON对象包含text字段然后使用--dataset json和--dataset_config path/to/your/data.jsonl。预处理完成后脚本会输出保存路径格式如preprocessed_data/c4_t5-base_512。这个路径就是后续训练时--dataset_path需要指定的值。注意事项预处理大型数据集如完整的C4可能耗时很长并且需要可观的磁盘空间数百GB。建议在后台运行并使用nohup或tmux保持进程。同时务必检查生成的文件是否完整以及token数量是否符合预期。4. 完整训练流程详解从Warm-up到ReLoRA循环理解了原理准备好了数据和环境我们就可以开始真正的训练了。这个过程分为两个核心阶段全参数Warm-up和ReLoRA主训练。4.1 第一阶段全参数Warm-up这个阶段的目标是获得一个“还不错”的初始模型。我们使用标准的全参数训练。export DATA_PATH./preprocessed_data/c4_your_tokenizer_512 torchrun --nproc-per-node 8 \ --nnodes 1 \ torchrun_main.py \ --model_config configs/llama_250m.json \ --dataset_path $DATA_PATH \ --batch_size 32 \ # 每张GPU上的批次大小 --total_batch_size 1024 \ # 全局有效批次大小 --lr 5e-4 \ # 学习率 --max_length 512 \ --save_every 1000 \ --eval_every 1000 \ --num_training_steps 20000 \ # Warm-up总步数 --tags warm_up_250M参数调校经验--total_batch_size与--batch_size这是分布式训练的关键。脚本会自动计算梯度累积步数Gradient Accumulation Steps。total_batch_size batch_size * num_gpus * gradient_accumulation_steps。不要直接设置--gradient_accumulation让脚本自己算更省心。--lr论文和代码作者建议ReLoRA阶段的学习率大约是常规训练的2倍。因此在Warm-up阶段我们使用一个相对常规的学习率例如5e-4。这个值需要根据模型大小调整对于250M模型可能偏大对于1B模型可能偏小建议参考类似规模模型的常用配置。--num_training_stepsWarm-up需要多少步论文中对1B模型使用了约10K步从某个10K步的检查点开始。对于250M模型我个人经验是5000到20000步都可以目标是让损失曲线明显下降并开始趋于平缓。你可以通过观察eval_loss来判断。--dtype默认为float32。如果你的GPU支持如V100、A100、H100强烈建议使用--dtype bfloat16它能大幅减少显存占用且几乎不影响精度。训练结束后在checkpoints目录下会生成一个带时间戳的文件夹里面保存了最终的模型检查点例如model_20000。记下这个路径它是下一阶段的输入。4.2 第二阶段ReLoRA主训练这是核心阶段。我们加载Warm-up好的模型冻结其主干开始LoRA的循环训练。torchrun --nproc-per-node 8 \ torchrun_main.py \ --model_config configs/llama_250m.json \ --dataset_path $DATA_PATH \ --batch_size 64 \ # 注意由于显存节省可以增大batch_size --total_batch_size 1024 \ --lr 1e-3 \ # 学习率提高约为Warm-up阶段的2倍 --max_length 512 \ --use_peft \ # 启用PEFTLoRA --lora_r 128 \ # LoRA的秩rank --lora_alpha 256 \ # 缩放因子通常设置为r的2倍 --relora 5000 \ # 关键每5000步合并并重置一次LoRA --cycle_length 5000 \ # 学习率调度器周期长度通常等于relora步数 --restart_warmup_steps 100 \ # 每次重置后的学习率热身步数 --scheduler cosine_restarts \ # 使用带重启的余弦退火 --warmup_steps 500 \ # 初始学习率热身步数 --reset_optimizer_on_relora True \ # 每次重置LoRA时也重置优化器 --num_training_steps 100000 \ # ReLoRA阶段总步数 --save_every 5000 \ --eval_every 5000 \ --warmed_up_model ./checkpoints/llama_250m-2023-xx-xx-xx-xx-xx/model_20000 \ --tags relora_250M_mainReLoRA专属参数深度解析--relora这是最重要的参数定义了“合并-重置”的周期。如何选择论文中提供了一个经验公式并建议让总训练步数能被relora步数整除。例如总步数143K他们选择了5320。一个简单的原则是周期不宜过短否则频繁合并重置开销大且可能不稳定也不宜过长否则单次低秩更新时间太长可能限制表达能力。对于250M模型从2000到10000步都可以尝试。建议从5000步开始。--reset_optimizer_on_relora默认True即每次合并后完全重置Adam优化器的状态动量、方差。这是最稳妥的方案。你也可以设置为False并配合--optimizer_magnitude_pruning 0.9这会对优化器状态进行“幅度修剪”保留90%最大的状态值丢弃10%最小的。论文发现两者效果相近完全重置更简单。--scheduler cosine_restarts与--cycle_length必须使用这种带重启的余弦退火调度器。cycle_length需要设置为和relora相同的值以保证每次LoRA重置时学习率也正好完成一个周期并重启热身。--restart_warmup_steps每次周期重启后学习率从0线性增长到最大值所需的步数。这个值通常很小50-200目的是让模型在参数空间发生跳跃合并操作后能平稳地开始新的探索。--lora_rLoRA的秩。在预训练场景下这个值需要比微调时大得多。微调常用4,8,16而ReLoRA论文中对1B模型使用了128甚至256。更大的秩意味着更强的单周期表达能力但参数量和计算量也会增加。这是一个需要权衡的超参数。--batch_size注意由于主干参数被冻结显存占用大幅下降你通常可以将batch_size提高到Warm-up阶段的1.5倍甚至2倍从而加快训练速度。这是ReLoRA带来的直接红利。5. 分布式训练配置与性能调优对于大模型单卡训练是不现实的。项目使用PyTorch原生的DDP进行单机多卡分布式训练配置起来非常直观。5.1 基础分布式启动命令torchrun --nproc-per-node NUM_GPUS \ --nnodes 1 \ --node-rank 0 \ --master-addr localhost \ --master-port 29500 \ torchrun_main.py 你的所有参数--nproc-per-node你机器上的GPU数量。--nnodes 1节点数为1即单机。其他master-*参数在单机情况下通常用默认值即可。一个针对调试的简化示例torchrun --nproc-per-node 2 torchrun_main.py \ --model_config configs/llama_35m.json \ --use_peft \ --lora_r 64 \ --relora 500 \ --cycle_length 500 \ --reset_optimizer_on_relora False \ --optimizer_magnitude_pruning 0.9 \ --lr 0.001 \ --batch_size 120 \ # 单卡batch --total_batch_size 480 \ # 全局batch 120 * 2 * 2(梯度累积) --num_training_steps 5000 \ --dtype bfloat16 \ --tags debug_run5.2 性能调优与监控最大化吞吐量核心是增大--batch_size直到GPU显存用满。使用nvidia-smi监控显存。在ReLoRA阶段由于大部分参数被冻结你可以尝试非常大的batch_size。同时确保开启了Flash Attention查看日志输出以获得最佳速度。梯度累积通过--total_batch_size间接控制。如果你想保持较大的全局批次但单卡显存不足脚本会自动增加梯度累积步数。监控日志中的gradient_accumulation_steps值。混合精度训练务必使用--dtype bfloat16。与float16相比bfloat16具有与float32相同的指数位动态范围更大在训练中更加稳定是现代AI训练尤其是大模型的事实标准。日志与监控代码使用Weights Biases (wandb)进行日志记录。在命令行添加--wandb参数并在环境中设置WANDB_API_KEY就可以在网页上实时查看损失曲线、学习率、吞吐量等指标。这对于调整relora周期、学习率等超参数至关重要。6. 常见问题排查与实战心得在实际操作中你几乎一定会遇到各种问题。下面是我总结的一些典型情况及其解决方案。6.1 训练不收敛或损失异常高这是最令人头疼的问题。请按以下顺序排查检查tokenizer90%的异常问题源于tokenizer不匹配确认pretokenize.py使用的tokenizer与model_config.json中指定的tokenizer_name完全一致。一个快速验证的方法是加载预处理后的一个数据样本用模型的tokenizer解码一下看看是不是可读的英文或其他语言。检查学习率学习率过大或过小都会导致问题。对于ReLoRA阶段尝试从1e-3开始如果损失出现NaNNot a Number立即停止将学习率降低一个数量级1e-4再试。Warm-up阶段的学习率也应保守一些。检查Warm-up是否充分如果Warm-up步数太少模型基础能力太弱ReLoRA阶段可能无法有效学习。观察Warm-up结束时的评估损失它应该已经下降到相对较低且平稳的水平。检查relora周期周期太短如100步会导致训练过程被频繁的重置打断无法有效学习周期太长则失去了低秩更新的意义。首次尝试建议将relora设置为总训练步数的1/20到1/10。关闭Flash Attention虽然罕见但某些环境下的Flash Attention实现可能存在兼容性问题导致精度错误。可以尝试移除flash-attn库回退到PyTorch的原生Attention实现看看问题是否消失。6.2 显存占用高于预期确认是否启用了--use_peft如果在ReLoRA阶段显存和Warm-up阶段一样大很可能是因为忘记添加--use_peft参数导致仍然在进行全参数训练。检查--dtype确保是bfloat16而非float32。检查模型配置确认model_config.json中的hidden_size,num_hidden_layers,num_attention_heads等参数是否符合你的预期没有因为配置错误导致模型比想象中大得多。梯度累积增大--total_batch_size而单卡--batch_size不变时脚本会增加梯度累积步数这不会增加峰值显存但会延长计算时间。6.3 评估损失在ReLoRA重置后剧烈波动这是正常现象每次执行合并操作Relora Reset时模型参数发生了一次离散的跳跃损失曲线通常会有一个向上的尖峰。紧接着由于学习率重启并进入短暂的热身期损失会快速下降。你应该关注的是多个周期平均下来的损失下降趋势而不是单个点的瞬时值。只要整体趋势是向下的就说明训练是有效的。6.4 如何保存和加载最终的模型训练完成后最终的检查点例如model_100000是一个完整的模型它已经包含了所有合并进去的LoRA权重。你可以像加载任何普通PyTorch模型一样加载它用于后续的推理或作为基础模型进行微调。import torch from model import Transformer config ... # 加载你的模型配置 model Transformer(config) state_dict torch.load(path/to/model_100000/pytorch_model.bin) model.load_state_dict(state_dict) # 现在model就是一个标准的、可用于推理的Transformer如果你想提取出纯主干网络不含任何LoRA结构代码中可能提供了合并脚本或者你需要手动将LoRA权重合并到线性层后保存。具体可以查看项目仓库的utils/merge_peft.py或类似脚本。6.5 我的独家心得从小模型开始实验不要一上来就用1B参数跑。用configs/llama_35m.json3500万参数这个配置做实验快速验证你的数据流水线、tokenizer和超参数选择是否正确。一两小时就能跑完一个实验周期效率极高。使用Wandb做超参数扫描relora周期、lora_r、lr这三个是最关键的超参数。用Wandb的sweep功能设计一个小范围的网格搜索能帮你快速找到适合你模型和数据的最优组合。关注“有效吞吐量”由于ReLoRA每隔N步就要做一次合并操作这涉及GPU同步和参数更新这会引入额外的开销。计算“每秒处理的样本数”时要包含这个开销。如果relora周期设得太短比如100步这个开销占比会很大反而降低整体效率。通常周期在几千步时开销可以忽略不计。理解其局限性ReLoRA是一种内存优化技术它用更长的训练时间由于低秩更新可能需要更多步数达到相同效果和略微增加的算法复杂度换取了大幅降低的显存占用。它不是为了“更快”地训练模型而是为了“在有限资源下训练更大”的模型。如果你的目标是在8张A100上最快地预训练一个7B模型全参数训练配合ZeRO-3可能仍是首选。但如果你只有2张4090却想尝试预训练一个1B模型ReLoRA几乎是唯一可行的路径。最后我想说的是ReLoRA代表了一种非常实用的研究思路在算力约束下寻找创新。它没有提出全新的模型结构而是通过巧妙地改造训练流程极大地降低了大型模型预训练的门槛。对于广大中小实验室和个人开发者而言这类工作比那些动辄需要千卡集群的“巨无霸”模型往往具有更直接的参考价值和启发性。我强烈建议你亲手跑一遍代码从3500万参数的小模型开始感受一下这种“循环累积”的力量。当你看到损失曲线在一次次重置后依然稳步下降时你会对大规模模型训练有更深刻的理解。