FlexGen 仓库中的 Performer 高效注意力微调实战：基于 Flax/JAX 的掩码语言模型训练指南

推理引擎大模型【免费下载链接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.项目地址https://gitcode.com/gh_mirrors/fl/FlexGen点击查看免费下载导读本文围绕 FlexGen 仓库内置的 Hugging Face Transformers research_projects 示例完整讲解如何基于PerformerFAVOR 快速注意力机制在 Flax/JAX 生态下对 BERT 进行掩码语言建模MLM微调。你将掌握两个开箱即用的训练脚本简单百科与完整英文维基百科、全部核心命令行参数--performer、--reinitialize、--wandb_user_name等并通过阅读源码理解 Performer 用随机特征映射把注意力复杂度从二次降为线性的底层原理从而能够在资源受限环境下训练超长序列模型。一、项目背景这份示例在 FlexGen 仓库中的位置该 Performer 微调示例位于 benchmark/third_party/transformers/examples/research_projects/performer 目录下是 FlexGen 仓库 benchmark 体系中以third_party方式内置的 Hugging Face Transformers 研究项目。它展示了与 FlexGen 主项目在单 GPU 上以吞吐为导向运行大语言模型主题高度互补的另一个维度在训练/微调阶段通过高效注意力降低显存与算力消耗。目录内共包含 6 个文件构成一套完整的最小可运行研究项目文件作用README.md项目说明、依赖、示例与关键参数sanity_script.sh快速验证脚本bert-base-cased Simple Wikipediafull_script.sh完整实验脚本bert-large-cased English Wikipediarun_mlm_performer.py训练入口参数解析、数据加载、MLM 数据整理器、训练/评估循环modeling_flax_performer.pyFlax 版 Performer 模型定义注意力替换为 FAVORmodeling_flax_performer_utils.py快速注意力核心算法随机特征映射、低秩分解二、环境依赖与安装前提原文档明确列出运行时依赖为datasets、flax和jax并说明wandb集成是内置的可选启用。结合 run_mlm_performer.py 的导入语句实际依赖清单如下datasets从 Hugging Face Hub 加载维基百科数据集load_datasetflax与jax模型定义、jax.pmap多设备并行、jax.lax底层算子全部基于 JAX 生态transformers提供BertConfig、FlaxBertForMaskedLM、AutoTokenizer、HfArgumentParser、TrainingArguments等基础组件numpy、tqdm数据处理与进度显示可选tensorboard脚本在启动时会调用is_tensorboard_available()检测未安装则打印提示并降级为不记录指标可选wandb只有传入--wandb_user_name时才在运行时import wandb并初始化。一个值得注意的细节是脚本在模型加载阶段使用的是jnp.float32dtype并没有默认开启混合精度如果希望加速训练可以关注TrainingArguments中的 fp16 相关开关。这保证了示例代码在任何 JAX 可用的 CPU/GPU/TPU 环境都能直接运行。三、两个开箱即用的微调脚本原文档提供了两个脚本分别用于快速验证和完整实验其命令均可直接复制运行。3.1 快速验证脚本 sanity_script.shsanity_script.sh 从bert-base-cased检查点出发在Simple Wikipedia 数据集datasets 提供的一个小规模、用简单英语编写的维基百科子集上微调TOKENIZERS_PARALLELISMtrue python run_mlm_performer.py --output_dir experiments \ --dataset_name wikipedia --dataset_config_name 20200501.simple \ --model_name_or_path bert-base-cased --tokenizer_name bert-base-cased \ --do_train --overwrite_output_dir \ --per_device_train_batch_size 4 --learning_rate 5e-4 \ --warmup_steps 100 --num_train_epochs 3 --performer3.2 完整实验脚本 full_script.shfull_script.sh 使用更大的bert-large-cased检查点在English Wikipedia 完整数据集上微调适合真正验证 Performer 在长序列、大数据量下的效果TOKENIZERS_PARALLELISMtrue python run_mlm_performer.py --output_dir experiments \ --dataset_name wikipedia --dataset_config_name 20200501.en \ --model_name_or_path bert-large-cased --tokenizer_name bert-large-cased \ --do_train --overwrite_output_dir \ --per_device_train_batch_size 4 --learning_rate 5e-4 \ --warmup_steps 100 --num_train_epochs 3 --performer两个脚本的差异仅在--dataset_config_name20200501.simplevs20200501.en与--model_name_or_path/--tokenizer_namebase vs large。命令开头的TOKENIZERS_PARALLELISMtrue用于关闭 tokenizer 并行警告属于 Transformers 生态的常见惯例。四、核心命令行参数详解原文档列出了五个关键参数这里结合 run_mlm_performer.py 中的ModelArguments、DataTrainingArguments、WandbArguments三个 dataclass 逐一展开4.1--performer启用 FAVOR 注意力核心开关原文档说明移除--performer参数即可使用标准 Bert 模型。在源码中这个开关直接决定模型类别的选择lm_class FlaxPerformerForMaskedLM if model_args.performer else FlaxBertForMaskedLM即传--performer时加载 modeling_flax_performer.py 中的FlaxPerformerForMaskedLM其注意力层会调用make_fast_softmax_attention构造的快速注意力函数不传则回退为标准的FlaxBertForMaskedLM。这使你可以用同一套数据与训练逻辑直接对比标准注意力与 Performer 的精度和吞吐差异——这正是研究型实验最需要的 A/B 能力。4.2--reinitialize从空白模型开始训练原文档说明添加--reinitialize将从空白模型blank model而非 Bert 检查点开始。源码对应逻辑为if model_args.reinitialize: model lm_class(configBertConfig.from_pretrained(model_args.model_name_or_path)) else: model lm_class.from_pretrained(model_args.model_name_or_path, ...)传--reinitialize时仅用model_name_or_path读取BertConfig的架构参数层数、头数、隐层维度等权重全部随机初始化不传则加载预训练权重。这用于回答Performer 在从头训练时的收敛性如何这类研究问题。4.3--model_name_or_path更换 BERT 规模原文档指出可将该参数换成 Hugging Face Hub 上任意预训练检查点来改变 BERT 规模。源码中该参数有双重用途不传--reinitialize时作为权重初始化的来源同时通过BertConfig.from_pretrained(model_args.model_name_or_path)决定模型架构。配套参数--tokenizer_name允许指定与模型不同的 tokenizer若不传则回退使用--model_name_or_path对应的 tokenizer。4.4--wandb_user_name触发 Weights Biases 日志原文档说明传入你的用户名将触发 wandb 日志记录。源码在WandbArguments中定义了wandb_user_name: Optional[str] field( defaultNone, metadata{help: The WandB user name for potential logging. If left None, no logging}, ) wandb_project_name: Optional[str] field( defaultperformer-experiments, metadata{help: The WandB project name for potential logging}, )训练循环中每完成一个 batch 会记录Training loss每完成一个 epoch 记录Eval loss项目名默认为performer-experiments可用--wandb_project_name覆盖。只要不传--wandb_user_namewandb 完全不会被激活因此默认运行无需登录 wandb 账号。4.5--dataset_name与--dataset_config选择数据集原文档说明可通过这两个参数选择数据集并建议使用 Hub 的数据集查看器辅助定位。源码的DataTrainingArguments中dataset_name是数据集名称dataset_config_name是配置名如维基百科的语言/时间快照版本。加载逻辑为datasets load_dataset(data_args.dataset_name, data_args.dataset_config_name) if validation not in datasets.keys(): datasets[validation] load_dataset(..., splitftrain[:{data_args.validation_split_percentage}%]) datasets[train] load_dataset(..., splitftrain[{data_args.validation_split_percentage}%:])当数据集本身没有 validation 划分时脚本会自动按--validation_split_percentage默认 5从训练集切分出验证集。此外脚本同样支持本地数据文件传入--train_file/--validation_file支持 csv、json、txt 三种格式即可脱离 Hub 训练脚本会取名为text的列或第一列作为语料。4.6 更多可调参数从源码补充--mlm_probability掩码概率默认 0.15--max_seq_length最大序列长度默认取模型最大输入长度。这是 Performer 相对标准 BERT 优势最明显的场景——标准注意力在序列长度增长时显存呈平方增长而 Performer 可支撑更长的序列--pad_to_max_length是否将所有样本填充到max_seq_length默认 False动态按 batch 内最大长度填充--overwrite_cache/--preprocessing_num_workers控制数据集预处理缓存与并行进程数--use_fast_tokenizer是否使用 fast tokenizer默认 True其余训练参数batch size、学习率、warmup、epochs、seed 等全部继承自 Transformers 的TrainingArguments。五、源码级拆解训练脚本如何工作run_mlm_performer.py 完整实现了一个 Flax 版 MLM 训练器主流程如下参数解析用HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, WandbArguments))解析四组参数若命令行只传一个以.json结尾的参数则按 JSON 配置文件解析数据集准备如上节所述从 Hub 或本地文件加载并切分训练/验证集然后 tokenizereturn_special_tokens_maskTrue、截断、可选填充模型加载按--performer与--reinitialize选择模型类与初始化方式数据整理器FlaxDataCollatorForLanguageModeling负责动态 padding 与掩码生成掩码策略遵循经典 MLM 约定——80% 替换为[MASK]、10% 替换为随机词、10% 保持原词且只对被掩码的 token 计算 loss未掩码 token 的标签设为 -100优化器与学习率调度使用 Flax 的Adam优化器并实现了一个因子可组合的学习率调度器create_learning_rate_scheduler支持constant、linear_warmup、rsqrt_decay、rsqrt_normalized_decay、decay_every、cosine_decay六种因子默认组合为constant * linear_warmup * rsqrt_decay即学习率先线性热身再按步数平方根倒数衰减warmup 步数取max(warmup_steps, 1)并行训练jax.pmap将training_step/eval_step映射到所有本地设备jax_utils.replicate复制优化器参数训练时通过jax.lax.pmean做跨设备梯度平均——这意味着脚本天然支持多 GPU/TPU 数据并行指标记录每个 epoch 计算 loss 与 accuracy保存 TensorBoard 标量若可用可选同步到 wandb。六、Performer 模型实现注意力如何从 O(n²) 降到 O(n)6.1 模型结构BERT 骨架 快速注意力插槽modeling_flax_performer.py 完整复刻了 BERT 的组件层次FlaxPerformerEmbeddings词/位置/类型三路 embedding 求和 LayerNorm→FlaxPerformerEncoderN 层FlaxPerformerLayer→ MLM 头。与标准 BERT 的唯一结构性差异在注意力层class FlaxPerformerAttention(nn.Module): num_heads: int head_size: int nn.compact def __call__(self, hidden_state, attention_mask): single_head_dim self.head_size // self.num_heads fast_softmax_attention make_fast_softmax_attention(qkv_dimsingle_head_dim) self_att nn.attention.SelfAttention( num_headsself.num_heads, qkv_featuresself.head_size, nameself, attention_fnfast_softmax_attention )(hidden_state, attention_mask) layer_norm FlaxPerformerLayerNorm(namelayer_norm)(self_att hidden_state) return layer_norm关键在nn.attention.SelfAttention的attention_fn参数——它把 Flax 原生自注意力的点积注意力函数替换成了make_fast_softmax_attention返回的快速注意力函数从而让换注意力机制变成了一次函数注入其余 BERT 结构FFN、残差、LayerNorm、pooler完全复用。6.2 权重迁移PyTorch 检查点如何进入 Flax脚本允许从 PyTorch 版 BERT 检查点初始化 Flax 权重modeling_flax_performer.py 中的convert_from_pytorch静态方法负责这一映射处理了四大类差异全连接层PyTorch 的dense.weight→ Flax 的dense.kernel注意力头分解query/key/value的权重按num_attention_heads重塑并转置以匹配 FlaxSelfAttention的头分解存储层归一化LayerNorm.weight/bias→layer_norm.gamma/beta参数转置intermediate.dense.kernel、output.dense.kernel、pooler.dense.kernel等需要转置attention.output.dense与attention.output.LayerNorm则要消除一层嵌套。6.3 快速注意力核心算法随机特征映射与低秩分解modeling_flax_performer_utils.py 的头部注释说明该文件复制自 Google Research 的fast_self_attention.py核心思路是利用结构化随机特征映射RFM技术对注意力矩阵做低秩分解从而近似快速 softmax 注意力。make_fast_softmax_attention是构造入口关键可调参数包括nb_features256随机特征数量特征越多近似越精确、计算开销越大ortho_featuresTrue默认使用高斯正交随机矩阵GaussianOrthogonalRandomMatrix而非非结构化高斯矩阵GaussianUnstructuredRandomMatrix正交矩阵可降低近似方差nonnegative_featuresTrue默认使用非负 softmax 核特征nonnegative_softmax_kernel_feature_creator即用exp(投影 - 范数项) eps构造非负特征来近似 softmax若置 False 则改用 sin/cos 特征sincos_softmax_kernel_feature_creatorredraw_featuresTrue每个注意力调用根据 query 的和重新抽取投影矩阵保证注意力是 permutation equivariantrenormalize_attentionTrue对结果做重归一化以匹配 softmax 的归一化特性numerical_stabilizer0.000001数值稳定项。在FastAttentionviaLowRankDecomposition.dot_product_attention中算法将序列维度上的O(L²)注意力矩阵乘法重构为两个低秩步骤先算Z (K)ᵀVkey 特征与 value 的缩并再算W QZquery 特征与 Z 的缩并配合R Q(K)ᵀ1计算归一化因子。这样每个位置的计算量与序列长度 L 呈线性关系而非平方关系——这正是 Performer 支持超长序列、降低显存占用的根本原因。归一化后的输出还会经过jnp.reciprocal与数值稳定化处理保证训练稳定性。另外工具模块还提供了make_fast_generalized_attention可将 softmax 注意力推广到任意核函数如jax.nn.relu支持ortho、iid、deterministic三种特征类型为扩展实验留有余地。七、运行与验证建议先跑 sanity_script.shSimple Wikipedia 规模小在单卡上几分钟内即可完成 3 个 epoch用于验证环境jax/flax/datasets 版本兼容、模型加载、wandb 可选集成是否就绪对比基线去掉--performer跑一遍同一命令即可获得标准FlaxBertForMaskedLM的 loss/accuracy 曲线作为对照验证长序列能力增大--max_seq_length观察 Performer 在序列变长时的显存增长是否明显慢于标准注意力尝试从头训练加--reinitialize验证随机初始化下 Performer 的收敛行为监控指标不加--wandb_user_name时关注终端输出的逐 epochLoss/Acc若机器装有 TensorBoard指标会写入--output_dir/logs目录。需要说明的是本项目属于 research_projects 研究示例追求的是算法原理验证与快速迭代而非生产级训练框架在 FlexGen 仓库上下文中它适合作为评估高效注意力方案对后续推理吞吐影响的实验前置环节。八、总结这份 Performer 微调示例虽然目录不大却是一条完整的研究闭环sanity_script.sh/full_script.sh提供可直接运行的实验入口run_mlm_performer.py提供与标准 BERT 无缝切换的训练框架modeling_flax_performer.py用函数注入的方式将 FAVOR 快速注意力嵌入 BERT 骨架modeling_flax_performer_utils.py则落地了随机特征映射与低秩分解的核心数学。对希望在有限显存下训练长序列 Transformer 的开发者而言这份示例既是可复用的 Flax 训练脚手架也是理解高效注意力机制最直接的源码教材。赞分享推理引擎大模型【免费下载链接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.项目地址https://gitcode.com/gh_mirrors/fl/FlexGen点击查看免费下载相关推荐FlexGen 基准仓库内 JAX/Flax 流式预训练实战用 Dataset Streaming 从零训练 RoBERTa 掩码语言模型FlexGen 基准仓库内 JAX/Flax 流式预训练实战用 Dataset Streaming 从零训练 RoBERTa 掩码语言模型 本篇技术指南围绕推理引擎大模型FlexGen 仓库中的 Flax 分词分类微调实战基于 run_flax_ner.py 的 NER/POS 训练与评估指南FlexGen 仓库中的 Flax 分词分类微调实战基于 run_flax_ner.py 的 NER/POS 训练与评估指南 导读 本文围绕 benchmar推理引擎大模型FlexGen 仓库中的 Flax GLUE 文本分类微调实战指南基于 run_flax_glue.py 的序列分类训练全流程FlexGen 仓库中的 Flax GLUE 文本分类微调实战指南基于 run_flax_glue.py 的序列分类训练全流程 导读 本文围绕 benchma推理引擎大模型上一篇ThinkPad X230黑苹果完美教程轻松实现macOS体验下一篇BG3ModManager5步解决博德之门3模组管理难题创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考