使用基于自研芯片的 Amazon EC2 Inf2 实例构建高效经济的大语言模型 GPT-J-6B

在 re:Invent 2019上，亚马逊云科技发布了 Inferentia 芯片和 Inf1 实例。Inferentia 是一种高性能[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)推理芯片，由亚马逊云科技定制设计，其目的是提供具有成本效益的大规模低延迟预测。时隔四年，2023年4月亚马逊云科技发布了 Inferentia2 芯片和 Inf2 实例，旨在为大型模型推理提供技术保障。 Inf2 实例提供高达 2.3 petaflops 的 DL 性能和高达 384 GB 的总加速器内存以及 9.8 TB/s 的带宽。Amazon Neuron SDK 与 PyTorch 和 TensorFlow 等流行的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)框架原生集成。因此，用户可以继续使用现有框架和应用程序代码在 Inf2 上进行部署。开发人员可以在 Amazon Deep Learning AMI、Amazon Deep Learning 容器或 [Amazon Elastic Container Service](https://aws.amazon.com/cn/ecs/?trk=cndc-detail) ([Amazon ECS](https://aws.amazon.com/cn/ecs/?trk=cndc-detail))、[Amazon Elastic Kubernetes Service](https://aws.amazon.com/cn/eks/?trk=cndc-detail) ([Amazon EKS](https://aws.amazon.com/cn/eks/?trk=cndc-detail)) 和 [Amazon SageMaker](https://aws.amazon.com/cn/sagemaker/?trk=cndc-detail) 等托管服务中使用 Inf2 实例。  [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)Inf2 实例的核心是 Amazon Inferentia2 设备，每个设备包含两个 NeuronCores-v2。每个 NeuronCore-v2 都是一个独立的异构计算单元，具有四个主要引擎：张量（Tensor）、向量（Vector）、标量（Scalar） 和 GPSIMD 引擎。张量引擎针对矩阵运算进行了优化。标量引擎针对 ReLU（整流线性单元）函数等元素运算进行了优化。向量引擎针对非元素向量操作进行了优化，包括批量归一化或池化。下图显示了 Amazon Inferentia2 设备架构的内部工作原理。  Amazon Inferentia2 支持多种数据类型，包括 FP32、TF32、BF16、FP16 和 UINT8，因此用户可以根据工作负载选择最合适的数据类型。它还支持新的可配置 FP8 (cFP8) 数据类型，这与大型模型特别相关，因为它减少了模型的内存占用和 I/O 要求。 Amazon Inferentia2 嵌入了支持动态执行的通用数字信号处理器（DSP），因此无需在主机上展开或执行控制流运算符。Amazon  Inferentia2 还支持动态输入形状，这对于输入张量大小未知的模型（例如处理文本的模型）来说非常关键。 Amazon  Inferentia2 支持用 C++ 编写的自定义运算符。Neuron Custom C++ Operators 使用户能够编写在 NeuronCores 上本机运行的 C++ 自定义运算符。使用标准 PyTorch 自定义运算符编程接口将 CPU 自定义运算符迁移到 Neuron 并实现新的实验运算符，所有这些都不需要对 NeuronCore 硬件有深入了解。 Inf2 实例是 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)上的第一个推理优化实例，可通过芯片之间的直接超高速连接（NeuronLink v2）支持分布式推理。NeuronLink v2 使用集体通信 （Collective Communications）运算符（例如 all-reduce）在所有芯片上运行高性能推理管道。 # **Neuron SDK** Amazon Neuron 是一种 SDK，可优化在 Amazon Inferentia 和 Trainium 上执行的复杂神经网络模型的性能。Amazon Neuron 包括深度学习编译器、运行时和工具，这些工具与 TensorFlow 和 PyTorch 等流行框架原生集成，它预装在 Amazon Deep Learning AMI 和 Deep Learning Containers 中，供客户快速开始运行高性能且经济高效的推理。 Neuron 编译器接受多种格式（TensorFlow、PyTorch、XLA HLO）的[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)模型，并优化它们以在 Neuron 设备上的运行。Neuron 编译器在[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)框架内调用，其中模型由 Neuron Framework 插件发送到编译器。生成的编译器工件称为 NEFF 文件（Neuron 可执行文件格式），该文件又由 Neuron 运行时加载到 Neuron 设备。 Neuron 运行时由内核驱动程序和 C/C++ 库组成，后者提供 API 来访问 Inferentia 和 Trainium Neuron 设备。TensorFlow 和 PyTorch 的 Neuron ML 框架插件使用 Neuron 运行时在 NeuronCores 上加载和运行模型。Neuron 运行时将编译的深度学习模型（也称为 Neuron 可执行文件格式（NEFF））加载到 Neuron 设备，并针对高吞吐量和低延迟进行了优化。  ### **Inf2 实例的应用场景** 使用 Inf2 实例运行流行的应用程序，例如文本摘要、代码生成、视频和图像生成、语音识别、个性化等。Inf2 实例是 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)中的第一个推理优化实例，引入了由 NeuronLink (一种高速、非阻塞互连) 支持的横向扩展分布式推理。用户现在可以在 Inf2 实例上跨多个加速器高效部署具有数千亿个参数的模型。Inf2 实例的吞吐量比其他类似的 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)实例高出三倍，延迟低八倍，性价比高出40%。为了实现可持续发展目标，与其他类似的 [Amazon EC2 ](https://aws.amazon.com/cn/ec2/?trk=cndc-detail)实例相比，Inf2 实例的每瓦性能提高了50%。 ### **使用 Inf2 实例运行 GPT-J-6B 模型** GPT-J-6B 是由一组名为 EleutherAI 的研究人员创建的开源自回归语言模型。它是 OpenAI 的 GPT-3 最先进的替代方案之一，在聊天、摘要和问答等广泛的自然语言任务中表现良好。 该模型由28层组成，模型维度为4096，前馈维度为16384。模型维度分为16个头，每个头的维度为256。旋转位置嵌入（RoPE）应用于每个头的64个维度。使用与 GPT-2/GPT-3 相同的一组 BPE，使用50257的标记化词汇训练模型。  GPT-J-6B 有60亿个参数，非常适合大语言模型（LLM）学习的入门版本，进行文本生成测试。下文我们就以 GPT-J-6B 为例在 Inf2 实例上进行部署。在部署过程中，我们用到了 Neuron SDK 和 transformers-neuronx。transformers-neuronx 是由 Amazon Neuron 团队构建的开源库，可帮助使用 Amazon Neuron SDK 运行转换器解码器推理工作流程。目前，它提供了 GPT2、GPT-J 和 OPT 模型类型的演示脚本，它们的前向函数在编译过程中重新实现，以进行代码分析和优化。客户可以基于同一个库实现其他模型架构。Amazon Neuron 优化的转换器解码器类已使用称为 PyHLO 的语法在 XLA HLO (高级操作) 中重新实现。该库还实现了张量并行（Tensor Parallelism），以跨多个 NeuronCore 对模型权重进行分片。 部署流程如下： ##### **1. 准备创建 Inf2 实例** 进入到 Amazon 控制台，区域选择 us-east-1 或者 us-west-2，实例类型建议选择 8xlarge，配置为 CPU:32 核、内存:128G、GPU:32G、存储:500G，操作系统选择 Amazon Linux 2（作者使用的是 us-east-1 的 AMI：ami-0aa7d40eeae50c9a9）。Inf2 实例也可以使用 Amazon Deep Learning AMI 进行部署，好处是开箱即用，AMI 中内置了 Inferentia 的驱动程序和 Neuron SDK。为了让读者更加深入的了解 Inf2 的工作原理，我们并没有采用 Amazon Deep Learning AMI 。当 Inf2 实例创建完成后，通过 SSH 客户端登陆到 Inf2 实例。 ##### **2. 安装驱动和 Neuron 工具** ```js # Configure Linux for Neuron repository updates sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF [neuron] name=Neuron YUM Repository baseurl=https://yum.repos.neuron.amazonaws.com enabled=1 metadata_expire=0 EOF sudo rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB # Install OS headers sudo yum install kernel-devel-\$(uname -r) kernel-headers-\$(uname -r) -y # Update OS packages sudo yum update -y # Install Neuron Driver and tools sudo yum install aws-neuronx-dkms sudo yum install aws-neuronx-tools # Add PATH export PATH=/opt/aws/neuron/bin:\$PATH ``` ##### **3. 安装 PyTorch Neuron (torch-neuronx)** ```js # Install Neuron Runtime sudo yum install aws-neuronx-collectives-2.* -y sudo yum install aws-neuronx-runtime-lib-2.* -y # Install Python venv sudo yum install -y python3.7-venv gcc-c++ # Create Python venv python3.7 -m venv aws_neuron_venv_pytorch # Activate Python venv source aws_neuron_venv_pytorch/bin/activate python -m pip install -U pip # Install Jupyter notebook kernel pip install ipykernel python3.7 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python (torch-neuronx)" pip install jupyter notebook pip install environment_kernels # Set pip repository pointing to the Neuron repository python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com # Install wget, awscli python -m pip install wget python -m pip install awscli # Install Neuron Compiler and Framework python -m pip install neuronx-cc==2.* torch-neuronx torchvision ``` 您也可以参考文档： https\://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/setup/pytorch-install.html#pytorch-neuronx-install?trk=cndc-detail 进行安装 ##### **4. 更新 transformers 并安装 transformers-neuronx** ```js pip install --upgrade pip transformers pip install "git+https://github.com/aws-neuron/transformers-neuronx.git" ``` ##### **5. 下载 GPT-J-6B 模型** `gptj_demo --model_name="EleutherAI/gpt-j-6B" save gpt-j-6B-split` ##### **6. 编译并运行 GPT-J-6B 演示程序** `gptj_demo run gpt-j-6B-split` ##### **7. 观察输出和运行结果**  ##### **8. gptj_demo 接受如下参数，读者可以根据需要自行调整**  例如运行下列命令就是，序列长度为100，并设置输入批量大小为1。 `gptj_demo run gpt-j-6B-split --n_positions 100 --batch_size 1` ##### **9. 修改模型运行代码，使用循环接受连续输入** ```js cd /home/ubuntu/aws_neuron_venv_pytorch/lib/python3.8/site-packages/transformers_neuronx vi gpt_demo.py ``` 注释到99-108行 ```js ''' with torch.inference_mode(): encoded_text = tokenizer.encode(prompt_text) input_ids = torch.as_tensor([encoded_text]) input_ids = torch.cat([input_ids for _ in range(args.batch_size)], dim=0) print('running model.sample') generated_sequence = model.sample(input_ids, sequence_length=args.n_positions) print('generated_sequence=', generated_sequence) outputs = [tokenizer.decode(gen_seq) for gen_seq in generated_sequence] print(outputs) ''' ``` 在后面加入 ```js prompt_text = input('Please input a prompt: ') while prompt_text!='exit': #print(prompt_text) with torch.inference_mode(): encoded_text = tokenizer.encode(prompt_text) input_ids = torch.as_tensor([encoded_text]) input_ids = torch.cat([input_ids for _ in range(args.batch_size)], dim=0) #print('running model.sample') generated_sequence = model.sample(input_ids, sequence_length=args.n_positions) #print('generated_sequence=', generated_sequence) outputs = [tokenizer.decode(gen_seq) for gen_seq in generated_sequence] print(outputs) prompt_text = input('\\nPlease input a prompt: ') ``` 再次运行以下命令 `gptj_demo run gpt-j-6B-split`  ##### 10. 至此，我们看到 GPT-J-6B 模型成功在 Inferentia2 芯片上完成了文本生成任务。运行完成后，用户不要忘记删除 EC2 Inf2 实例以节约成本。 # **总结** 本文介绍了亚马逊云科技的最新发布推理实例 Inf2，并演示了如何编译运行大语言模型 GPT-J-6B，为正在调研 Inf2 实例的用户提供参考。 由于 Inferentia2 芯片使用了与 Inferentia1 完全不同的架构设计，Neuron SDK 会同时包含对 Inferentia1 和Inferentia2 的支持，Inferentia1 对应的 SDK 名称为 neuron，Inferentia2 对应的 SDK 名称为 neuronx。Inf2 不会取代 Inf1。亚马逊云科技将继续向客户推荐 Inf1 和 Inf2 推理工作负载。一般来说，对于模型参数量较小的CV、NLP 模型，Inf1 更适合；而 Inf2 为较大的 ML 模型提供更好的性能和性价比。对 Inf2 的 Neuron SDK 模型支持在整个2023年的功能路线图中不断迭代，详细路线图可以参阅参考资料。 <!--StartFragment--> # **参考资料** 1.  https\://awsdocs-neuron.readthedocs-hosted.com/en/latest/index.html?trk=cndc-detail 2.  https\://aws.amazon.com/blogs/aws/amazon-ec2-inf2-instances-for-low-cost-high-performance-generative-ai-inference-are-now-generally-available/?trk=cndc-detail 3. https\://github.com/orgs/aws-neuron/projects/1/views/1?trk=cndc-detail 4. https\://aws.amazon.com/blogs/machine-learning/deploy-large-language-models-on-aws-inferentia2-using-large-model-inference-containers/?trk=cndc-detail ### **本篇作者**  **张铮** *亚马逊云科技[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)产品技术专家，负责基于亚马逊云科技加速计算和 GPU 实例的咨询和设计工作。专注于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)大规模模型训练和推理加速等领域，参与实施了国内多个[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)项目的咨询与设计工作。*  **郎建英** *亚马逊云科技资深解决方案架构师，专注于[机器学习](https://aws.amazon.com/cn/machine-learning/?trk=cndc-detail)和高性能计算在加速计器上面的应用，致力于云原生应用推广、落地。具有15年以上的 AIML/HPC 行业专业经验，在加入亚马逊云科技之前曾就职于 Intel 公司。*