本章涵盖:

  • 专门化开源 LLM 在代码生成中所扮演的角色
  • 优化和量化这些 LLM,以提升推理性能
  • 在 4-bit 量化后,在笔记本电脑上运行这些 LLM

在本章中,我们将把到目前为止介绍过的概念应用到真实世界的领域专用 LLM 上,包括优化和量化。我们会聚焦于为 Python 代码生成和编程辅助而调优的模型。虽然这一类开源模型大多支持多种语言,但我们将使用 Python,这样你可以更容易判断输出质量。

7.1 使用 Transformers 生成代码

除了闭源、专有的基于 LLM 的代码助手之外,也已经有几个为编码任务设计的流行开源 LLM 被发布。在本章中,我们会动手实验其中一些选项,看看它们的优缺点,并了解如何优化它们,以及在需要时如何量化它们,使其能够在普通硬件上运行推理,理想情况下是在我们的笔记本电脑上运行,同时保持良好性能。不过,首先我们先思考一下,在这个代码助手时代,人类程序员的现在与未来。

鉴于商业和开源 copilot 的可用性,许多 CEO 和其他高管都预测,程序员会被 prompt engineer 取代,或者声称每个人,无论背景如何,都会成为软件架构师。这有多真实?我们先从大型通用模型开始看。一些研究工作,例如 SWE-bench,会基于真实代码中的真实编程挑战,评估 GPT-4 和 Claude 2 这类大型商业模型。该论文聚焦 Python,包含来自 GitHub 仓库的 2,294 个 issue 及相关 pull request,但网上也可以找到许多其他语言的 benchmark。即便加入知识检索,这两个模型在这些问题上的得分也并不好。更多 prompt engineering 帮助有限:即使上下文更大,结果仍可能混乱或随机,而且会消耗更多计算能力和用户时间。自该论文发表以来,许多大小模型,无论是否使用检索增强生成,也就是 RAG,都已经被 benchmark。更新后的排行榜可以在 SWE-bench 网站上看到,每周一更新。随着时间推移,新模型和支持技术提升了结果,但所有被 benchmark 的模型仍然不能在没有专家人工监督的情况下使用,见图 7.1。

image.png

图 7.1 —— SWE-bench 排行榜,2025 年初快照

几个因素解释了这种有限表现。解决软件工程问题通常需要协调多个函数、类和文件之间的变更,需要与执行环境交互,还需要进行远超直接代码生成的推理。这些仍然是人类优于机器学习模型的领域。而且,由于这些模型是作为通用模型训练的,它们很少能匹配专业软件工程师在复杂编程任务中所具备的深度专业能力。

专门用于编程任务的模型,例如 Microsoft Copilot,通常比通用模型更准确,也能处理更复杂的工作。不过,对一些组织来说,它们可能成本较高,并且需要人工监督和 guardrail,以防止数据泄漏。它们也可能缺乏训练数据来源方面的透明度。这些工具会引发本书前面讨论过的同类隐私、知识产权和安全问题,尤其是当它们依赖大型闭源模型时。

幸运的是,专门化于这类任务的 LLM 已经可用,而专门化会显著降低模型大小。我们可以从行业越来越倾向于更小、更聚焦的代码 LLM 这一趋势中看到这一点:这些模型虽然只有几十亿参数,却能在评估 benchmark 上匹配,甚至有时超过更大的模型。这并不令人意外:如果主要任务是在一种或几种编程语言中写代码,那么大型通用模型中的大量无关知识都可以被移除。一个更小、更聚焦的数据池更容易保持相关性,训练成本更低,也更不容易无意中包含受版权保护或未经授权的数据。

回到主要问题:生成式 AI 会取代人类程序员吗?很可能不会,即使模型在编码任务上达到 98% 到 100% 的准确率。编码只是软件开发的一部分。软件开发还包括投入新事物、理解用户需求和需要、把想法转化为可构建的东西,以及打磨最终结果。这些步骤需要人类判断,不能被简化为 next-token generation。在我看来,即便生成式模型变得高度准确,软件工程师也不会转向主要负责审查和验证 AI 生成代码的角色。这些模型仍然可以提供巨大帮助,我也不会劝阻专业人士使用它们。但要用好它们,领域知识始终是必要的:一个优秀的软件工程师可以借助代码助手成为更好的软件工程师;但一个几乎没有软件工程知识的人不会。就 AI 助手提供支持而言,软件工程与其他领域并没有不同。

7.2 使用 Transformer 架构生成 Python 代码

我们开始动手。我们将使用一个开源预训练代码生成模型作为 baseline,看看会得到什么结果。

7.2.1 使用 CodeGen 生成 Python 代码

我们首先考虑的开源预训练模型是 Salesforce 的 CodeGen。它使用标准 Transformer 架构进行对话式程序合成:每个问题会通过多个交互步骤解决,每一步都包含来自用户的自然语言 specification,以及系统合成出的一个子程序。虽然现在已经有更好的开源代码模型,包括 CodeGen 2.0 和 2.5,但 CodeGen 对这类专门化 LLM 来说仍然是一个有用示例,因为它突出了一些架构选择,而这些选择可能需要在 Hugging Face Transformers、ONNX 及相关库的典型用法之外做额外工作。

从概念上看,CodeGen 可以与 OpenAI 的 Codex 相比较。在 Salesforce 发布自己的方案时,Codex 正是 GitHub Copilot 背后的模型。CodeGen 提供两个预训练模型家族:multi,训练用于生成六种语言的代码,包括 C、C++、Go、Java、JavaScript 和 Python;以及 mono,只训练用于生成 Python。每个家族包含四种规模:350M,也就是 3.5 亿参数;2B,也就是 27 亿参数;6B,也就是 61 亿参数;以及 16B,也就是 161 亿参数。为了展示 CodeGen 在发布时的规模,图 7.2 将其最大模型与当时可用于同类任务的其他开源或闭源 LLM 进行了比较。

image.png

图 7.2 —— CodeGen 发布时最流行代码助手 LLM 的规模比较

如本章前面所说,我们只聚焦 Python 代码生成,因此本节中所有对 CodeGen 模型的引用,都是指 mono 家族。

即便在优化之前,350M 预训练模型的推理速度也很快。在单块 NVIDIA T4 GPU 上的 benchmark 显示,生成少于 50 个 token 的请求平均耗时 0.8 秒,生成 50 到 250 个 token 的请求平均耗时 0.9 秒。不过,350M 模型并不适合生产环境;本章后面会讨论生成结果的质量。

2B 模型是一个更稳健的候选模型,更大的版本也是如此。这个模型家族不仅可以生成使用 Python 标准库的代码,也可以生成使用第三方开源库的代码。即便如此,350M 版本仍然很适合用来解释本章涵盖的概念。

本章第一个 Colab notebook 包含一个完整的端到端示例,用于优化 CodeGen 350M mono 的推理。我们会使用这个较小模型来说明,即使是几亿参数规模的模型,也仍然有推理优化空间,而且相同思路也适用于更大的 mono 模型。

所有 CodeGen checkpoint,包括 mono 和 multi,都可以通过 Hugging Face Hub 和 Transformers 库获得,因此我们可以像往常一样下载预训练模型和 tokenizer:

import torch
from transformers import AutoTokenizer, CodeGenForCausalLM

device = "cpu"
model_id = "Salesforce/codegen-350M-mono"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = CodeGenForCausalLM.from_pretrained(model_id).to(device)
model.eval()

这里唯一的区别是使用了 CodeGenForCausalLM,这是 Hugging Face Transformers 中用于 CodeGen 模型的 auto class。

这个示例不需要硬件加速。我们假设模型和数据都加载并运行在 CPU 上。小型 CodeGen mono 模型的工作方式如下:

prompt = "def hello_world():"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

generated_ids = model.generate(input_ids, max_length=12)
print(tokenizer.decode(generated_ids[0], 
                       skip_special_tokens=True,
                       pad_token_id=50256))

给定一个具有描述性名称的函数定义作为 prompt,它会生成如下函数体:

def hello_world():
    print("Hello World")

这是提示 CodeGen 模型的一种方式。对 350M 模型来说,这是唯一效果较好的方法,不过其他方式也是可能的。对于 2B 及更大的 CodeGen 模型,也可以使用纯自然语言 prompt,仅限英语,例如:

Create bar chart with matplotlib

它会生成类似如下内容:

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# create a DataFrame
df = pd.DataFrame({'sales': [10, 5, 6, 9, 7],
                  'sale_of_sales': [20, 10, 15, 25, 30]})

# plot the bar chart
df.plot(kind='bar', figsize=(10, 6), color='blue')

# Add title and labels
plt.title("Sales")
plt.xlabel("Month")
plt.ylabel("Sales")

# show the plot
plt.show()

你也可以使用单行注释,例如:

# Create empty tensor with PyTorch

这会产生类似如下结果:

tensor = torch.empty(1, dtype=torch.float32)

如果使用多行注释:

"""

    Create a basic web app with Dash

"""

你会得到类似如下结果:

import dash
import dash_core_components as dcc
import dash_html_components as html
import plotly.graph_objects as go

app = dash.Dash()

app.layout = html.Div([
    html.H1('Dash App'),
    html.Div([
        dcc.Dropdown(
            id='dropdown',
            options=[
                {'label': 'New York City', 'value': 'NYC'},
                {'label': 'Montreal', 'value': 'MTL'},
                {'label': 'San Francisco', 'value': 'SF'}
            ],
            value='MTL'
        ),
        dcc.Graph(id='graph', animate=True)
    ])
])

@app.callback(
    dash.dependencies.Output('graph', 'figure'),
    [dash.dependencies.Input('dropdown', 'value')])
def update_figure(selected_dropdown_value):
    print(f"Selected dropdown value: {selected_dropdown_value}")
    return {
        'data': [
            {'x': [1, 2, 3], 'y': [4, 1, 2], 'type': 'bar', 'name': 'SF'},
            {'x': [1, 2, 3], 'y': [2, 4, 5], 'type': 'bar', 'name': 
            ➥'Montreal'},
        ],
        'layout': {
            'title': 'Dash Data Visualization'
        }
    }

if __name__ == '__main__':
    app.run_server(debug=True)

你还可以仅根据函数定义,函数名可以任意,即使不太具描述性,再加上一个行内注释,让模型根据注释生成代码:

def increment_elements(int_list):

    """

        Return list with elements incremented by 1

    """

这会生成类似如下代码:

def increment_elements(int_list):
    """
        Return list with elements incremented by 1
    """
    for i in range(len(int_list)):
        int_list[i]+=1
    return int_list

这些模型还可以验证代码,并在输入数据上运行用户编写或模型生成的代码:

int_list=[1, 3, 5]

>>> print(len(int_list))

它会生成如下内容:

int_list=[1, 3, 5]

>>> print(len(int_list))

3

你也可以给定一些代码和一条注释作为 prompt,让模型实现单元测试:

def create_empty_array(n, m):

    arr = np.zeros((n, m))

    return arr

#Write unit test for create_empty_array()

模型会根据注释中的指导为代码生成单元测试:

self.assertEqual(create_empty_array(2,3), [[0,0,0],[0,0,0]])
self.assertEqual(create_empty_array(5,4), np.zeros((5,4)))
self.assertEqual(create_empty_array(6,7), np.zeros((6,7)))
self.assertEqual(create_empty_array(3,4), np.zeros((3,4)))

根据第 6 章学到的内容,你可能会预期可以使用 Optimum 将这个模型转换为 ONNX,然后进行量化。现在这是正确的,但直到 2024 年初,CodeGen 架构还没有被 Optimum 支持——它是最近才被加入的。我们先使用 Optimum 将 CodeGen 模型导出为 ONNX 并量化;稍后再看一种替代方法,用于 Optimum 不支持某个模型架构的情况。

安装支持 ONNX Runtime 的 Optimum 后,可以从 Hugging Face Hub 下载模型,并自动转换为 ONNX:

from optimum.onnxruntime import ORTModelForCausalLM

model_id = 'Salesforce/codegen-350M-mono'
model = ORTModelForCausalLM.from_pretrained(model_id,
                                            export=True,
                                            provider="CPUExecutionProvider"
)

我们使用 Optimum 的 ORTModelForCausalLM 类,因为 CodeGen 是 causal LM。我们还在下载时指定了 ONNX execution provider,本例中是 CPUExecutionProvider

然后可以把模型保存到磁盘:

from pathlib import Path

onnx_path = Path("onnx")
model.save_pretrained(onnx_path)

现在可以通过 Optimum API 对它进行动态量化,如第 6 章所解释:

from optimum.onnxruntime import ORTQuantizer
from optimum.onnxruntime.configuration import AutoQuantizationConfig

dynamic_quantizer = ORTQuantizer.from_pretrained(model)
dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False,
                                              per_channel=False)

model_quantized_path = dynamic_quantizer.quantize(
    save_dir=onnx_path,
    quantization_config=dqconfig,
)

ONNX 转换后的模型为 1.33 GB;8-bit 量化版本为 346.22 MB。

现在我们来 benchmark 执行性能,看看 ONNX 转换和量化是否改善了延迟和吞吐。我们会分别对三个模型运行几百轮推理,并比较结果。

为了保持代码清晰可读,我们会使用 Transformers pipeline 对每个待测试的 CodeGen 模型版本运行推理。下面是创建 vanilla 模型 pipeline 的代码:

from transformers import pipeline, GenerationConfig

generation_config = GenerationConfig(
    pad_token_id=50256,
    truncation=True,
    max_length=12
)
pipe = pipeline("text-generation", 
                model=model, 
                tokenizer=tokenizer,
                generation_config=generation_config
                )

我们还需要实现两个简单工具函数,用来在 benchmark 执行期间为三个模型版本收集指标。首先定义 track_infer_time 函数,它测量单次推理运行的持续时间,并会应用于完整 benchmark 循环:

from contextlib import contextmanager
from time import perf_counter

@contextmanager
def track_infer_time(time_buffer):
    start_time = perf_counter()
    yield
    end_time = perf_counter()

    time_buffer.append(end_time - start_time)

第二个工具是 BenchmarkInferenceResult 类,用于在 benchmark 中收集指标:

from dataclasses import dataclass

@dataclass
class BenchmarkInferenceResult:
    model_inference_time: [int]  
    optimized_model_path: str

vanilla CodeGen 模型的 benchmark 代码随后就很简单:

from tqdm import trange 

PROVIDERS = {
    ("cpu", "PyTorch CPU"),
}

results = {}

inference_runs = 200
for device, label in PROVIDERS:
    time_buffer = []
    for _ in trange(inference_runs, desc=f"Tracking 
    ➥inference time (PyTorch vanilla model)"):
      with track_infer_time(time_buffer):
        pipe(prompt)

    results[label] = BenchmarkInferenceResult(
        time_buffer, 
        None
    )

用于 benchmark 该 CodeGen 模型 ONNX 版本的代码与原始模型相同。唯一不同的是传给 pipeline 的模型,也就是 ONNX 转换版本,以及用于计算和显示推理指标的 provider tag:

PROVIDERS = {
    ("CPUExecutionProvider", "ONNX CPU"),
}

图 7.3 展示了两个模型版本在几百次运行中的平均推理时间。所有 benchmark 都在无硬件加速的免费 Google Colab VM 上运行,硬件为 Intel Xeon CPU,两个虚拟 CPU,13 GB RAM。

image.png

图 7.3 —— PyTorch 中 vanilla CodeGen 模型与其 ONNX 转换版本的 CPU 平均推理时间

在我们的硬件上,转换为 ONNX 将平均推理时间从约 1186 ms 降低到约 958 ms。图 7.4 在同一 benchmark 上比较两个模型版本,展示延迟,包括平均值和百分位,以及吞吐量。以 ONNX 格式运行推理降低了延迟,并提升了吞吐。

image.png

图 7.4 —— CodeGen 350M mono vanilla 模型与其 ONNX 转换版本的延迟和吞吐指标

图 7.5 展示了两个待测模型版本在 100 次运行中的推理持续时间箱线图;其中也显示了单次运行时间。

image.png

图 7.5 —— CodeGen 350M mono vanilla 模型及其 ONNX 转换版本在 100 次运行中的推理耗时

在这个箱线图中可以看到,转换为 ONNX 并应用相关优化后,标准差小于原始模型。迁移到 ONNX 后,不同运行之间的持续时间更加一致。

通过用 ONNX 量化模型重复 benchmark,也就是把量化模型替换到 pipeline 中并设置 provider tag,我们可以看到进一步改善:平均推理时间大约降低 60%,见图 7.6。

image.png

图 7.6 —— ONNX CodeGen 模型及其量化版本的 CPU 平均推理时间

图 7.7 在同一 benchmark 上比较两个 ONNX 模型版本的延迟,包括平均值和百分位,以及吞吐量。

image.png

图 7.7 —— CodeGen 350M mono ONNX 模型及其量化版本的延迟和吞吐指标

与 vanilla CodeGen 模型上的同一 benchmark 相比,使用量化 ONNX 模型运行推理降低了延迟,并提升了吞吐。

注意
本节中的所有图表都是使用 Python 开源 Plotly 库生成的。我选择 Plotly 不仅是因为与 Matplotlib 相比,它的可视化更优雅,还因为它内置交互性。CodeGen benchmark 以及本节可视化的完整源代码,都可以在配套 Colab notebook 中找到。

我们已经看到,即使对于小型语言模型,将模型导出为 ONNX 和 / 或 ONNX 量化版本,也可以显著改善延迟和吞吐。但我们如何评估优化和量化对生成代码质量的影响呢?7.2.3 节将探索用于评估 Python 代码生成语言模型输出的策略。

7.2.2 对 Optimum 不支持的模型使用 ONNX

有些 LLM 架构尚未被 Optimum 包支持,例如截至 2024 年初,CodeGen 模型仍未被支持。在这些情况下,可以使用 ONNX 量化的 mid-level 方法。先从 Hugging Face Hub 下载模型并保存到磁盘,然后按如下方式运行 Transformers 库中的 transformers.onnx 脚本:

!python -m transformers.onnx --feature "causal-lm" --framework pt 
➥-export_with_transformers --model=local-pt-checkpoint onnx/

--feature 选项设置模型类别,对于 CodeGen 模型来说是 causal-lm;其他模型请参考 Transformers 文档。--model 选项指定下载后的 checkpoint 保存目录。

转换完成后,可以使用 ONNX Runtime API 执行 8-bit 量化:

import os
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

def quantize_onnx_model(onnx_model_path, quantized_model_path):    
    onnx_opt_model = onnx.load(onnx_model_path)
    quantize_dynamic(onnx_model_path,
                     quantized_model_path,
                     weight_type=QuantType.QInt8)

quantize_onnx_model(onnx_model_path, quantized_model_path)

7.2.3 模型评估

有几个开源 Python 包,例如 Hugging Face 的 Evaluate 和 EleutherAI 的 LLM Evaluation Harness,分别在第 6 章和第 4 章讨论过,再加上公开数据集,都可以用于 LLM 评估。但生成 Python 或其他代码这样的非结构化输出,需要额外选择,才能让评估稳健且有意义。

本节中,我们将基于 7.2.1 节介绍的 CodeGen 350M mono 模型及其 ONNX 转换版本和 8-bit 量化变体,走一遍一种评估方法。你也可以将它应用到其他专门用于 Python 代码生成的 LLM 上。代码来自 Amazon Science 的 ReCode 论文,并在本章 Colab notebook 中共享。我选择这种方法,是因为 ReCode 通过测量不同 prompt 扰动下的最坏情况行为,定义了代码生成模型的鲁棒性指标。语言模型对 prompt 变化很敏感,因此这种方法提供了一种实用方式,用于评估生产使用中的代码质量。

我们将使用 OpenAI 的 HumanEval 评估数据集,该数据集在《Evaluating Large Language Models Trained on Code》中提出。HumanEval 也可以在 Hugging Face Hub 上获得。它包含 164 个 Python 编程问题,每个问题都有函数签名、docstring、函数体和若干单元测试。注释和 docstring 使用英语。数据集字段如下:

  • task_id——数据样本标识符
  • prompt——模型输入,包含函数头和 docstring
  • canonical_solution——该 prompt 对应问题的代码解法
  • test——用于验证生成代码正确性的测试函数
  • entry_point——测试入口点

你可以通过 Hugging Face Datasets API 下载该数据集,然后将其转换为 ReCode 评估期望的格式。或者,也可以从 ReCode GitHub 仓库下载一个即用 JSON 副本:

!mkdir -p datasets/nominal/
%cd ./datasets/nominal/
!wget https://raw.githubusercontent.com/amazon-
➥science/recode/refs/heads/main/datasets/nominal/HumanEval.jsonl
%cd ../..

你可以在有无 GPU 的情况下对这个测试集运行端到端评估,但建议使用硬件加速。在 NVIDIA T4 GPU 上,vanilla CodeGen 350M mono 模型对全部 164 个测试样本运行一次需要超过一小时,条件是将 prompt 追加到答案中,并设置较大的最大生成长度。因为完整代码很长,这里不列出;本节会讲解工作流,让你了解生成代码如何被评估。

从 Hugging Face Hub 下载 CodeGen 340M mono 模型及其 tokenizer 后,把 HumanEval JSON 样本拆分为单独 JSON 文件,每个文件包含一次生成运行的 prompt。对于每个样本,先使用 CodeGen tokenizer 对 prompt 分词。然后使用如下配置生成代码:

from transformers import GenerationConfig

generation_config = GenerationConfig(
    pad_token_id=50256,
    truncation=True,
    max_length=1000,
    max_context_length=1000,
    use_cache=True,
    return_dict_in_generate=True,
    output_scores=True
)   

为某个样本生成代码后,需要检查它是否是有效 Python。未通过验证的结果不会写入结果文件。对于这个检查,我们使用 Python 内置的 ast 包,将生成代码解析为抽象语法树,也就是 AST;这也是 Python 在编译为 bytecode 前所做的事情。

为了看看 AST 长什么样,考虑一个简单示例。假设生成模型的 prompt 是:

def hello_world():

而该函数生成的函数体是:

print('Hello world')

你可以通过 ast 类的方法获得后者的 AST。在 Colab 代码中,它是从 is_valid_python 中调用的:

import ast

def is_valid_python(code):
    try:
        try:
            pared_code = ast.parse(code)
        except SyntaxError:
            return False
    except Exception as e:
        print("Exception: ", e)
        return False
    return pared_code

生成代码返回的 AST 看起来如下:

<ast.Module object at 0x000001458A1D3880>
{'body': [<ast.Expr object at 0x000001458A22FA60>], 'type_ignores': []}
children: [<ast.Expr object at 0x000001458A22FA60>]\n

<ast.Expr object at 0x000001458A22FA60>
{'value': <ast.Call object at 0x000001458A22FA30>, 'lineno': 1,
 'col_offset': 0, 'end_lineno': 1, 'end_col_offset': 20}
children: [<ast.Call object at 0x000001458A22FA30>]\n

<ast.Call object at 0x000001458A22FA30>
{'func': <ast.Name object at 0x000001458A22FA00>, 'args': [<ast.Constant
 object at 0x000001458A22F9D0>], 'keywords': [], 'lineno': 1, 
'col_offset': 0, 'end_lineno': 1, 'end_col_offset': 20}
children: [<ast.Name object at 0x000001458A22FA00>, <ast.Constant 
object at 0x000001458A22F9D0>]\n

<ast.Name object at 0x000001458A22FA00>
{'id': 'print', 'ctx': <ast.Load object at 0x0000014589ABBC70>,
 'lineno': 1, 'col_offset': 0, 'end_lineno': 1, 'end_col_offset': 5}
children: [<ast.Load object at 0x0000014589ABBC70>]\n

<ast.Constant object at 0x000001458A22F9D0>
{'value': 'hello world', 'kind': None, 'lineno': 1, 'col_offset': 6, 'end_lineno': 1, 'end_col_offset': 19}
children: []\n

<ast.Load object at 0x0000014589ABBC70>
{}
children: []\n

在这个例子中,树中有六个节点。生成 AST 是 Python 编译过程中的标准步骤。当你运行一个 .py 文件时,Python 首先把源码转换为 AST,然后转换为 bytecode,也就是 .pyc 文件。随后 Python 虚拟机会解释该 bytecode。AST 也用于其他用途;例如,流行 Python IDE 使用它们解析文件并理解代码结构,从而提供编辑器功能。我们也可以用同样方式,验证 LLM 生成的代码。

使用 AST 验证生成的 Python 代码后,我们会把结果收集到一个 JSON 文件中,以简化评估。最后一步依赖 HumanEval 包,因此我们从源码安装它,因为它无法通过 Python 包管理器安装:

!git clone https://github.com/openai/human-eval 
%cd human-eval 
!pip install .
%cd ..

对于每个生成样本,我们通过执行代码并运行其 HumanEval 测试来评估正确性。作为防止恶意代码执行的安全措施,HumanEval 在 execution.py 中注释掉了运行生成代码的那一行。在我们的 Colab notebook 中,我们覆盖了部分代码,以启用对模型生成代码的评估。运行这些评估时要谨慎。

我们会聚合结果,并将其与生成代码一起保存到 JSON 文件中。我们还会添加一个 passed attribute,值为 truefalse

虽然 Colab notebook 只覆盖 vanilla CodeGen 350M mono 模型,但你可以对其 ONNX 导出版本、8-bit 量化版本,以及其他用于 Python 代码生成的开源小型语言模型,重复相同评估流程和代码。

7.2.4 使用更好的模型生成 Python 代码

Salesforce 发布 CodeGen 2.0 版本时,再次提供了 mono 和 multi 两个家族,每个家族都有四种规模,大致匹配第一代规模。

最新一代是 CodeGen 2.5,它只有单一规模——mono 和 multi 版本都是 70 亿参数。它具备稳健的 infill sampling,也就是可以读取当前位置左侧和右侧的文本,并且通过 FlashAttention 优化了快速采样,使其适合高效 serving 和本地部署,包括在个人机器上部署。FlashAttention 是传统 self-attention 的替代方案,可以让 LLM 更高效,实现更快训练和推理。它不是反复从 GPU HBM,也就是大但相对慢的内存中读取和写入 keys、queries 和 values,而是将它们一次性加载到更小但更快的 SRAM 中,融合 attention 操作,然后再写回结果。

所有 CodeGen 1 和 2 模型都以宽松的 Apache 2.0 许可证发布,CodeGen 2.5 multi 和 mono 也使用相同许可证。注意,CodeGen 2.5 instruct 模型仅供研究使用,不可商用。

使用 CodeGen 2.5 mono 模型生成 Python 代码的方式,与本章前面讨论过的 CodeGen 1 mono 模型相同。首先,从 Hugging Face Hub 下载模型和 tokenizer:

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Salesforce/codegen25-7b-mono"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained model_id)

然后可以对分词后的 prompt 调用 generate 方法来生成代码:

text = "def hello_world():"
input_ids = tokenizer(text, return_tensors="pt").input_ids
generated_ids = model.generate(input_ids, max_length=128)
print(tokenizer.decode(generated_ids[0], skip_special_tokens=True)

如前所述,CodeGen 2.5 也很适合 code infilling,即填补不完整代码片段中的缺失部分。它要求 prompt 遵循如下格式约定:

  1. 在光标位置插入 <mask_1> token。
  2. 添加 <sep> token 标记边界。
  3. 再插入一个 <mask_1> token,指定要填补哪个 mask。

示例如下:

def format(prefix, suffix):
  return prefix + "<mask_1>" + suffix + "<|endoftext|>"
  ➥+ "<sep>" + "<mask_1>"

然后生成 Python 函数体,如下例所示:

prefix = "def hello_world():\n    "
suffix = "    return greeting_message" 

需要用自定义 format 函数格式化 prompt,对其分词,然后调用模型的 generate 方法:

text = format(prefix, suffix)
input_ids = tokenizer(text, return_tensors="pt").input_ids
generated_ids = model.generate(input_ids, max_length=128)
print(tokenizer.decode(generated_ids[0],
➥ skip_special_tokens=False)[len(text):])

还有其他用于代码生成的小型语言模型被提出。一个很强的例子是 StarCoder2,这是一个由 BigCode 开发的 30 亿参数模型。BigCode 是一个开放科学协作项目,专注于负责任地开发和使用面向代码的大型语言模型,并由 ServiceNow、Hugging Face 和 NVIDIA 支持。StarCoder2 在 2024 年论文中提出,基于 Stack v2 数据集中超过 3 万亿 token 训练而成,覆盖 17 种编程语言,包括 Python,以及来自 Wikipedia、arXiv 和 GitHub issue 的经过整理的自然语言文本。StarCoder2 使用 grouped-query attention 来加速推理,并具备 16,384-token 上下文窗口和 4,096-token sliding-window attention。由于它没有进行 instruction tuning,所以最适合在 prompt 中指定函数头之后生成函数体,不过官方 GitHub 仓库包含调优脚本。

你可以像往常一样使用 Hugging Face 包生成 StarCoder2 代码。首先从 Hub 下载模型和 tokenizer:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "bigcode/starcoder2-3b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id,
                                             device_map='auto',
                                             torch_dtype=torch.bfloat16)
model.eval()

然后可以对分词后的 prompt 调用 AutoModelForCausalLM 类的 generate 方法生成代码:

prompt = "def fibonacci(n):"
inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
outputs = model.generate(inputs)
print(tokenizer.decode(outputs[0]))

Transformers pipeline 也受支持:

from transformers import pipeline, GenerationConfig

generation_config = GenerationConfig(
    do_sample=True,
    use_cache=True,
    temperature=0.2,
    top_p=0.95,
    max_length=12
)
pipe = pipeline("text-generation",
                model=model,
                tokenizer=tokenizer,
                generation_config=generation_config
                )

result = pipe(prompt)
print(result[0]['generated_text'])

你可以使用 Hugging Face 的 Bitsandbytes 库创建该模型的 8-bit 量化版本:

from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

model_id = "bigcode/starcoder2-3b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
quantized_model = AutoModelForCausalLM.from_pretrained(model_id,

quantization_config=quantization_config)
quantized_model.eval()

也可以用相同方式创建 4-bit 量化版本,只需要将第二行替换为:

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

本章配套 Colab notebook benchmark 了 StarCoder v2 3B 模型 BF16 和 8-bit 量化版本的推理性能。benchmark 代码与前面展示的 CodeGen 示例类似,因此这里不重复。当你在 GPU 上运行 benchmark 时,会看到得益于 grouped-query attention 实现,该模型提供了比 CodeGen 350M 和 2B mono 模型更好的吞吐和延迟。

7.3 在普通硬件上进行编码辅助

7.2 节讨论的 CodeGen 和 StarCoder 模型展示了良好的推理性能,并且仍有进一步优化空间,使它们适合部署在 API 后面,为客户端应用提供代码生成能力。如果能让其中一个模型在笔记本电脑上作为代码助手平稳运行,也会很有用。本节中,我们将验证一个优化后的小型语言模型能否在本地作为代码助手运行,并达到可接受性能。我们的目标笔记本电脑是一台搭载 Apple M1 芯片、拥有 8 GB 统一内存的 MacBook Air。我选择了一台性能相对较弱的 Apple M 系列笔记本,以增加挑战,也展示即使在这种硬件上,你仍然可以优化推理并平稳运行有用的领域专用 SLM。我们的参考模型将是 StarCoder v2 3B。

如果浏览 Hugging Face Hub,你会发现几个 LLM,包括 StarCoder v2 3B,已经以 GGUF 格式可用。尝试这些量化模型进行笔记本推理的一个自然方式,是像使用标准 FP32 或 FP16 模型一样使用 Transformers API:

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "second-state/StarCoder2-3B-GGUF"
filename = "starcoder2-3b-Q2_K.gguf"

tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
model = AutoModelForCausalLM.from_pretrained(model_id, gguf_file=filename)

这里唯一不同的是要下载的量化模型 checkpoint 文件名。GGUF 提供多种量化选项,因此同一个模型可能有多个量化版本可用。

虽然 Hugging Face Hub 可以托管 GGUF 模型,Transformers 库也可以下载它们,但 Transformers 仍然不能直接使用量化 GGUF 格式的模型。GGUF 下载完成后,模型会自动反量化,推理只能以原始格式运行。这会让我们回到起点,因为原始模型通常太大,无法在 CPU/GPU 和内存受限的本地环境中运行。如果想通过 Transformers API,唯一变通办法是在下载后把模型保存到磁盘:

model_checkpoint_dir = 'starcoderv2-checkpoint-dir'
tokenizer.save_pretrained(model_checkpoint_dir)
model.save_pretrained(model_checkpoint_dir)

然后我们再对它应用一种量化技术。

但还有更聪明的选项。llama.cpp 是一个开源 C/C++ 库和项目,可以在多种硬件上为多个 LLM 启用推理,包括没有 GPU 的设备,例如笔记本电脑和手机。GGUF 是 llama.cpp 使用的二进制文件格式。

为了保持简单,我们会使用 llama-cpp-python,它为 llama.cpp 提供 Python binding,并支持 Apple 的 Metal,也就是 MPS,Metal Performance Shaders backend。安装它需要 Python 3.8+,以及适用于你机器的 C 编译器,例如 macOS 上的 Xcode。创建虚拟环境后,可以从 PyPI 安装 llama-cpp-python

CMAKE_ARGS="-DGGML_METAL=on" pip install llama-cpp-python

前面的命令会构建并安装该库,同时也会构建和安装 llama.cpp。你还需要 huggingface-hub 库来下载模型:

!pip install huggingface-hub

现在一切准备就绪,可以开始编写代码,使模型可用于本地推理。首先导入所需依赖:

from huggingface_hub import hf_hub_download
from llama_cpp import Llama

然后从 Hugging Face Hub 下载模型:

model_name = "second-state/StarCoder2-3B-GGUF "
model_file = "starcoder2-3b-Q4_K_M.gguf "
model_path = hf_hub_download(model_name, filename=model_file)

本示例中使用的 GGUF 转换模型,是 Second State 的 StarCoder v2 3B 的 4-bit K_M 量化版本。下载到笔记本电脑后,可以使用 llama-cpp-python 库的 Llama 类将其加载到内存中,如下所示:

llm = Llama(
    model_path=model_path,
    n_ctx=3072,
    n_threads=4
    n_gpu_layers=0
)

这里指定了模型 checkpoint 的本地路径、上下文大小,StarCoder v2 3B 为 3,072,要使用的 CPU core 数量,以及要卸载到 GPU 的模型层数。本例中不卸载任何层,因为我们希望只在 CPU 上运行推理。

接下来,可以设置代码生成参数,例如最大新 token 数、stop policy、是否在输出中 echo prompt、解码策略,以及其他推理选项:

generation_kwargs = {
    "max_tokens":20,
    "stop":["</s>"],
    "echo":False,
    "top_k":1
}

现在可以设置 prompt 并运行推理:

prompt = "def hello_world():"
res = llm(prompt, **generation_kwargs)

前面的命令会输出一个 OpenAI 兼容格式的字典:

{
  "id": "cmpl-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx",
  "object": "text_completion",
  "created": xxxxxxxxxx,
  "model": "./models/starcoder2-3b-Q4_K_M.gguf",
  "choices": [
    {
      "text": "\n    return 'Hello World!\n\n",
      "index": 0,
      "logprobs": None,
      "finish_reason": "length"
    }
  ],
  "usage": {
    "prompt_tokens": 5,
    "completion_tokens": 20,
    "total_tokens": 25
  }
}

如果只想返回生成文本并省略其他响应字段,可以使用如下命令:

print(res["choices"][0]["text"])

第一次运行这段代码会更慢,因为它必须把模型 checkpoint 下载到本地。第一次之后,就可以完全离线生成代码。

你可以用和 CodeGen 模型相同的方式 benchmark 这个本地模型的性能。使用较小模型时相同配置重复推理测试,会看到类似如下延迟和吞吐结果:

             GGUF 4-bit K_M
Average_latency (ms) 432.545684
Latency_P50          432.409188
Latency_P75          436.077177
Latency_P90          438.191962
Latency_P95          438.622171
Latency_P99          440.215217
Throughput               2.311895

整体性能只使用四个线程且不使用 GPU,与 4-bit 量化 CodeGen 模型相当,对于本地使用来说是可以接受的。你可以按照 7.2.3 节解释的方法评估生成代码。

到目前为止,我们都是通过执行带 prompt 的 Python 脚本,从 CLI 运行推理。你也可以在 chat mode 下运行,llama-cpp-python 为此提供了 API。该库还包含一个 Web server,因此你可以把 llama.cpp 兼容模型与任何 OpenAI 兼容客户端一起使用。这种模式让本地模型很容易与 IDE 集成,用于代码补全和建议。

在本节中,我们假设所需代码生成模型已经在 Hugging Face Hub 上以 GGUF 格式可用。但如果你想要的模型只有单精度或半精度版本,或者虽然有 GGUF 版本但性能不佳,该怎么办?在这些情况下,可以先把源模型 checkpoint 和支持文件下载到本地,然后使用 llama.cpp GitHub 仓库中的 convert_hf_to_gguf.py 脚本,将模型量化并转换为 GGUF:

python llama.cpp/convert-hf-to-gguf.py <checkpoints_dir> --outfile
➥<output_file_name>.gguf --outtype <quantization_type>

该脚本需要下载后的模型 checkpoint 和文件所在目录、目标 GGUF 文件名,以及量化类型,具体选项请参考 llama.cpp 文档。你也可以把生成的 GGUF 模型上传到 Hugging Face Hub。量化之后,重复本节前面的步骤即可在本地生成代码。

总结

  • 专门化开源代码生成模型可以用更少参数匹配通用模型性能。
  • CodeGen 模型支持多种编程语言,也通过 mono 家族支持仅 Python 的代码生成。
  • 与标准 PyTorch 实现相比,将代码生成模型转换为 ONNX 可以将平均推理时间降低约 19%。
  • ONNX 模型的 8-bit 量化可以在不牺牲质量的情况下,进一步减少 60% 推理时间。
  • Optimum 不支持的模型需要使用 transformers.onnx 包进行 mid-level 转换,并手动配置。
  • 抽象语法树,也就是 AST,解析可以在执行前检查生成的 Python 代码是否语法正确。
  • HumanEval 数据集包含 164 个 Python 编程问题,用于评估代码生成模型。
  • 代码生成模型支持函数签名、自然语言注释和多行 docstring 等 prompt。
  • StarCoder2 使用 grouped-query attention 加速推理,同时降低计算开销。
  • llama-cpp-python 提供 Python binding,可以在本地无 GPU 情况下运行 GGUF 模型。
  • Code infilling 需要 mask 和 separator token,使模型理解上下文。
  • 4-bit 量化可以将模型大小减少约 75%,同时保持可接受的代码生成质量。
  • 本地部署代码生成模型可以在不调用外部 API 的情况下提供离线编码帮助。
  • Apple 的 Metal Performance Shaders,也就是 MPS backend,可以优化 Apple Silicon 处理器上的推理。
  • 评估代码生成需要同时进行语法验证和用于功能正确性的单元测试。
  • ReCode 评估框架为不同 prompt 扰动下的代码生成提供鲁棒性指标。