higgsfield - 高容错且可扩展的GPU调度与机器学习框架

higgsfield 项目介绍

higgsfield 是一个开源的、容错的、高度可扩展的 GPU 编排和机器学习框架，专为训练拥有数十亿到数万亿参数的模型而设计，如大型语言模型（LLMs）。这个项目旨在简化大规模模型的训练过程，为开发者提供一套多功能且强大的工具集。

核心功能

higgsfield 作为 GPU 工作负载管理器和机器学习框架，主要提供以下五个核心功能：

资源分配：为用户的训练任务分配独占和非独占的计算资源（节点）。
高效分片支持：支持 ZeRO-3 deepspeed API 和 PyTorch 的完全分片数据并行 API，实现万亿参数模型的高效分片。
训练框架：提供一个框架用于在分配的节点上启动、执行和监控大型神经网络的训练。
资源竞争管理：通过维护实验运行队列来管理资源竞争。
持续集成：通过与 GitHub 和 GitHub Actions 的无缝集成，促进机器学习开发的持续集成。

安装和使用

安装 higgsfield 非常简单，只需要通过 pip 命令即可完成：

$ pip install higgsfield==0.0.3

使用 higgsfield 训练模型也非常直观。以训练 LLaMa 模型为例，开发者只需要几行代码就可以在分布式环境中进行训练：

from higgsfield.llama import Llama70b
from higgsfield.loaders import LlamaLoader
from higgsfield.experiment import experiment

import torch.optim as optim
from alpaca import get_alpaca_data

@experiment("alpaca")
def train(params):
    model = Llama70b(zero_stage=3, fast_attn=False, precision="bf16")
    optimizer = optim.AdamW(model.parameters(), lr=1e-5, weight_decay=0.0)
    dataset = get_alpaca_data(split="train")
    train_loader = LlamaLoader(dataset, max_words=2048)

    for batch in train_loader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()

    model.push_to_hub('alpaca-70b')

工作原理

higgsfield 的工作流程主要包括以下步骤：

在服务器上安装所需的工具（Docker、项目部署密钥、higgsfield 二进制文件）。
为实验生成部署和运行工作流。
代码一旦上传到 GitHub，就会自动部署到指定节点。
通过 GitHub 访问实验运行界面，启动实验并保存检查点。

设计理念

higgsfield 遵循标准的 PyTorch 工作流程，因此开发者可以轻松集成其他工具，如 deepspeed、accelerate，甚至可以从头实现自定义的 PyTorch 分片。

项目还解决了两个常见的问题：

环境地狱：不再需要担心不同版本的 PyTorch、NVIDIA 驱动程序或数据处理库。higgsfield 允许轻松编排实验及其环境，记录和跟踪所有依赖项的特定版本和配置，确保实验的可重复性。
配置地狱：不再需要为实验定义大量参数或使用复杂的 YAML 配置。higgsfield 提供了一个简单的接口来定义实验，开发者只需设计交互方式即可。

兼容性

higgsfield 对节点有以下要求：

Ubuntu 操作系统
SSH 访问
具有 sudo 权限的非 root 用户（无需密码）

已经在 Azure、LambdaLabs 和 FluidStack 等云平台上进行了测试。

入门指南

higgsfield 提供了详细的入门指南，包括项目初始化、环境设置、Git 设置、节点设置以及运行第一个实验等步骤。此外，还提供了涵盖分布式模型、数据准备、模型参数优化、模型保存、训练稳定技术和监控等主题的教程。

总的来说，higgsfield 为大规模模型训练提供了一个强大、灵活且易于使用的框架，可以帮助研究人员和开发者更高效地进行机器学习实验和开发。