Felafax -- 在 Google Cloud TPU 上微调 LLaMa3.1，成本降低 30%，并可无缝扩展！

Felafax 是一个基于 XLA 运行时 的开源大语言模型持续训练与微调框架。我们负责必要的运行时环境搭建，并提供开箱即用的 Jupyter 笔记本，让您立即上手。

• 易于使用。
• 能够轻松配置训练的各个方面（专为机器学习研究人员和黑客设计）。
• 训练规模可从单个拥有 8 个核心的 TPU 虚拟机，无缝扩展到包含 6000 个 TPU 核心的完整 TPU Pod（1000 倍）！

我们在 felafax 的目标是构建基础设施，使在非 NVIDIA 硬件（TPU、AWS Trainium、AMD GPU 和 Intel GPU）上运行 AI 工作负载变得更加容易。

✨ 免费微调

只需添加您的数据集，点击“运行全部”，即可在 Google Colab 上免费使用 TPU 资源进行训练！

Felafax 支持	免费笔记本
Llama 3.1 (1B, 3B)	▶️ 在 Google Colab TPU 上免费开始

目前支持的模型

• LLaMa-3.1 JAX 实现 $${\color{red}新!}$$

  • 从 PyTorch 转换为 JAX，以提升性能
  • 支持 1B、3B、8B、70B、405B 的全精度和 LoRA 训练。
  • 通过 JAX 经过硬件优化的 XLA 后端，在多种硬件（TPU、AWS Trainium、NVIDIA、AMD）上高效运行
  • 可通过跨多个加速器分片，无缝扩展以处理更大的上下文长度和数据集

• LLaMa-3/3.1 PyTorch XLA

  • 支持 LoRA 和全精度训练
  • codepointer

通过 Felafax CLI 运行微调 $${\color{red}新!}$$

只需几个简单步骤，即可使用 Felafax CLI 开始对您的模型进行微调。

第一步：安装 CLI 并进行身份验证

首先安装 CLI。

pip install pipx
pipx install felafax-cli

然后生成身份验证令牌：

• 访问 preview.felafax.ai，创建或登录您的账户。
• 导航到 Tokens 页面，创建一个新的令牌。

最后，使用您的令牌对 CLI 会话进行身份验证：

felafax-cli auth login --token <your_token>

第二步：设置微调配置

首先，生成一个默认的微调配置文件。此命令会在当前目录下生成一个包含默认超参数值的 config.yml 文件。

felafax-cli tune init-config

其次，根据您的超参数更新配置文件：

• HuggingFace 参数：

  • 提供您的 HuggingFace 令牌和仓库 ID，以便上传微调后的模型。

• 数据集管道和训练参数：

  • 调整用于微调数据集的 batch_size 和 max_seq_length。
  • 如果希望训练遍历整个数据集，将 num_steps 设置为 null。如果设置了具体的步数，训练将在指定步数后停止。
  • 设置用于微调的 learning_rate 和 lora_rank。
  • eval_interval 是两次评估之间的步数。

第三步：开始微调任务

运行以下命令查看您可以微调的基础模型列表，目前我们支持所有 LLaMA-3.1 的变体。

felafax-cli tune start --help

现在，您可以使用上述列表中的选定模型以及来自 HuggingFace 的数据集名称（如 yahma/alpaca-cleaned）开始微调过程：

felafax-cli tune start --model <your_selected_model> --config ./config.yml --hf-dataset-id <your_hf_dataset_name>

一个示例命令可以帮助您入门：

felafax-cli tune start --model llama3-2-1b --config ./config.yml --hf-dataset-id yahma/alpaca-cleaned

启动微调任务后，Felafax CLI 会负责启动 TPU、运行训练，并将微调后的模型上传到 HuggingFace Hub。

其他实用命令

监控微调任务

您可以流式传输实时日志来监控微调任务的进度：

# 使用您在开始微调后获得的作业名称 `<job_name>`。
felafax-cli tune logs --job-id <job_name> -f

列出您的微调模型

微调完成后，您可以列出所有已微调的模型：

felafax-cli model list

与您的微调模型对话（再次在 TPU 上运行！）：

您可以启动一个交互式终端会话，与您的微调模型进行对话：

# 将 `<model_id>` 替换为您之前运行的 `model list` 命令中的模型 ID。
felafax-cli model chat --model-id <model_id>

使用帮助探索更多命令！

CLI 分为三大命令组：

• tune：用于启动或停止微调任务。
• model：用于管理和与您的微调模型互动。
• files：用于上传和查看您的数据集文件。

使用 --help 标志可以了解更多关于任何命令组的信息：

felafax-cli tune --help

AMD 405B 微调运行

我们最近使用 JAX 而不是 PyTorch，在 8 张 AMD MI300x GPU 上对 llama3.1 405B 模型进行了微调。JAX 先进的分片 API 使我们能够实现出色的性能。请查看我们的 博客文章 ，了解我们使用的设置和分片技巧。

我们以 bfloat16 精度对所有模型权重和 LoRA 参数进行了 LoRA 微调，LoRA 阶数为 8，LoRA alpha 为 16：

• 模型大小： LLaMA 模型权重占用约 800GB 的显存。
• LoRA 权重 + 优化器状态： 大约 400GB 的显存。
• 总显存占用： 占总显存的 77%，约为 1200GB。
• 限制： 由于 405B 模型体积庞大，批次大小和序列长度的空间有限。使用的批次大小为 16，序列长度为 64。
• 训练速度： 约 35 tokens/秒
• 内存效率： 一直保持在 70% 左右
• 扩展性： 使用 JAX，在 8 张 GPU 上的扩展几乎是线性的。

GPU 利用率和显存利用率图表如下所示。然而，我们仍需计算模型 FLOPs 利用率 (MFU)。 注：由于基础设施和显存限制，我们无法运行 405B 模型的 JIT 编译版本（这一点还需要进一步研究）。整个训练过程都是在 JAX eager 模式下执行的，因此仍有很大的性能提升空间。

• GPU 利用率：
• 显存利用率：
• rocm-smi 数据可在此处找到 这里。

致谢：

• Google Deepmind 的 Gemma 仓库。
• EasyLM 和 EleutherAI 在 JAX 中对 Llama 模型所做的出色工作。
• PyTorch XLA FSDP 和 SPMD 测试由 HeegyuKim 完成。
• 示例来自 PyTorch-XLA 仓库。

联系方式

如果您有任何问题，请通过 founders@felafax.ai 与我们联系。

Felafax 快速上手指南

Felafax 是一个基于 XLA 运行时的框架，专为在 Google Cloud TPU 及其他非 NVIDIA 硬件（如 AMD GPU、AWS Trainium）上持续训练和微调开源大语言模型（如 LLaMa 3.1）而设计。它支持从单台 TPU VM 到数千核 TPU Pod 的无缝扩展，并提供免费的 Google Colab TPU 资源供初学者使用。

环境准备

系统要求

• 操作系统：Linux 或 macOS（Windows 用户建议使用 WSL2 或 Docker）
• Python 版本：Python 3.8+
• 硬件支持：
  • Google Cloud TPU（推荐用于免费试用）
  • AWS Trainium
  • AMD GPU (MI300x 等)
  • NVIDIA GPU (通过 JAX/XLA 后端)
• 网络要求：需能访问 Hugging Face 和 Google Cloud 服务（国内用户可能需要配置代理或使用加速镜像）

前置依赖

• pip 和 pipx 包管理工具
• Hugging Face 账号及 Token（用于上传微调后的模型）
• Felafax 平台账号（用于 CLI 认证）

提示：国内开发者若访问 Hugging Face 受阻，建议在配置中使用国内镜像源或在具备科学上网环境的服务器上进行操作。

安装步骤

1. 安装 Felafax CLI

使用 pipx 安装命令行工具，确保环境隔离：

pip install pipx
pipx install felafax-cli

2. 获取认证令牌

1. 访问 preview.felafax.ai 注册或登录账号。
2. 进入 Tokens 页面创建一个新的 API Token。

3. 登录 CLI

使用生成的 Token 完成认证：

felafax-cli auth login --token <your_token>

基本使用

以下流程将演示如何使用 Felafax CLI 快速启动一个 LLaMa 3.1 模型的微调任务。

第一步：初始化配置文件

生成默认的超参数配置文件 config.yml：

felafax-cli tune init-config

编辑 config.yml，填入你的 Hugging Face Token 和仓库 ID，并根据需要调整 batch_size、learning_rate 等参数。

第二步：启动微调任务

查看支持的基座模型列表：

felafax-cli tune start --help

执行微调命令（以 LLaMa-3.2-1B 和 Alpaca 数据集为例）：

felafax-cli tune start --model llama3-2-1b --config ./config.yml --hf-dataset-id yahma/alpaca-cleaned

说明：执行后，Felafax 会自动调配 TPU 资源、运行训练，并在完成后将模型上传至你的 Hugging Face 仓库。

第三步：监控与交互

实时监控日志：

felafax-cli tune logs --job-id <job_name> -f

查看已微调的模型：

felafax-cli model list

在线对话测试（直接在 TPU 上运行）：

felafax-cli model chat --model-id <model_id>

---

更多高级用法（如文件上传、自定义分片策略）可通过 felafax-cli --help 探索。