Alpa入门指南 - 自动并行化大规模神经网络训练与部署系统

Alpa简介

Alpa是一个用于训练和部署大规模神经网络的自动并行化系统。它可以自动将单设备代码扩展到分布式集群上,实现数据、算子和流水线并行。Alpa的主要特点包括:

自动并行化:只需几行代码即可实现分布式训练
卓越性能:可线性扩展至数十亿参数规模的模型训练
与机器学习生态系统紧密集成:基于JAX、XLA、Ray等开源高性能库

Alpa logo

快速上手

模型部署

使用Hugging Face Transformers接口和Alpa分布式后端进行大模型推理非常简单:

from transformers import AutoTokenizer
from llm_serving.model.wrapper import get_model

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-2.7b")
tokenizer.add_bos_token = False

# 加载模型,Alpa会自动下载权重
model = get_model(model_name="alpa/opt-2.7b", path="~/opt_weights/")

# 生成文本
prompt = "Paris is the capital city of"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
output = model.generate(input_ids=input_ids, max_length=256, do_sample=True)
generated_string = tokenizer.batch_decode(output, skip_special_tokens=True)

print(generated_string)

模型训练

只需添加@alpa.parallelize装饰器,即可将单设备训练代码扩展到分布式集群:

import alpa

@alpa.parallelize
def train_step(model_state, batch):
    def loss_func(params):
        out = model_state.forward(params, batch["x"])
        return jnp.mean((out - batch["y"]) ** 2)

    grads = grad(loss_func)(model_state.params)
    new_model_state = model_state.apply_gradient(grads)
    return new_model_state

# 训练循环自动在指定集群上运行
model_state = create_train_state()
for batch in data_loader:
    model_state = train_step(model_state, batch)