注意:Alpa 目前没有被积极维护。它作为一个研究成果可用。Alpa 的核心算法已合并到仍在维护的 XLA 中。https://github.com/openxla/xla/tree/main/xla/hlo/experimental/auto_sharding
Alpa 是一个用于训练和服务大规模神经网络的系统。
将神经网络扩展到数百亿个参数已实现了诸如 GPT-3 这样的重大突破,但训练和服务这些大规模神经网络需要复杂的分布式系统技术。Alpa 旨在通过几行代码实现大规模分布式训练和服务的自动化。
Alpa 的关键特性包括:
💻 自动并行化。Alpa 自动将用户的单设备代码在数据、操作符和流水线并行的分布式集群上进行并行化。
🚀 卓越性能。Alpa 在分布式集群上训练具有数十亿参数的模型时,实现了线性扩展。
✨ 与机器学习生态系统的紧密集成。Alpa 由开源、高性能且具备生产就绪的库支持,如 Jax、XLA 和 Ray。
服务
以下代码展示了如何使用 huggingface/transformers 接口和 Alpa 分布式后端进行大模型推理。 详细文档请参见 使用 Alpa 服务 OPT-175B。
from transformers import AutoTokenizer
from llm_serving.model.wrapper import get_model
# 加载分词器
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 = "巴黎是"
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
将单设备训练代码扩展到分布式集群。
有关安装说明、教程、示例等,请查看 文档 网站和
examples 文件夹。
import alpa
# 通过简单地使用装饰器在 Jax 中并行化训练步骤
@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 slack 联系 Alpa 开发者。
- 如果您有兴趣贡献代码,请阅读 贡献者指南。
许可证
Alpa 依据 Apache-2.0 许可证 授权。