这两年大模型训练的门槛下降得很快但一旦真要动手把自己的小模型训练出来很多小型实验室还是会被硬件预算卡住。最近看到 Puro-2B 这类项目感觉它把很多团队心里想但没有写出来的方案变成了一个可参考的标题Poor Labs Qwen2-1.5B Trained on RTX 5090 within $5090。翻译过来就是一个预算有限的实验室用一张 RTX 5090 显卡基于 Qwen2-1.5B 底座在约 5090 美元的总预算内完成 2B 级模型训练实验。预算数字又恰好和显卡型号撞名算是一点技术人的幽默。这篇文章就围绕这条“单卡训练 2B 模型”的路径展开。先讲清楚 Puro-2B 这类项目背后的技术选型再把 RTX 5090 上训练 Qwen2-1.5B 的环境搭建、显存估算、LoRA 与全参训练脚本、常见报错和工程实践完整拆解一遍。适合预算有限又想做模型训练实验的学生团队、小型创业公司也适合刚入手 RTX 5090、想知道这块卡到底能做什么的同学。1. Puro-2B 是什么小实验室的 2B 级模型探索1.1 项目名称背后的含义从项目名可以看出几个关键信息首先是主体Poor Lab即经费有限的实验室其次是底座Qwen2-1.5B选择的是阿里通义千问 Qwen2 系列里最小的尺寸之一再就是训练设备一张 RTX 5090最后是预算约束整体控制在 5090 美元附近。这种命名的价值在于它把“没有 H100 集群也能做模型训练”这件事摆到了台面上。过去大家默认训练大语言模型至少需要多卡 A100消费级显卡只能做推理。但参数量降到 1.5B 到 2B 这个量级之后单张 32GB 显存的高端消费卡就有了操作空间。RTX 5090 恰好属于这个档位Blackwell 架构、32GB GDDR7 显存既比专业计算卡便宜又比大多数旧旗舰卡显存大。Puro-2B 具体是如何从 Qwen2-1.5B 变成 2B 的需要看项目仓库的 config 和训练日志才能确定。常见做法有两种一种是在 Qwen2-1.5B 基础上通过 LoRA 或全参微调得到约 2B 规模的模型另一种是做模型扩容Model Surgery比如扩大中间层维度或词表再继续预训练。本文后面提供的实训流程覆盖前一种主流做法扩容思路会在原理部分说明。1.2 为什么选 Qwen2-1.5B 作为底座Qwen2-1.5B 是通义千问 Qwen2 系列中参数量较小的一员发布于 2024 年采用标准的 decoder-only Transformer 结构使用 GQA分组查询注意力、SwiGLU 激活函数和 RMSNorm。它有两个对低成本训练非常友好的特点。一是许可证宽松。Qwen2 系列使用 Apache 2.0 协议可以商用也允许基于它微调后再发布只需要保留版权声明和许可信息。这对小型团队来说很重要不用担心闭源协议带来的合规风险。二是生态成熟。Hugging Face 上可以直接下载权重transformers 从 4.40 版本开始就支持 Qwen2 架构Peft、TRL、DeepSpeed、vLLM 等常用工具都能直接对接。相比去复现一个冷门架构用 Qwen2-1.5B 起步可以把精力集中到训练本身。1.5B 这个尺寸也处于一个性价比拐点0.5B 虽然更快但语言能力和推理能力偏弱很多下游任务表现不够7B 模型在单卡上做全参训练则非常吃力。1.5B 往上扩到 2B 左右既能做一些真实的领域能力验证又不会让显存直接爆炸。1.3 需要区分的几个训练概念在实际动手之前先把几个容易混淆的术语分清楚继续预训练Continued Pre-training在原始预训练语料或领域语料上继续训练让模型补充某个领域的知识。指令微调SFT用 Instruction 和 Response 配对数据训练让模型学会按指令回答。偏好对齐DPO/RLHF让模型的回答更符合人类偏好通常需要额外构建偏好数据。LoRA低秩适配冻结原模型权重只训练一小部分新增的低秩矩阵显存占用远低于全参训练。蒸馏Distillation用一个大模型生成数据或软标签训练一个小模型。Puro-2B 标题里写的 Trained往往不是单一操作而是上述多种手段的组合。低成本实验室通常的路线是先做领域继续预训练再做 SFT 对齐条件允许再用 DPO 调一轮偏好。本文实战部分先给出最基础也最容易跑通的 LoRA SFT 方案再补充全参训练的思路。2. 环境准备RTX 5090 单卡训练环境搭建2.1 硬件选型与预算估算RTX 5090 的海外建议零售价为 1999 美元国内实际价格受市场供需影响波动较大。按项目标题的 5090 美元总预算来规划显卡大概占一半左右剩余预算要覆盖 CPU、主板、内存、电源、硬盘和散热。有几个硬件细节建议优先考虑电源RTX 5090 整卡功耗较高官方规格在 575W 左右瞬时功耗可能更高建议搭配额定 1000W 以上、支持 ATX 3.1 / 12V-2x6 接口的电源。内存训练时除了 GPU 显存CPU 内存也要加载数据集和做预处理32GB 起步64GB 更稳。硬盘模型权重、数据集和中间 checkpoint 都很占空间建议 2TB NVMe SSD训练时把数据放在本地而不是网络盘。散热长时间训练要注意机箱风道训练过程中可以用nvidia-smi观察温度超过 85℃ 时考虑加强散热或降低功耗墙。如果不想自购硬件也可以租用 RTX 5090 云服务器按小时付费短周期实验更划算。2.2 软件版本说明RTX 5090 属于 NVIDIA Blackwell 架构对软件版本要求比旧显卡更高。本文示例以 Linux 环境为例使用 Ubuntu 22.04/24.04Python 3.10 或 3.11。版本组合如下需要根据你的实际项目调整组件建议版本说明NVIDIA 驱动支持 Blackwell 的最新稳定版以nvidia-smi能正确识别 RTX 5090 为准CUDA12.8 及以上Blackwell 架构需要较新 CUDAPyTorch2.7 及以上选择 cu128 构建新卡尽量用新版 PyTorchtransformers4.40 及以上建议直接用最新稳定版peft0.11 及以上LoRA 训练datasets2.20 及以上数据集加载accelerate0.34 及以上Trainer 依赖bitsandbytes最新版做量化时使用注意 Blackwell 兼容性Windows 也能训练但很多底层库对 Blackwell 的适配先在 Linux 上完成遇到奇怪的编译错误时WSL2 或纯 Linux 环境会省心很多。2.3 创建虚拟环境并验证 GPU建议用 conda 或 venv 建一个独立环境避免污染系统 Python。以下是 conda 示例conda create -n puro2b python3.11 -y conda activate puro2b pip install torch2.7.1 --index-url https://download.pytorch.org/whl/cu128 pip install transformers datasets accelerate peft bitsandbytes tensorboard安装完成后先做一次基础验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))预期输出类似2.7.1cu128 True NVIDIA GeForce RTX 5090如果torch.cuda.is_available()返回 False优先检查驱动版本和 PyTorch 的 CUDA 构建是否匹配不要急着怀疑显卡坏了。3. 核心原理单卡训练 2B 模型的显存账本3.1 参数量与显存估算为什么单张 32GB 显卡能训练 2B 级别模型先算一笔显存账。全参训练时显存主要消耗在四部分模型权重、梯度、优化器状态和激活值。以 2B 参数模型、BF16 精度、AdamW 优化器为例模型权重2B × 2 字节 4GB梯度2B × 2 字节 4GB优化器状态AdamW 需要保存 fp32 精度的参数副本、一阶动量、二阶动量约 2B × 12 字节 24GB三部分相加已经达到 32GB刚好顶到 RTX 5090 的显存上限而激活值还没有计算。这说明全参训练 2B 模型在 32GB 单卡上非常紧张必须靠梯度检查点、梯度累积、优化器 CPU offload 等手段压显存。如果改用 LoRA情况完全不同。LoRA 冻结原始权重只训练低秩矩阵可训练参数通常只占模型总参数的 0.5% 到 2%。2B 模型的 LoRA 可训练参数可能只有几千万优化器状态只有几百 MB原始权重 4GB 常驻显存即可训练变得非常宽裕。3.2 BF16 混合精度为什么是首选Blackwell 架构本身支持 FP8 甚至 FP4但训练主流的做法仍然是 BF16 混合精度。原因在于数值稳定性。FP16 只有 5 位指数位表示范围的整数部分有限在梯度数值跨度大的场景容易出现溢出或下溢BF16 有 8 位指数位和 FP32 相同的指数范围虽然尾数精度下降但训练稳定性明显更好。在 transformers Trainer 里bf16True即可开启。需要确认显卡驱动和 PyTorch 构建支持 BF16RTX 5090 完全没问题。3.3 省显存的关键技巧当全参训练遇到显存不足时按优先级尝试以下手段开启梯度检查点Gradient Checkpointing不保存前向过程的全部激活值反向时重新计算显存可能减少一半代价是训练速度变慢。Trainer 里设置gradient_checkpointingTrue。调整批次大小并配合梯度累积per_device_train_batch_size降到 1 或 2通过gradient_accumulation_steps保证等效批次大小不变。优化器状态 CPU OffloadDeepSpeed ZeRO-Offload 或 accelerate 的 CPU offload把 fp32 优化器状态放到内存显存压力大幅下降但会引入 CPU-GPU 通信开销。缩短序列长度2B 模型虽然支持 32K 上下文但训练时如果不需要长文本把max_length设为 1024 或 2048能显著减少激活值。使用 LoRA / QLoRA这是单卡训练最稳妥的方案建议作为第一选择。3.4 从 1.5B 到 2B 的扩容思路如果目标是把 Qwen2-1.5B 真正变成 2B 参数模型而不是仅仅做适配就需要模型扩容。一个常见思路是保留预训练权重通过修改 config 增大隐藏层维度、中间层维度或层数然后把原始权重映射到新结构里未初始化的部分随机初始化最后在大规模语料上继续预训练。这种操作在工程上叫 Model Surgery难点不在扩参本身而在于扩容后必须投入足够的训练数据才能让新参数真正生效否则模型能力不升反降。单张 RTX 5090 做这种实验不是不行但训练数据规模和训练时长都会受限于单卡算力建议先用小规模数据验证流程再决定是否投入完整训练。4. 实战基于 Qwen2-1.5B 的 LoRA 微调全流程4.1 项目结构与数据准备先建立一个清晰的目录结构puro-2b-lab/ ├── data/ │ ├── train.jsonl │ └── eval.jsonl ├── scripts/ │ └── train_lora_sft.py ├── output/ │ └── puro-2b-lora/ └── README.md训练数据使用 JSONL 格式一条数据一行。这里以指令微调为例每行包含 instruction 和 output 两个字段{instruction: 用一句话解释什么是梯度下降, output: 梯度下降是一种通过沿损失函数梯度的反方向更新参数逐步最小化损失函数的优化算法。}数据准备要注意两点一是数据量不是越大越好几千条高质量指令就能让模型变化明显二是训练集和验证集要分开避免用同一批数据评估效果。4.2 编写训练脚本下面是完整的 LoRA SFT 脚本文件路径为scripts/train_lora_sft.py。核心流程是加载模型与分词器构造 LoRA 配置格式化数据使用 Trainer 训练。import torch from datasets import load_dataset from transformers import ( AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling, Trainer, TrainingArguments, ) from peft import LoraConfig, get_peft_model model_name Qwen/Qwen2-1.5B output_dir ./output/puro-2b-lora # 1. 加载模型和分词器 tokenizer AutoTokenizer.from_pretrained(model_name, trust_remote_codeTrue) if tokenizer.pad_token is None: tokenizer.pad_token tokenizer.eos_token model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.bfloat16, trust_remote_codeTrue, ) # 2. 配置 LoRA lora_config LoraConfig( r16, lora_alpha32, target_modules[q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj], lora_dropout0.05, biasnone, task_typeCAUSAL_LM, ) model get_peft_model(model, lora_config) model.print_trainable_parameters() # 3. 加载并格式化数据 dataset load_dataset(json, data_filesdata/train.jsonl, splittrain) eval_dataset load_dataset(json, data_filesdata/eval.jsonl, splittrain) def format_example(example): text ( f|im_start|user\n{example[instruction]}|im_end|\n f|im_start|assistant\n{example[output]}|im_end|\n ) return {text: text} def tokenize_function(examples): return tokenizer( examples[text], truncationTrue, max_length1024, paddingFalse, ) dataset dataset.map(format_example).map(tokenize_function, remove_columns[instruction, output, text]) eval_dataset eval_dataset.map(format_example).map(tokenize_function, remove_columns[instruction, output, text]) data_collator DataCollatorForLanguageModeling(tokenizertokenizer, mlmFalse) # 4. 训练参数 training_args TrainingArguments( output_diroutput_dir, per_device_train_batch_size2, per_device_eval_batch_size2, gradient_accumulation_steps8, learning_rate2e-4, num_train_epochs3, logging_steps10, save_steps200, eval_strategysteps, eval_steps200, save_total_limit3, bf16True, gradient_checkpointingTrue, optimadamw_torch, report_to[tensorboard], seed42, ) # 5. 训练 trainer Trainer( modelmodel, argstraining_args, train_datasetdataset, eval_dataseteval_dataset, data_collatordata_collator, tokenizertokenizer, ) trainer.train() trainer.save_model(output_dir)脚本里手动拼接的|im_start|格式与 Qwen2 官方聊天模板一致也可以改用tokenizer.apply_chat_template实现同样的效果。注意训练阶段不加add_generation_prompt否则会多出一个空的 assistant 提示符。4.3 运行训练在项目根目录执行conda activate puro2b python scripts/train_lora_sft.py训练开始时model.print_trainable_parameters()会输出类似下面的信息trainable params: 18,874,368 || all params: 1,569,317,376 || trainable%: 1.2027这说明模型总参数约 1.57BLoRA 只训练了约 1887 万参数占比 1.2% 左右。正是这个比例让单卡训练变得轻松。训练过程中可以打开 TensorBoard 观察 loss 曲线tensorboard --logdir ./output/puro-2b-lora/runs预期 loss 整体呈下降趋势。具体数值取决于数据集质量、学习率和训练步数不要拿别人博客里的 loss 数值当标准关键是看曲线是否稳定收敛。4.4 加载模型并验证效果训练完成后用下面的脚本做一次生成验证import torch from transformers import AutoModelForCausalLM, AutoTokenizer from peft import PeftModel model_name Qwen/Qwen2-1.5B adapter_path ./output/puro-2b-lora base_model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.bfloat16, device_mapauto, trust_remote_codeTrue, ) model PeftModel.from_pretrained(base_model, adapter_path) model model.merge_and_unload() tokenizer AutoTokenizer.from_pretrained(model_name, trust_remote_codeTrue) prompt |im_start|user\n用一句话解释什么是梯度下降|im_end|\n|im_start|assistant\n inputs tokenizer(prompt, return_tensorspt).to(cuda) outputs model.generate( **inputs, max_new_tokens128, temperature0.7, top_p0.9, do_sampleTrue, ) print(tokenizer.decode(outputs[0], skip_special_tokensTrue))输出是示意性质的例如模型会给出一个完整的定义解释。实际效果和你的训练数据强相关如果生成结果不符合预期先检查数据格式是否正确再考虑增加训练轮次或数据量。4.5 全参继续预训练的配置思路如果要做全参继续预训练核心配置和 LoRA 脚本类似但需要额外处理显存压力。推荐使用 DeepSpeed ZeRO-Offload在scripts/ds_config_zero2_offload.json写入{ zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, contiguous_gradients: true, overlap_comm: true }, bf16: { enabled: true }, train_batch_size: auto, train_micro_batch_size_per_gpu: auto, gradient_accumulation_steps: auto }然后在 TrainingArguments 里加一行deepspeedscripts/ds_config_zero2_offload.json。训练数据换成纯文本格式去掉指令模板用 DataCollatorForLanguageModeling 做语言建模。注意全参训练时学习率通常要调低建议 1e-5 到 5e-5 之间warmup 比例也要适当加大。5. 常见问题与排查思路RTX 5090 上市时间还不长软件生态适配仍在完善中训练过程中遇到问题很正常。以下按故障现象整理常见的排查路径问题现象常见原因解决思路torch.cuda.is_available() 返回 False驱动版本过老或 PyTorch 的 CUDA 构建不匹配升级 NVIDIA 驱动安装 cu128 以上版本的 PyTorchCUDA out of memory批次过大、序列过长或未开启梯度检查点降低 batch size、缩短 max_length、开启 gradient checkpointingbitsandbytes 量化报错旧版 bitsandbytes 缺少 Blackwell kernel升级 bitsandbytes 到最新版或改用 BF16 训练FlashAttention 编译失败早期 FA2 不支持 Blackwell暂时使用 attn_implementationsdpaloss 不下降或剧烈震荡学习率过高、数据质量差、批次过小降低学习率、清洗数据、增大等效 batch 或增加 warmup训练中显卡温度过高降频散热不足或机箱风道差控制功耗墙、改善散热、监控温度5.1 数据格式错误导致 loss 不降指令微调数据最容易踩的坑是格式不统一。比如有的数据没写完整 assistant 回复有的把换行符丢失有的 text 字段拼接后无法被 tokenizer 正确切分。建议在训练前单独写一个脚本随机打印 20 条格式化后的样本人工确认格式无误再启动训练。5.2 梯度检查点与显存优化冲突gradient_checkpointingTrue和某些并行策略或model.compile()组合时可能报错。如果开启后报自定义函数相关错误可以先关闭 torch.compile只用梯度检查点。单卡训练场景下这两项优化并不是必须叠加的。5.3 保存和加载 adapter 路径混乱Trainer 保存到./output/puro-2b-lora后目录里会有adapter_config.json、adapter_model.safetensors和原始模型文件。加载时用PeftModel.from_pretrained指定 adapter 目录而不是再from_pretrained加载整个目录否则可能重复加载权重。合并权重后如果准备部署可以用model.save_pretrained(merged)单独保存一份合并模型。6. 最佳实践与工程建议6.1 预算和实验资源管理$5090 的预算听起来不多但如果规划合理足够支撑一个持续 2 到 3 个月的实验周期。建议把预算拆成三块显卡和整机硬件占一半以上数据集整理和标注留一部分剩余作为云服务器应急备份。租赁显卡按小时计费适合大规模短期训练自购硬件适合长期小步快跑。在 RTX 5090 上训练时不要把per_device_train_batch_size盲目调大。2B 模型的 LoRA 训练batch size 2 配合梯度累积 8等效 batch 已经是 16对大多数指令数据集足够。全参训练则优先考虑 CPU offload而不是强行压缩 batch。6.2 实验跟踪与可复现性训练实验最怕“改了一堆参数最后不知道哪个改动生效”。建议每个实验固定以下信息训练数据版本、超参数、随机种子、模型 commit hash、训练时长。可以用 wandb 或 TensorBoard 记录简单做法是在训练脚本里用argparse把所有超参写成命令行参数配合一个实验记录表格。模型 checkpoint 建议保存save_total_limit3保留最近三个中间状态。大模型训练过程中如果断电或 OOM可以从最近的 checkpoint 用trainer.train(resume_from_checkpointTrue)恢复训练避免前功尽弃。6.3 数据质量优先于数据数量低成本训练的容错空间很小数据质量直接影响最终模型效果。继续预训练阶段优先去重和过滤低质量文本指令微调阶段优先保证格式统一、答案正确而不是堆大量噪声数据。一个常见的工程做法是先用大模型生成一批候选数据人工抽样检查后再进入训练集。6.4 合规与安全边界使用 Qwen2-1.5B 做底座时Apache 2.0 协议允许商用和修改但需要保留原始 LICENSE、NOTICE 和版权声明。训练数据要确认来源授权不要使用未经授权的爬虫数据或侵权内容。发布微调模型时建议在模型卡里明确写出底座模型、训练数据来源、训练方法方便他人评估和使用。这里的核心原则是训练环境越受限越要在数据合规和授权边界上做足功课避免模型训练出来却因为数据问题无法继续使用。6.5 性能优化优先级排序单卡训练的性能优化建议按以下顺序判断收益先确认数据加载不是瓶颈num_workers 0数据预处理提前完成。再用 SDPA 或可用的 FlashAttention 替代朴素注意力减少 attention 计算时间。然后考虑 torch.compileBlackwell 架构在较新 PyTorch 版本上支持度更好但首次编译时间长。最后才考虑分布式或多卡方案单卡没跑通之前不要引入多卡复杂度。记住训练速度和显存占用经常是矛盾的不要在同一个实验里尝试所有优化逐项验证收益。7. 最后给想复现这条路的人一点实在建议Puro-2B 这类项目的意义不在于它的模型参数比别人多也不在于跑分比别人高而在于验证了一条低成本路径一张 RTX 5090、一个开源的 Qwen2-1.5B 底座、一份干净的数据集就能完成完整的 2B 级模型训练闭环。如果你准备自己动手建议按三步推进。第一步用 LoRA SFT 小规模数据跑通整个流程确认环境、脚本、验证链路都没有问题第二步把数据规模和训练时长逐步扩大观察 loss 收敛情况和生成效果第三步再考虑全参继续预训练或模型扩容这类更重的手段。每一步都保留好实验记录和 checkpoint方便回退对比。RTX 5090 的 Blackwell 生态还在快速完善中新版本 PyTorch、bitsandbytes、FlashAttention 都在陆续适配。这篇文章里的版本建议只能代表写作时点的状态动手前最好到官方仓库确认最新版本遇到编译问题优先升级依赖而不是去改复杂的底层代码。祝你的 2B 模型顺利出炉。