ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

GPTBigCode 模型解析:Multi-Query Attention、FIM 填中与 Flash Attention 2 加速实践

GPTBigCode 模型解析:Multi-Query Attention、FIM 填中与 Flash Attention 2 加速实践 GPTBigCode 模型解析Multi-Query Attention、FIM 填中与 Flash Attention 2 加速实践【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformersGPTBigCode 是 Transformers 中专门为 BigCode 项目多语言代码模型如 SantaCoder、StarCoder设计的 GPT-2 变体实现。本文将以 docs/source/en/model_doc/gpt_bigcode.md 为主线结合仓库内 源码 与 配置实现系统讲解其与 GPT-2 的全部实现差异、核心配置参数以及如何加载预训练权重并结合 Flash Attention 2 / SDPA 进行推理加速。读完本文你将能读懂该模型的源码设计动机并独立完成代码生成与填充任务的部署与加速配置。GPTBigCode 是什么面向代码生成的 BigCode 系列实现GPTBigCode 模型由 BigCode 项目于 2023 年提出对应论文SantaCoder: dont reach for the stars!目标是负责任地开发大规模代码大语言模型。BigCode 团队使用来自 The Stack 数据集的 Java、JavaScript、Python 子集训练 1.1B 参数的模型并在 MultiPL-E 文本到代码基准上评测。研究还表明更激进的近似重复过滤能提升性能而从 GitHub 星标 5 的仓库中选择文件反而会显著损害性能。在 Transformers 仓库中该模型于 2023 年 4 月 10 日被合入实现 PR 对应 modeling_gpt_bigcode.py。从架构上看GPTBigCode 本质是一个经过深度优化的 GPT-2 模型核心增强点是支持Multi-Query AttentionMQA并集成了 Fill-in-the-MiddleFIM训练所依赖的因果掩码设计。仓库中的典型预训练权重包括bigcode/gpt_bigcode-santacoderSantaCoder 主权重以及bigcode/starcoderStarCoder 权重。模型在当前仓库中被完整注册到自动加载体系GPTBigCodeConfig通过 modeling_auto.py 中的映射接入AutoConfigGPTBigCodeModel、GPTBigCodeForCausalLM等类接入AutoModel/AutoModelForCausalLMtokenizer 自动映射见 tokenization_auto.py。因此你可以完全通过AutoModelForCausalLM.from_pretrained(...)使用而无需显式指定类名。实现细节GPTBigCode 相对 GPT-2 的优化清单官方文档列出 GPTBigCode 与 GPT-2 的差异全部可以在 modeling_gpt_bigcode.py 中得到源码级验证。理解这些差异是后续进行二次开发或排查推理行为的前提。1. 支持 Multi-Query AttentionMQA这是 GPTBigCode 最具代表性的架构改动。传统 MHAMulti-Head Attention下每个注意力头各有一份独立的 K/V 投影而 MQA 让所有 Query 头共享同一份 Key 和 Value。在 GPTBigCodeAttention 中kv_heads 1 if self.multi_query else self.num_heads见 modeling_gpt_bigcode.py即 KV 头数被压缩为 1。对应的 QKV 线性层c_attn输出维度也随之变为embed_dim 2 * kv_dim见 modeling_gpt_bigcode.py相比 MHA 的3 * embed_dim显著减少了 KV 缓存与计算量——这正是代码模型在长序列生成与批处理中提升吞吐的关键。代码中的 MQA 前向拆分为见 modeling_gpt_bigcode.pyif self.multi_query: query, key, value ( self.c_attn(hidden_states).unsqueeze(1).split((self.embed_dim, self.kv_dim, self.kv_dim), dim3) ) query query.view(*input_shape, -1, self.head_dim).transpose(1, 2)同时为了向统一的 attention 后端分发SDPA / Flash Attention 2 / eager代码通过repeat_kv见 modeling_gpt_bigcode.py在eager_attention_forward见 modeling_gpt_bigcode.py中将共享 KV 广播回num_heads份后再参与 matmul。测试佐证GPTBigCodeMQATest.test_mqa_reduces_to_mha通过把 MQA 权重扩展复制为 MHA 权重验证了两者在相同输入下输出逐元素一致rtol1e-5, atol1e-5证明 MQA 在数学上是 MHA 的精确特例见 tests/models/gpt_bigcode/test_modeling_gpt_bigcode.py。2. 使用gelu_pytorch_tanh而非经典 GELUGPTBigCode 默认激活函数是 GELU 的 tanh 近似实现gelu_pytorch_tanh而非 GPT-2 常用的精确 GELU。该默认值直接写死在 configuration_gpt_bigcode.py并由 MLP 层通过ACT2FN[config.activation_function]查表实例化见 modeling_gpt_bigcode.py。采用 tanh 近似通常带来微小的数值差异但在部分精度方案下更利于算子融合与训练稳定性。3. 避免不必要的同步文档指出参考实现中存在的若干同步点在此被消除该优化后续也已合并回 GPT-2对应 PR #20061。实践中这减少了内核启动间的主机-设备同步开销对短序列高频调用场景收益明显。4. 用nn.Linear替代Conv1DGPT-2 传统上用 1D 卷积Conv1D实现全连接投影GPTBigCode 直接改用标准nn.Linear如self.c_attn、self.c_proj、self.c_fc、self.c_proj见 modeling_gpt_bigcode.py 与 modeling_gpt_bigcode.py。线性层通常能获得更好的算子调度与显存布局但权重排布与 GPT-2 不同因此与原始openai-community/gpt2等 GPT-2 权重不兼容——加载时必须使用配套的 GPTBigCode 权重文件。5. 合并注意力分支、统一 softmax 路径并缓存 maskGPT-2 中_attn与_upcast_and_reordered_attn两条路径被合并matmul 与缩放常数始终合一原reorder_and_upcast_attn更名为attention_softmax_in_fp32配置位。缓存 attention mask 值避免每层重复构造且整个模型只在forward中通过create_causal_mask预计算一次因果掩码见 modeling_gpt_bigcode.py再逐层复用避免每层重复生成 mask。模型还持有一个常量下三角掩码 bufferself.biaspersistentFalse不进 checkpoint并在_init_weights中初始化见 modeling_gpt_bigcode.py。6. 用 TorchScript jit 融合 fp32 上转、掩码与 softmaxGPTBigCode 将注意力中“fp32 上转 → 掩码 → softmax → 缩放”整合成数个torch.jit.script函数源码中可见 modeling_gpt_bigcode.py 的upcast_masked_softmax、upcast_softmax与masked_softmax。注释也说明为可融合性将不同条件是否需要掩码、是否上转拆成了独立 kernel 分支避免把条件判断写进单个函数导致 jit 放弃融合。同时 eager 路径也坚持softmax(dim-1, dtypetorch.float32)保证数值稳定后再转回原精度见 modeling_gpt_bigcode.py。7. Key/Value 缓存合并与 MHA 张量布局调整解码器把 Key 与 Value 合并进同一缓存结构显著减少逐 token 生成时缓存管理的开销。在 MHA 模式下multi_queryFalseQKV 张量采用(num_heads, 3, head_dim)内存布局而非 GPT-2 的(3, num_heads, head_dim)见 modeling_gpt_bigcode.py。这样在 KV 合并后按头切分更高效但同样使 checkpoint 与原始 GPT-2 权重不兼容。8. 残差路径缩放初始化_init_weights遵循 OpenAI GPT-2 论文的残差网络初始化方案对GPTBigCodeMLP/GPTBigCodeAttention中的c_proj权重按std initializer_range / sqrt(2 * n_layer)重新初始化见 modeling_gpt_bigcode.py以抵消深层堆叠时残差路径上的方差累积。GPTBigCodeConfig核心配置参数与默认值GPTBigCodeConfig定义于 configuration_gpt_bigcode.pymodel_type gpt_bigcode。为兼容 GPT-2 习惯命名它通过attribute_map将 HF 通用名映射回 GPT 风格命名hidden_size→n_embd、max_position_embeddings→n_positions、num_attention_heads→n_head、num_hidden_layers→n_layer两类命名在使用中等价。默认值速览均来自 configuration_gpt_bigcode.py参数默认值含义vocab_size50257词表大小n_positions/max_position_embeddings1024最大位置编码长度n_embd/hidden_size768隐藏层维度n_layer/num_hidden_layers12Transformer 层数n_head/num_attention_heads12注意力头数n_innerNone自动取 4×hidden_sizeFFN 中间维度activation_functiongelu_pytorch_tanh激活函数resid_pdrop0.1残差路径 dropoutembd_pdrop0.1embedding dropoutattn_pdrop0.1注意力权重 dropoutlayer_norm_epsilon1e-5LayerNorm epsiloninitializer_range0.02初始化标准差scale_attn_weightsTrue注意力分数是否除以 sqrt(head_dim) 缩放attention_softmax_in_fp32True是否在 fp32 中执行融合 softmaxscale_attention_softmax_in_fp32True是否在 fp32 中完成 softmax 缩放multi_queryTrue是否启用 Multi-Query Attentionuse_cacheTrue生成时是否缓存 KVbos_token_id/eos_token_id50256BOS/EOS tokenpad_token_idNonepadding token默认未定义add_cross_attentionFalse是否加入交叉注意力层tie_word_embeddingsTrue是否共享输入/输出词嵌入两个值得单独说明的注意点multi_query与num_key_value_heads__post_init__中会根据multi_query自动派生num_key_value_heads 1 if self.multi_query else self.n_head见 configuration_gpt_bigcode.py。当multi_queryFalse时退化为标准 MHA。attention_softmax_in_fp32与scale_attention_softmax_in_fp32源码中scale_attention_softmax_in_fp32 config.scale_attention_softmax_in_fp32 and config.attention_softmax_in_fp32见 modeling_gpt_bigcode.py即“缩放是否在 fp32 中进行”以“整体上转 fp32”为前提二者存在联动关系。从配置初始化模型的官方示例from transformers import GPTBigCodeConfig, GPTBigCodeModel # 使用默认参数初始化一个随机权重配置 configuration GPTBigCodeConfig() # 从配置构建随机初始化模型 model GPTBigCodeModel(configuration) # 读取模型使用的配置 configuration model.config从配置构建模型用户通常直接from transformers import GPTBigCodeConfig, GPTBigCodeModel后从预训练权重加载。加载流程见 configuration_gpt_bigcode.py from transformers import GPTBigCodeConfig, GPTBigCodeModel # 初始化一个 GPTBigCode 配置 configuration GPTBigCodeConfig() # 从配置初始化模型随机权重 model GPTBigCodeModel(configuration) # 读取模型配置 configuration model.config将 StarCoder 与 Flash Attention 2 结合使用GPTBigCode 是首批在 Transformers 中系统性验证 Flash Attention 2 的架构之一模型基类声明了_supports_flash_attn True与_supports_sdpa True见 modeling_gpt_bigcode.py并支持_supports_attention_backend True可以接入统一的 attention 后端分发机制。官方文档给出的完整步骤如下。第一步安装 Flash Attention 2含滑动窗口特性pip install -U flash-attn --no-build-isolation注意确保硬件兼容 Flash-Attention 2目前主要面向 NVIDIA Ampere 及以上架构的 GPU详情参见 flash-attn 官方说明加载模型时应使用半精度如torch.float16。第二步加载并运行以 SantaCoder 为例from transformers import AutoModelForCausalLM, AutoTokenizer model AutoModelForCausalLM.from_pretrained(bigcode/gpt_bigcode-santacoder, device_mapauto, attn_implementationflash_attention_2) tokenizer AutoTokenizer.from_pretrained(bigcode/gpt_bigcode-santacoder) prompt def hello_world(): model_inputs tokenizer([prompt], return_tensorspt).to(model.device) model.to(model.device) generated_ids model.generate(**model_inputs, max_new_tokens30, do_sampleFalse) tokenizer.batch_decode(generated_ids)[0] # def hello_world():\n print(hello world)\n\nif __name__ __main__:\n print(hello world)\n|endoftext|官方文档同时给出以bigcode/starcodercheckpoint 在两种序列长度下对比原生实现与 Flash Attention 2 纯推理时间的加速图结论是注意力后端替换带来的端到端提速显著原图位于原文档 Expected speedups 小节。从源码看这种加速得益于两处设计缓存/掩码前置因果 mask 在模型级统一计算一次并贯穿所有层见 modeling_gpt_bigcode.py避免每层重复构建减轻了 Flash 后端逐层切换的开销后端抽象注意力通过ALL_ATTENTION_FUNCTIONS.get_interface(...)按config._attn_implementation分发见 modeling_gpt_bigcode.pyattn_implementationflash_attention_2时直接落到 Flash 融合 kernel无需修改任何模型逻辑。若你的环境不支持 Flash Attention 2也可以不传attn_implementation默认走 SDPA/eager模型行为保持一致——文档顶部也以徽章形式标注了该模型同时支持 FlashAttention 与 SDPA 两种实现。模型家族与任务适配官方文档为本模型族注册了四个入口类全部位于 src/transformers/models/gpt_bigcode 并被AutoAPI 覆盖。各类的forward文档与行为要点如下。GPTBigCodeModel裸 Transformer返回BaseModelOutputWithPastAndCrossAttentions。forward接受input_ids/inputs_embeds二选一见 modeling_gpt_bigcode.py、past_key_values首次调用时若use_cacheTrue会自动创建DynamicCache见 modeling_gpt_bigcode.py、attention_mask、token_type_ids、position_ids、encoder_hidden_states与encoder_attention_mask。模型由wtetoken 嵌入wpe位置嵌入drop 若干GPTBigCodeBlock 末尾ln_f组成见 modeling_gpt_bigcode.py并默认使用标准因果掩码兼容 FIM / 前缀 LM 训练中非零attention_mask的情形。GPTBigCodeForCausalLM因果语言建模在裸模型之上叠加无偏置的lm_head权重与输入嵌入transformer.wte绑定_tied_weights_keys见 modeling_gpt_bigcode.py。forward语义要点labels在模型内部自动移位即可以简单传labelsinput_ids标签中设为-100的位置被忽略loss 只作用于[0, ..., vocab_size]区间内的标签logits_to_keep控制只计算尾部 k 个 token 的 logits配合前缀缓存减少生成/评测开销见 modeling_gpt_bigcode.py。它混入GenerationMixin因此可以直接调用.generate(...)。GPTBigCodeForSequenceClassification序列分类与 GPT-1 等因果模型一致取最后一个有效 token的隐状态做分类。要点若配置了pad_token_id会定位每行最右侧的非 padding token否则直接取每行最后一位当未配置pad_token_id且 batch_size 1 时直接抛错防止静默得到错误结果见 modeling_gpt_bigcode.py传inputs_embeds时无法推断 padding 位置会给出 warning 并按最后一位处理loss 支持回归num_labels1MSE、单标签分类CrossEntropy与多标签分类BCEWithLogits三种problem_type自动判定见 modeling_gpt_bigcode.py。GPTBigCodeForTokenClassificationtoken 级分类在隐藏状态上先 dropout 再过classifier线性层输出逐 token 的 logits训练时使用CrossEntropyLoss在(batch*seq, num_labels)上计算见 modeling_gpt_bigcode.py适合代码 token 级任务如类型标注、token 质量打分。工程佐证仓库测试如何验证 GPTBigCode测试文件 提供了两类关键验证可作为二次开发时的回归参考慢速生成测试GPTBigCodeModelLanguageGenerationTest从bigcode/gpt_bigcode-santacoder加载真实权重验证代码补全输出例如 promptdef print_hello_world():应生成def print_hello_world():\n print(Hello World!)\n...见 test_modeling_gpt_bigcode.py批处理场景需设置pad_token eos_token且padding_side left见 test_modeling_gpt_bigcode.py另有tiny_starcoder_py上的换行回归测试用于防止注意力缩放问题导致输出刷屏空行见 test_modeling_gpt_bigcode.py。MQA 等价性测试GPTBigCodeMQATesttest_mqa_reduces_to_mha在 5 组随机种子 × 训练/推理两种模式下把 MQA 的共享 KV 权重 expand 成 MHA 权重断言二者输出一致rtol/atol1e-5见上文已引用的 test_modeling_gpt_bigcode.py。实践要点速查权重不通用GPTBigCode 的 checkpoint 与 GPT-2 权重不兼容nn.Linear布局、(num_heads, 3, head_dim)QKV 布局不同必须使用bigcode/*系列对应权重。推理加速优先级硬件允许 →attn_implementationflash_attention_2 半精度否则使用默认 SDPA 后端同样可获得不错的性能且无需改动任何调用代码。代码生成配置左填充 pad_tokeneos_token是批处理生成的推荐姿势FIM 场景下保持模型默认因果行为并正确传入注意力掩码即可。微调/评测labels自动移位直接传labelsinput_idslogits_to_keep可减少长前缀场景的冗余 logits 计算。GPTBigCode 的源码是理解“如何把一个成熟架构GPT-2针对特定场景代码生成做工程化取舍”的绝佳样本从 MQA 的显存/带宽权衡到 jit 融合与掩码缓存这些毫秒级优化再到对 Flash Attention 2 的完整适配——这些设计思路可以直接迁移到其他自回归代码模型的部署与优化工作中。【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表