ALPA: 自动并行化大规模神经网络训练和服务的开源框架

ALPA: 革新大规模神经网络训练和服务的开源利器

在人工智能领域,大规模神经网络的出现如GPT-3等模型带来了突破性进展。然而,训练和部署这些动辄数十亿参数的庞大模型,对分布式系统技术提出了巨大挑战。为了让更多研究者和开发者能够便捷地驾驭大模型,来自加州大学伯克利分校等机构的研究人员开发了ALPA(Automated Layout Parallelism Analyzer)开源框架。ALPA旨在通过自动并行化技术,让用户只需几行代码就能在分布式集群上高效训练和部署大型模型。

ALPA的核心特性

ALPA具有以下几个突出特点:

自动并行化: ALPA能够自动将用户的单设备代码并行化到分布式集群上,支持数据并行、算子并行和流水线并行等多种并行策略。
卓越性能: 在训练数十亿参数的大模型时,ALPA可以在分布式集群上实现近乎线性的扩展。
与机器学习生态系统紧密集成: ALPA基于开源的高性能生产级库如JAX、XLA和Ray构建,可以无缝对接现有的机器学习工作流。

ALPA logo

ALPA的使用方法

使用ALPA非常简单,只需在原有的JAX训练代码上添加一个装饰器即可:

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)

对于模型推理,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)