AWS AI芯片助力低成本高效能的Llama 31模型

关键要点

今天，我们很高兴地宣布，AWS Trainium 和 AWS Inferentia 支持对Llama 31模型进行微调和推断。Llama 31家族的多语言大型语言模型LLMs包含8B、70B和405B三种预训练和指令调优的生成模型。在之前的一篇博文，我们讨论了如何在基于AWS Trainium和Inferentia的实例上部署Llama 3模型。本文将概述如何在AWS AI芯片上开始微调和部署Llama 31模型，以实现其性价比优势。

Llama 31模型概览

Llama 31家族的多语言LLMs是一组包含8B、70B和405B三种规模的预训练和指令调优的生成模型文本输入/文本和代码输出。所有模型都支持长上下文长度128k，并优化了带有分组查询注意力GQA的推理效果。

Llama 31指令调优模型8B、70B、405B针对多语言对话用例进行了优化，在许多公开的聊天模型上表现优异，特别在常见的行业基准测试中。它们经过训练可以生成特定工具的调用，支持如搜索、图像生成、代码执行和数学推理等功能。此外，它们支持零样本工具使用。

Llama 31 405B被Meta称为世界上最大的公开LLM。该模型为人工智能AI设立了新的标准，非常适合企业级应用及研发。它特别适用于合成数据生成等任务，其中模型的输出可以用于改进较小的Llama模型，并通过模型蒸馏将知识转移到较小的405B模型。这款模型在通用知识、长文本生成、多语言翻译、机器翻译、编码、数学、工具使用、增强上下文理解以及高级推理与决策方面表现出色。

在架构上，Llama 3和Llama 31的核心LLM采用相同的致密架构。它们是自回归语言模型，使用优化的变换器架构。调优版本通过监督微调SFT和带有人工反馈的强化学习RLHF来与人类的需求保持一致。

Meta的负责任使用指南可以帮助你实施额外的微调，以便定制和优化模型并实施适当的安全减轻措施。

Trainium在Amazon Bedrock和Amazon SageMaker上的应用

在AWS上开始使用Llama 31的最快方法是通过Amazon Bedrock，它由我们专门构建的AI基础设施，包括AWS Trainium支持。通过其完全托管的API，Amazon Bedrock提供了我们的专用AI基础设施的优势，便于访问这些强大的模型，使你能够专注于构建具有差异化的AI应用。

如果你需要更大的控制权，可以使用SageMaker对Llama 31模型进行微调和部署。Trainium在SageMaker JumpStart中的支持即将上线。

AWS Trainium和AWS Inferentia2为Llama 31模型提供高性能和低成本

如果你希望构建自己的机器学习流水线以进行训练和推断，能够获得更大的灵活性和控制权，你可以使用AWS弹性计算云Amazon EC2Trn1和Inf2实例开始Llama 31的使用。接下来，我们将看看如何使用AWS Neuron SDK开始使用新的Llama 31 8B/70B模型。

微调Llama 31

要开始微调Llama 31 8B或Llama 31 70B，你可以使用NeuronX Distributed库。NeuronX Distributed提供一些流行的分布式训练和推断技术的实现。要开始微调，你可以使用以下示例：

这两个示例都基于AWS ParallelCluster来管理Trainium集群基础设施和Slurm进行工作负载管理。以下是启动Llama 31 70B训练的示例Slurm命令：

bashsbatch exclusive nodes 32 cpuspertask 128 wrap=srun bash (pwd)/runllama370Btpppsh

在Slurm脚本中，我们在集群中启动分布式训练过程。在运行脚本中，我们加载了Meta提供的预训练权重和配置，并启动训练过程：

bashtorchrun DISTRIBUTEDARGS runllamanxdpy trainbatchsize BS usemetadeviceinit 1 trainingdir DATAPATH trainingconfig SCRIPTDIR/{MODELSIZE}configllama{LLAMAVERSION} maxsteps maxsteps seqlen SEQLEN pipelineparallelsize PPDEGREE tensorparallelsize TPDEGREE nummicrobatches NUMMICROBATCHES lr 0000015 minlr 1e06 beta1 09 beta2 095 weightdecay 01 warmupsteps 2000 constantsteps 0 usezero1optimizer 1 useselectivecheckpoint 1 useflashattention 1 qkvlinear 1 kvreplicator 4 pretrainedweight 1 saveloadxser 1 checkpointdir /shared/llama{LLAMAVERSION}{MODELSIZE}/ checkpointfreq checkpointfreq numkeptcheckpoint 1 loadingstep 1 tbdir tbdir amp tee LOGPATH/logexit {PIPESTATUS[0]}

在Trainium或Inferentia上部署Llama 31

当你的模型准备好进行部署时，可以通过更新之前Llama 3 8B神经元示例代码中的模型ID来完成。例如，以下代码在inf248xlarge实例上部署模型：

pythonmodelid = metallama/MetaLlama318Bneuronmodel = LlamaForSamplingfrompretrained(modelid neuronconfig=neuronconfig batchsize=1 tpdegree=24 amp=bf16 npositions=4096)neuronmodeltoneuron()

你可以使用相同的示例推断代码：

pythontokenizer = AutoTokenizerfrompretrained(modelid)prompt = Hello Im a language model and I like toinputids = tokenizerencode(prompt returntensors=pt)

使用topk采样进行推断

with torchinferencemode() start = timetime() generatedsequences = neuronmodelsample(inputids sequencelength=2048 topk=50) elapsed = timetime() start

generatedsequences = [tokenizerdecode(seq) for seq in generatedsequences]print(f生成的序列 {generatedsequences} 用时 {elapsed} 秒)

有关逐步详细信息，请参考新的Llama 31示例：

此外，如果你想使用Hugging Face的Optimum Neuron库，可以直接通过SageMaker快速部署模型。通过Llama 31模型卡Hub，选择部署，然后选择SageMaker，最后选择AWS Inferentia amp Trainium。将示例代码复制到SageMaker笔记本中，然后选择运行。

pythonimport jsonimport sagemakerimport boto3from sagemakerhuggingface import HuggingFaceModel gethuggingfacellmimageuri

try role = sagemakergetexecutionrole()except ValueError iam = boto3client(iam) role = iamgetrole(RoleName=sagemakerexecutionrole)[Role][Arn]

Hub模型配置 https//huggingfaceco/models

hub = { HFMODELID metallama/MetaLlama318B HFNUMCORES 2 HFAUTOCASTTYPE fp16 MAXBATCHSIZE 8 MAXINPUTLENGTH 3686 MAXTOTALTOKENS 4096 HFTOKEN lt替换为你的TOKENgt}

assert hub[HFTOKEN] != lt替换为你的TOKENgt 请将lt替换为你的TOKENgt替换为你的Hugging Face Hub API令牌

创建Hugging Face模型类

huggingfacemodel = HuggingFaceModel( imageuri=gethuggingfacellmimageuri(huggingfaceneuronx version=0023) env=hub role=role)

部署模型到SageMaker推断

predictor = huggingfacemodeldeploy( initialinstancecount=1 instancetype=mlinf2xlarge containerstartuphealthchecktimeout=1800 volumesize=512)

发送请求

predictorpredict( { inputs 法国的首都是哪里？ parameters { dosample True maxnewtokens 128 temperature 07 topk 50 topp 095 } })

如果你想使用vLLM部署模型，可以参考连续批次指南以创建环境。创建环境后，你可以使用vLLM在AWS Trainium或Inferentia上部署Llama 31 8/70B模型。以下是部署Llama 31 8B的示例：

pythonfrom vllm import LLM SamplingParams

示例提示语。

prompts = [ 你好，我的名字是 美国总统是 法国的首都是 AI的未来是]

创建采样参数对象。

samplingparams = SamplingParams(temperature=08 topp=095)

创建一个LLM。

llm = LLM( model=metallama/MetaLlama318B maxnumseqs=8 maxmodellen=128 blocksize=128 device=neuron tensorparallelsize=8)

从提示生成文本。输出是一个请求输出对象的列表，包含提示、生成的文本和其他信息。

outputs = llmgenerate(prompts samplingparams)

打印输出。

for output in outputs prompt = outputprompt generatedtext = outputoutputs[0]text print(f提示：{prompt!r} 生成的文本：{generatedtext!r})

结论

AWS Trainium和Inferentia为Llama 31模型的微调和部署提供了高性能和低成本的解决方案。我们期待你利用这些强大的模型和我们专门构建的AI基础设施，构建具有差异化的AI应用。有关如何开始使用AWS AI芯片的更多信息，请参考AWS Neuron文档中的模型示例和教程。

关于作者

John Gray 是Annapurna Labs的高级解决方案架构师，工作地点在西雅图。在这个角色中，John与客户合作处理他们的AI和机器学习用例，设计解决方案以具有成本效益地解决业务问题，并帮助他们使用AWS AI芯片构建可扩展的原型。

Pinak Panigrahi与客户合作，构建机器学习驱动的解决方案来解决AWS上的战略商业问题。在目前的角色中，他专注于优化生成AI模型在AWS AI芯片上的训练和推断。

Kamran Khan是AWS Inferentia/Trainium的业务开发主管，拥有超过十年的经验，帮助客户部署和优化深度学习训练和推断工作负载，使用AWS Inferentia和AWS Trainium。

Shruti Koparkar是AWS的高级产品营销经理，帮助客户探索、评估和采用Amazon EC2加速计算基础设施，以满足他们的机器学习需求。