PEFT × Transformer Engine 实战基于 ESM2 的 LoRA 蛋白质二级结构 Token 分类微调指南【免费下载链接】peft PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.项目地址: https://gitcode.com/gh_mirrors/pe/peft本篇技术指南以仓库中的 examples/lora_finetuning_transformer_engine 示例为蓝本系统讲解如何在 NVIDIA Transformer EngineTE加速的 ESM2 蛋白质语言模型上使用 PEFT 的 LoRA 适配器完成 token 分类蛋白质二级结构预测微调。读完本文你将掌握从 Docker/虚拟环境两种方式搭建 TE 运行环境、理解lora_finetuning_te.py的完整训练管线合成数据生成、Trust Remote Code加载、Trainer训练与评估以及如何切换到 Porter6 真实数据集进行更贴近实际的实验。示例概览TE 与 PEFT 的组合为什么值得关注Transformer EngineTE是 NVIDIA 提供的高性能 Transformer 算子库它能在 NVIDIA GPU 上自动选用最合适的融合算子与数值精度如 FP8、BF16来加速 Transformer 骨干网络。蛋白质领域的 ESM2 系列模型经 TE 改造后其自注意力中的layernorm_qkv等融合模块可以显著提升训练与推理吞吐。而 PEFTParameter-Efficient Fine-Tuning的价值在于只用极少量的可训练参数本例中只向目标模块注入低秩矩阵即可完成适配大幅降低显存与存储开销。两者结合便构成了TE 加速骨干 LoRA 轻量适配的高效蛋白质模型微调方案。本示例的核心流程可以概括为五步加载基于 Transformer Engine 的 ESM2 token 分类模型通过 PEFT 的LoraConfigget_peft_model注入 LoRA 适配器在代码内生成随机类蛋白序列并构造合成二级结构标签H、E、C使用 Hugging FaceTrainer完成训练与评估无需手工搭建 DDP保存 LoRA 适配器权重与 tokenizer 到输出目录。环境搭建两种方式任选其一方式 ADocker推荐基于 NVIDIA 官方公开的 PyTorch 容器镜像nvcr.io/nvidia/pytorch:26.01-py3构建自包含镜像。该镜像已内置 CUDA、cuDNN 与 Transformer Engine可以避免 TE 与本地 CUDA 工具链的版本匹配问题docker build -t lora-te examples/lora_finetuning_transformer_engine构建完成后直接运行训练挂载 GPUdocker run --gpus all --rm lora-te \ python lora_finetuning_te.py \ --base_model nvidia/esm2_t6_8M_UR50D \ --output_dir ./esm2_lora_output \ --num_train_samples 256 \ --num_eval_samples 64 \ --num_epochs 1也可以启动交互式会话进行探索式实验docker run --gpus all --rm -it lora-te bash镜像的构建逻辑十分简洁完整内容见 examples/lora_finetuning_transformer_engine/Dockerfile它以 NVIDIA PyTorch 容器为基础将工作目录设为/workspace/lora_finetuning先安装requirements.txt中的依赖再拷贝训练脚本。由于基础镜像自带 Transformer Engine因此 Dockerfile 中不再重复安装 TE。方式 B虚拟环境创建并激活 Python 虚拟环境后安装依赖python -m venv .venv source .venv/bin/activate pip install -r examples/lora_finetuning_transformer_engine/requirements.txt需要特别注意的是Transformer Engine 必须单独安装且版本必须与系统 CUDA 工具链版本匹配。这一点在 examples/lora_finetuning_transformer_engine/requirements.txt 中也有明确注释——torch、transformers、datasets、peft、accelerate、numpy、pandas、pyarrow、safetensors均由 pip 安装而transformer-engine一行被注释掉要求要么单独安装、要么使用系统自带的版本TE 官方安装指南对 CUDA 版本匹配有详细说明。因此方式 ADocker能帮你规避绝大多数 TE 与 CUDA 的兼容性问题。训练脚本的源码级拆解训练主脚本是 examples/lora_finetuning_transformer_engine/lora_finetuning_te.py下面按执行顺序逐段剖析。单卡锁定TE 模型与Trainer的兼容性前提脚本在导入torch之前就执行了一行关键代码os.environ.setdefault(CUDA_VISIBLE_DEVICES, 0)其原因是TE 支撑的模型与Trainer的DataParallel包装不兼容。通过把可见 GPU 固定为单卡torch.cuda.device_count()恒为 1从根上避免Trainer触发DataParallel分支。这是使用 TE 模型时容易踩坑、而示例脚本已经替你处理好的细节。标签体系与合成数据规则脚本定义了 SS33 态二级结构标签映射SS3_ID2LABEL {0: H, 1: E, 2: C} SS3_LABEL2ID {label: idx for idx, label in SS3_ID2LABEL.items()}其中H表示 α-螺旋helix、E表示 β-折叠beta strand、C表示卷曲/无规则coil。合成数据由两个层次构成build_synthetic_sequences从 20 种标准氨基酸ACDEFGHIKLMNPQRSTVWY中按min_seq_len~max_seq_len的随机长度生成随机序列sequence_to_synthetic_labels按氨基酸的理化倾向映射出伪标签——HELIX_AA set(AELMQKRH)中的残基标为HBETA_AA set(VIFYWT)中的残基标为E其余标为C。residue_to_ss_char与ss_char_to_label分别完成氨基酸→二级结构字符、二级结构字符→标签 ID 的转换未知字符回退到C。这套规则并非生物学上的真实预测而是为了让管线跑得通、能验证适合做快速冒烟测试与正确性检查。Token 化与标签对齐tokenize_and_align_labelsToken 分类任务的关键难点在于把每个残基的标签对齐到 token 位置。脚本的实现方式是labels [-100] * len(input_ids) usable_len min(len(sequence), len(label_str), len(input_ids) - 2) for idx in range(usable_len): labels[idx 1] ss_char_to_label(label_str[idx])要点解析ESM2 的 tokenizer 采用逐字符切分每个氨基酸大致对应一个 token因此可以采用idx 1的偏移跳过[CLS]/起始特殊 token做近似对齐长度取sequence、label_str、input_ids - 2三者的最小值避免越界超出有效区间的 label 保持-100这是 Hugging Face 生态中忽略该位置损失的约定值DataCollatorForTokenClassification与 loss 计算都会自动跳过。模型加载trust_remote_code是硬性前提脚本对 ESM2 系列模型做了前置校验if not args.trust_remote_code and esm2 in args.base_model.lower(): raise ValueError( fModel {args.base_model} requires remote code execution. Re-run with --trust_remote_code to confirm you trust this models code. )原因在于Hugging Face Hub 上的默认 ESM2 模型带有自定义建模代码custom modeling code必须显式传入--trust_remote_code才允许加载。这不是可选项而是运行前置条件——README 中也以醒目的 Note 形式强调了这一点。加载时脚本还完成了三件事config AutoConfig.from_pretrained(args.base_model, trust_remote_codeargs.trust_remote_code) config.num_labels 3 config.id2label SS3_ID2LABEL config.label2id SS3_LABEL2ID model AutoModelForTokenClassification.from_pretrained( args.base_model, configconfig, trust_remote_codeargs.trust_remote_code, dtypemodel_dtype, )将num_labels固定为 3与 SS3 三分类对应写入id2label/label2id映射便于推理时还原标签名依据torch.cuda.is_available() and torch.cuda.is_bf16_supported()自动选择bfloat16或float32精度use_bf16随后同步传给TrainingArguments(bf16...)在支持 BF16 的 Ampere 及以上架构 GPU 上获得更低显存占用与更高吞吐。LoRA 注入LoraConfig与get_peft_model这是整个示例的 PEFT 核心lora_config LoraConfig( task_typeTaskType.TOKEN_CLS, rargs.lora_r, lora_alphaargs.lora_alpha, lora_dropoutargs.lora_dropout, target_modules[layernorm_qkv], biasnone, inference_modeFalse, ) model get_peft_model(model, lora_config) model.print_trainable_parameters()各参数含义与底层实现对应 PEFT 源码 src/peft/tuners/lora/config.py参数示例默认值说明task_typeTaskType.TOKEN_CLS任务类型。PEFT 会依据它选择对应的包装类——从 src/peft/auto.py 的映射可见TOKEN_CLS对应PeftModelForTokenClassification从而自动适配Trainer的 token 分类 loss 与输出结构。rLoRA 秩8脚本默认低秩矩阵的维度秩越大可学习容量越高、可训练参数越多。lora_alpha16脚本默认缩放系数。PEFT 中适配器的实际缩放因子为lora_alpha / r用于控制注入分支对原始输出的影响强度。lora_dropout0.05脚本默认注入分支上的 dropout 概率用于缓解过拟合。target_modules[layernorm_qkv]注入目标模块。TE 化 ESM2 将 LayerNorm 与 QKV 投影融合为一个模块因此这里直接以模块名layernorm_qkv为目标——这是与普通nn.Linear模型通常 targetq_proj/v_proj最显著的差异点。biasnone是否训练偏置参数none表示不训练任何 bias进一步压缩可训练参数量。inference_modeFalse关闭推理模式允许适配器参与训练。注入完成后调用model.print_trainable_parameters()。该方法定义于 src/peft/peft_model.py会输出形如trainable params: 1,234,567 || all params: 92,345,678 || trainable%: 1.3370的统计方便第一时间确认 LoRA 把可训练参数压到了总参数的极小比例——这正是参数高效微调的直接证据。训练与评估Trainer全流程数据侧脚本支持两种来源合成数据默认make_synthetic_dataset按num_train_samples/num_eval_samples生成训练与评估集经 token 化后移除原始列Porter6 parquet 数据load_parquet_dataset用pandas.read_parquet读取训练/验证 parquet 文件要求包含Sequence与Secondary_structure两列映射后同样移除原列。Trainer侧的关键配置TrainingArgumentseval_strategysteps与save_strategysteps按eval_steps/save_steps默认 25 步周期评估与保存save_total_limit2最多保留 2 个 checkpoint控制磁盘占用load_best_model_at_endTruemetric_for_best_modeltoken_accuracygreater_is_betterTrue训练结束后自动回载评估指标最优的 checkpointremove_unused_columnsFalse保留数据集中的全部列token 分类场景必要bf16use_bf16与前面模型加载时选择的精度保持一致report_tonone不向 wandb 等外部平台上报适合无痕实验。评估指标由compute_metrics实现对 logits 取 argmax 得到预测用labels ! -100掩码过滤掉填充位置后计算 token 级准确率token_accuracy。训练结束后执行model.save_pretrained(args.output_dir) tokenizer.save_pretrained(args.output_dir)将 LoRA 适配器权重/配置与 tokenizer 文件一并落盘model此时是 PEFT 包装模型save_pretrained默认保存的是适配器而非完整骨干权重。运行示例在环境就绪的前提下直接运行训练脚本python examples/lora_finetuning_transformer_engine/lora_finetuning_te.py \ --base_model nvidia/esm2_t6_8M_UR50D \ --output_dir ./esm2_lora_output \ --num_train_samples 256 \ --num_eval_samples 64 \ --num_epochs 1默认模型nvidia/esm2_t6_8M_UR50D是 ESM2 系列中最小的 6 层、8M 参数模型UR50D 指 UniRef50 蛋白质数据库 D 版本训练非常适合用来验证整条管线。若想自行加载 Hub 上其他 ESM2 模型记得补上--trust_remote_code。自定义训练参数所有可调参数均通过命令行传入完整定义见 lora_finetuning_te.py 的parse_args下面是一次覆盖常用参数的示例python examples/lora_finetuning_transformer_engine/lora_finetuning_te.py \ --base_model nvidia/esm2_t6_8M_UR50D \ --trust_remote_code \ --output_dir ./esm2_lora_output \ --max_length 256 \ --batch_size 4 \ --learning_rate 3e-4 \ --lora_r 16 \ --lora_alpha 32 \ --lora_dropout 0.1完整的参数清单及默认值如下表源码中的parse_args即为权威出处参数默认值作用--base_modelnvidia/esm2_t6_8M_UR50DTE 化 ESM2 模型名或本地路径--output_dir./esm2_lora_output输出目录适配器 tokenizer checkpoint--max_length128最大序列长度超长截断--num_train_samples256合成训练样本数--num_eval_samples64合成评估样本数--min_seq_len/--max_seq_len32/96合成序列长度范围--batch_size8每设备批大小--num_epochs1训练轮数--learning_rate5e-4学习率--weight_decay0.01权重衰减--logging_steps10日志输出步频--eval_steps/--save_steps25/25评估/保存步频--seed42随机种子Python/numpy/torch 全链路固定见set_seed--lora_r8LoRA 秩--lora_alpha16LoRA 缩放系数--lora_dropout0.05LoRA dropout--trust_remote_codeFalse允许加载 Hub 自定义建模代码ESM2 必填--train_parquet/--val_parquetNonePorter6 训练/验证 parquet 路径二者必须同时给出或同时省略注意--train_parquet与--val_parquet存在互斥校验——只传其中一个会直接触发parser.error并终止程序这是脚本内建的防呆逻辑。数据集从合成数据切换到 Porter6脚本默认在运行时生成合成数据集随机类蛋白序列配随机生成的二级结构标签H/E/C。它适合快速冒烟测试与管线验证但标签并非真实生物注释不能用于衡量真实预测能力。若要做更贴近实际的评估可换用Porter6 二级结构数据集。Porter6 是蛋白质二级结构预测领域的常用基准数据集训练集约 55k 条序列、验证集为 2024 年构建的约 692 条序列。仓库中并未内置 Porter6 数据但提供了一个可行的获取路径BioNeMo 仓库bionemo-recipes的esm2_peft_te/data目录中附有prepare_porter6_dataset.py下载与转换脚本可据此生成训练/验证 parquet 文件。生成后将 parquet 路径传给训练脚本即可此时合成数据不再生成python examples/lora_finetuning_transformer_engine/lora_finetuning_te.py \ --base_model nvidia/esm2_t6_8M_UR50D \ --train_parquet porter6_train_dataset_55k.parquet \ --val_parquet porter6_val_dataset_2024_692.parquet \ --output_dir ./esm2_lora_output \ --num_epochs 3parquet 文件的读取逻辑位于load_parquet_dataset两列必须命名为Sequence氨基酸序列字符串与Secondary_structure等长二级结构标签字符串读取后通过Dataset.from_pandas(...).map(...)批量 token 化并移除原始列随后进入与合成数据完全相同的训练流程。真实数据规模较大建议配合更大的--num_epochs如示例中的 3以获得有意义的收敛结果。输出产物训练结束后--output_dir下会生成PEFT LoRA 适配器权重与配置adapter_model.safetensors等权重文件与adapter_config.json配置Tokenizer 文件tokenizer 配置与词表便于加载适配器后直接做推理训练过程中按--save_steps保存、受--save_total_limit2约束的中间 checkpoint。由于保存的是轻量适配器而非完整骨干权重适配器文件体积远小于 ESM2 原始权重方便分享、版本管理与多任务切换。需要推理时加载原 TE 模型后通过 PEFT 的PeftModel.from_pretrained或AutoPeftModelForTokenClassification挂载适配器即可复用。小结与延伸本示例演示了一条可完整跑通的TE 加速 ESM2 PEFT LoRA微调链路关键经验可以总结为三点环境先行优先使用基于nvcr.io/nvidia/pytorch的 Docker 镜像规避 TE 与 CUDA 版本匹配问题自建环境时需单独安装与 CUDA 匹配的 Transformer Engine。单卡约束TE 模型与Trainer的DataParallel不兼容务必通过CUDA_VISIBLE_DEVICES锁定单卡。适配器注入注意点TE 化模型的target_modules应指向融合模块layernorm_qkv而非传统q_proj/v_proj加载 Hub 上的 ESM2 模型必须携带--trust_remote_code。如需更多 Transformer Engine 加速 Transformer 的示例可以进一步参考 BioNeMo 框架的bionemo-recipes仓库中esm2_peft_te相关目录PEFT 侧的 LoRA 配置项细节则可查阅 src/peft/tuners/lora/config.py 与 src/peft/peft_model.py 的源码注释那里对lora_alpha缩放机制、target_modules语义、print_trainable_parameters统计口径都有权威说明。【免费下载链接】peft PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.项目地址: https://gitcode.com/gh_mirrors/pe/peft创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考