机器之心编辑部
随着 AI 模型的参数量越来越大，对算力的需求也水涨船高。
比如最近，Llama-3.1 登上了最强开源大模型的宝座，但超大杯 405B 版本的内存就高达 900 多 GB，这对算力构成了更加苛刻的寻衅。
如何降落算力的利用本钱和利用门槛，已经成为许多公司寻求打破的关键。Felafax 便是个中的一家创业公司，致力于简化 AI 演习集群的搭建流程。
Nikhil Sonti 和 Nikhin Sonti 创立了 Felafax，他们的口号是在构建开源 AI 平台，为下一代 AI 硬件做事，将机器学习的演习本钱降落 30%。
与英伟达比较，AMD 的 GPU，尤其是 MI300X 系列，供应了更高的性价比，按每美元打算，其性能表现更为出色。
最近，Felafax 的联合创始人 Nikhil Sonti 发布了一篇博客，详细分享了如何通过 8 张 AMD MI300X GPU 和 JAX 微调 LLaMA 3.1 405B 模型的方法，所有代码现已开源。
Github 链接：https://github.com/felafax/felafax
机器之心对博客内容进行了不改变原意的编译、整理，以下是博客内容：
JAX 尤其适宜非英伟达硬件
JAX 是一个强大的机器学习库，结合了类似 NumPy 的 API、自动微分功能以及 Google 的 XLA 编译器。它在模型并行化方面供应了精良的 API，因此非常适宜像 LLaMA 3.1 405B 这样的超大模型演习。
在利用 AMD 硬件时，JAX 有几个明显的上风：
多硬件并行支持：JAX 采取 XLA（加速线性代数）编译器，将打算编译为硬件无关的中间表示（HLO），这意味着同样的 JAX 代码无需修正便可高效运行在不同硬件后端，包括 AMD GPU。独立于底层硬件：XLA 编译器的优化策略是通用的，不针对某个特定的硬件平台。这使得任何支持 XLA 的硬件设备（如 CPU、GPU、TPU）都能受益于这些优化，得到更好的性能表现。极高的适应性：从 NVIDIA 转移到 AMD（或其他硬件）时，JAX 只需做极少的代码改动。而相较之下，PyTorch 与英伟达的 CUDA 生态系统紧密耦合，迁移过程相对繁芜。因此，JAX 成为了我们在非英伟达硬件上的最佳选择。
拉取 Docker 镜像：
docker pull rocm/jax:latest
启动 Docker 容器：
# Pull the Docker Image:
docker pull rocm/jax:latest
# Start the Docker Container:
docker run -it -w /workspace --device=/dev/kfd --device=/dev/dri --group-add video \
--cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 16G rocm/jax:latest
# Verify the Installation:
python3 -c 'import jax; print(jax.devices())'
验证安装
python3 -c 'import jax; print (jax.devices ())'
演习利用了一个配备了 8 张 AMD MI300x GPU 的 AMD 节点。每张 MI300x 拥有 192GB 的 HBM3 内存，性能表现与最新的英伟达 H100 GPU 比较非常出色。
与英伟达 H100 的比较，来源：TensorWave
演习 LLaMA 405B：性能与可扩展性
利用 JAX，可以成功地在 AMD GPU 上演习 LLaMA 405B 模型。我们利用 LoRA 微调，将所有模型权重和 LoRA 参数都设为 bfloat16，LoRA rank 设为 8，LoRA alpha 设为 16：
模型大小：LLaMA 模型的权重占用了约 800GB 的显存。LoRA 权重 + 优化器状态：大约占用了 400GB 的显存。显存总利用量：占总显存的 77%，约 1200GB。限定：由于 405B 模型的规模过大，batch 大小和序列长度的空间有限，利用的 batch size 为 16，序列长度为 64。JIT 编译：由于空间限定，无法运行 JIT 编译版本；它可能须要比迫切模式稍多的空间。演习速率：利用 JAX 迫切模式，约为 35 tokens / 秒。内存效率：稳定在约 70% 旁边。扩展性：在 8 张 GPU 上，利用 JAX 的扩展性靠近线性。由于硬件和显存的限定，我们无法运行 JIT 编译版本的 405B 模型，全体演习过程是在 JAX 的迫切模式下实行的，因此还有很大的进步空间。
下图中显示了在一次微调演习步骤中，8 张 GPU 的显存利用率和 rocm-smi 输出：
GPU 利用率：
显存利用率：
rocm-smi 输出：
演习设置
将 LLaMA 3.1 从 PyTorch 移植到 JAX
此前，Nikhil Sonti 分享过如何将 LLaMA 3.1 从 PyTorch 移植到 JAX。他指出，目前 90% 的大型措辞模型（LLM）都运行在 NVIDIA GPU 上，但实际上还有一些同样强大且性价比更高的替代方案。例如，在 Google TPU 上演习和支配 Llama 3.1 的本钱比 NVIDIA GPU 低约 30%。
然而，支持非 NVIDIA 硬件的开拓工具较为匮乏。Sonti 最初考试测验利用 PyTorch XLA 在 TPU 上演习 Llama 3.1，但过程并不顺利。XLA 与 PyTorch 的集成不足完善，短缺一些关键的库（如 bitsandbytes 无法正常运行），同时还碰着了一些难以办理的 HuggingFace 缺点。
为此，他决定调度策略，将 Llama 3.1 从 PyTorch 移植到 JAX，成功办理了这些问题。Sonti 还录制了详细的教程视频，并开源了所有代码：
方法演示：https://dub.sh/felafax-demo代码仓库：https://github.com/felafax/felafax
加载模型，并把模型参数分片
处理像 LLaMA 405B 这样的超大模型，须要在多个设备之间高效地进行参数分片。以下是如何通过 JAX 实现这一点的。
在 JAX 中进行参数分片
为了将巨大的 LLaMA 405B 模型高效地分布到 8 张 AMD GPU 上，须要利用 JAX 的设备网格（device mesh）功能。
支配代码：https://github.com/felafax/felafax/blob/e2a96a0e207e1dc70effde099fe33a9e42a7d5cb/llama3_jax/trainer_engine/jax_utils.py#L69
JAX 的设备网格可以帮助我们把可用的设备组织成一个网格，让我们可以指定如何把模型的参数和打算分配到不同的 GPU 上。
在本文的设置中，须要创建一个形状为（1, 8, 1）的网格，并将轴分别命名为数据并行（dp）、全分片数据并行（fsdp）和模型并行（mp）。然后，为模型的每个张量定义特定的分片规则，指定这些维度如何沿着这些网格轴进行分片。
DEVICES = jax.devices ()
DEVICE_COUNT = len (DEVICES)
DEVICE_MESH = mesh_utils.create_device_mesh ((1, 8, 1))
MESH = Mesh (devices=DEVICE_MESH, axis_names=("dp", "fsdp", "mp"))
可视化分片
可以利用以下代码来可视化分片结果，从而方便地验证分片规则是否按预期运用。
jax.debug.visualize_array_sharding
分片规则
模型不同组件的分片规则如下所示：
参数如何分片：参数要在 8 个 GPU 之间分配。例如，LM head（lm_head/kernel）张量有两个轴，按照 PS ("fsdp", "mp") 进行分片。在本例中是 8 和 1，因此可以看到该张量在第一个轴上沿着 8 个 GPU 被拆分。
Non-Replicated 参数：没有任何分片规范的参数会在所有设备上进行复制。例如，层归一化（attention_norm/kernel 和 ffn_norm/kernel）没有设置分片规范，是 PS (None)。
运用分片函数
在加载模型时，利用以下分片函数逐步对模型权重进行分片：
def make_shard_and_gather_fns (partition_specs):
def make_shard_fn (partition_spec):
out_sharding = NamedSharding (mesh, partition_spec)
def shard_fn (tensor):
return jax.device_put (tensor, out_sharding).block_until_ready ()
return shard_fn
shard_fns = jax.tree_util.tree_map (make_shard_fn, partition_specs)
return shard_fns
# Create shard functions based on partitioning rules
shard_fns = make_shard_and_gather_fns (partitioning_rules)
这使得我们能够将每个参数放置在指定的设备上，并按照设定的分片进行处理。
分片演习 Batch
最初，演习 Batch 是正常创建的，但在输入模型之前，须要按照下面的代码在 GPU 上进行分片：
train_batch = jax.device_put ( train_batch,
NamedSharding (self.mesh, PS ("dp", "fsdp")))
在这里，我们指定演习 Batch 该当在 "dp" 和 "fsdp" 轴上进行分片，在本例等分别对应于被分成 1 和 8 份，如果把结果可视化出来，如下所示：
分片前：
在调用 jax.device_put 之后：
加入 LoRA
LoRA 通过将权重更新分解为低秩矩阵，减少了可演习参数的数量，这对付微调大型模型特殊有效。以下是在 AMD GPU 上微调 Llama 3.1-405 的 LoRA 的要点：
将 LoRA 参数（lora_a 和 lora_b）与主模型参数分开。利用 jax.lax.stop_gradient (kernel) 来防止对主模型权重的更新。利用 lax.dot_general 进行快速、精确掌握的矩阵运算。LoRA 输出在添加到主输出之前会被缩放为 (self.lora_alpha/self.lora_rank)。LoRADense 层
在此设定一个自定义的 LoRADense 层，该层集成了 LoRA 参数：
class LoRADense (nn.Module):
features: int
lora_rank: int = 8
lora_alpha: float = 16.0
@nn.compact
def __call__(self, inputs: Any) -> Any:
# Original kernel parameter (frozen)
kernel = self.param ('kernel', ...)
y = lax.dot_general (inputs, jax.lax.stop_gradient (kernel), ...)
# LoRA parameters (trainable)
lora_a = self.variable ('lora_params', 'lora_a', ..., ...)
lora_b = self.variable ('lora_params', 'lora_b', ..., ...)
# Compute LoRA output
lora_output = lax.dot_general (inputs, lora_a.value, ...)
lora_output = lax.dot_general (lora_output, lora_b.value, ...)
# Combine original output with LoRA modifications
y += (self.lora_alpha/self.lora_rank) lora_output
return y.astype (self.dtype)
分片 LoRA 参数
为了高效地在设备之间分配 LoRA 参数，我们也通过 JAX 设定了分片规则，这确保了 LoRA 参数与主模型参数的分片同等，优化了内存利用和打算效率。
LoRA A matrices (lora_a)
LoRA A 矩阵（lora_a）
分片规则：PS ("fsdp", "mp")可视化结果：如下图所示，lora_a 参数被分片为 (8, 1)，这意味着第一个轴在 8 个设备上进行分片（"fsdp" 轴），而第二个轴未进行分片。LoRA B 矩阵（lora_b）
分片规则：PS ("mp", "fsdp")可视化结果：如下图所示，lora_b 参数被分片为 (1, 8)，这意味着第二个轴在 8 个设备上进行分片（fsdp 轴），而第一个轴未进行分片。这种分片策略优化了参数的分配，减少了通信开销，并在演习过程中增强了并行性。它确保每个设备仅持有一部分 LoRA 参数，使得大模型如 LLaMA 405B 的高效扩展成为可能。
仅更新 LoRA 参数
为了优化演习，在微调 LLaMA 405B 模型，只打算 LoRA 参数的梯度，保持主模型参数不变。这个方法减少了内存利用，并加速了演习，由于只更新较少的参数。可以移步 GitHub 仓库，查看实现细节。
在演习过程中，每一步都涉及将一批输入数据通过模型进行处理。由于只有 LoRA 参数是可演习的，因此模型的预测和打算的丢失仅依赖于这些参数，然后对 LoRA 参数进行反向传播。只更新这些参数简化了演习过程，使得在多个 GPU 上高效微调像 LLaMA 405B 这样的大型模型成为可能。
更多研究细节，请参考原博客。