导读
埃隆·马斯克一直是开源的支持者，对OpenAI逐渐闭源的商业路线不满。随着在X上的正面开怼：“告诉我们更多关于OpenAI “开放”的部分（Tell us more about the“Open”part of OpenAI）”，2024年3月17日xAI开源了大语言模型Grok。本文主要揭秘Grok-1在全国产GPU上的运行体验。
DCU，大语言模型应用的解决方案
海光DCU (Deep Computing Unit)是一款面向人工智能、科学计算的高性能全功能GPGPU加速卡。
海光DCU具有丰富的软件栈DTK（DCU ToolKit），它全面兼容ROCm GPU计算生态，以及CUDA生态、工具链和开发者环境等，能够完整支持大语言模型的训练、推理、部署等相关应用，已适配大量模型，达到国内领先水平。
Grok-1基于Transformer模型架构，拥有3140亿参数，共包含64个Transformer层，每层都包含一个解码器层，由多头注意力块和密集块组成。多头注意力块有48个头用于查询，8个头用于键/值（KV），KV大小为128。密集块（密集前馈块）的加宽因子为8，隐藏层大小为32,768。采用MoE混合专家系统设计，每个token从8个专家中选择2个进行处理。Grok-1的激活参数数量为860亿，表明其在处理语言任务时的潜在能力强大，同时使用旋转嵌入位置编码，这是一种处理序列数据的方法，可以提高模型处理长文本的能力，上下文长度最大支持8192个tokens。xAI能在较短时间内训练出Grok-1，得益于其模型结构采用了MoE设计。MoE具有以下特点：2.与具有相同参数数量的模型相比，具有更快的推理速度。3.需要大量显存，因为所有专家系统都需要加载到内存中。混合专家模型 (MoE) 的理念起源于1991年的论文 Adaptive Mixture of Local Experts。一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。作为一种基于Transformer架构的模型，混合专家模型主要由两个关键部分组成:- 稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构。
- 门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如，在下图中，“More”这个令牌可能被发送到第二个专家，而“Parameters”这个令牌被发送到第一个专家。有时，一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。
总结来说，在混合专家模型 (MoE) 中，将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层，其中 MoE 层由两个核心部分组成: 一个门控网络和若干数量的专家。在专家并行中，专家被放置在不同的节点上，每个节点处理不同批次的训练样本。对于非 MoE 层，专家并行的行为与数据并行相同。对于 MoE 层，序列中的tokens被发送到拥有所需专家的节点。https://github.com/xai-org/grok-1?tab=readme-ov-file#downloading-the-weights1# 1.拉取docker镜像
2 docker pull image.sourcefind.cn:5000/dcu/admin/base/jax:0.4.23-ubuntu20.04-dtk24.04-py310
3
4 # 2.创建容器
5 docker run -it --name your_container_name --privileged --shm-size=1024G -v /opt/hyhal:/opt/hyhal --device=/dev/kfd --device=/dev/dri/ --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ulimit memlock=-1:-1 --ipc=host --network host --group-add video -v path_to_your_directory:path_to_your_directory imageID /bin/bash
6
7# 3.拉取grok仓库代码
8 git clone https://github.com/xai-org/grok-1.git
9
10 # 4.安装所需库
11 # 需注意将jax库注释，已在镜像容器中采用DCU编译后的jax库。
12 cd grok-1
13 pip install -r requirements.txt
14
15# 5.推理
在海光DCU上运行Grok-1模型，不仅展示了其卓越的计算性能，还证明了国产GPU在大语言模型领域的巨大潜力。通过DCU的高效计算能力，推理过程更加流畅。DTK工具包的丰富功能和兼容性，使开发者能够轻松迁移和优化现有模型，从而减少开发时间和成本。此外，海光DCU的出色性能使得Grok-1能够处理更大规模的数据集和更复杂的模型结构。这为研究人员和企业提供了更高效的开发和部署环境，促进了AI技术在各个领域的应用。
