Haliax 开源项目教程-优快云博客

Haliax 开源项目教程

haliax Named Tensors for Legible Deep Learning in JAX 项目地址: https://gitcode.com/gh_mirrors/ha/haliax

1. 项目介绍

Haliax 是一个基于 JAX 的库，旨在通过使用命名张量（Named Tensors）来提高深度学习代码的可读性和组合性。命名张量通过使用命名轴（Named Axes）代替传统的基于位置的索引，使得张量操作更加直观和易于理解。Haliax 不仅关注代码的可读性，还注重性能，通常与纯 JAX 代码的性能相当。此外，Haliax 还支持分布式训练，包括全分片数据并行（FSDP）和张量并行，适用于训练大规模语言模型和其他基础模型。

2. 项目快速启动

安装 Haliax

首先，确保你已经安装了 JAX 和 Equinox 库。然后，你可以通过以下命令安装 Haliax：

pip install haliax

示例代码：注意力机制

以下是一个使用 Haliax 实现的简单注意力机制模块的示例代码：

import equinox as eqx
import jax
import jax.numpy as jnp
import haliax as hax
import haliax.nn as hnn

# 定义轴
Pos = hax.Axis("position", 1024)  # 序列长度
KPos = Pos.alias("key_position")
Head = hax.Axis("head", 8)  # 注意力头数
Key = hax.Axis("key", 64)  # 键大小
Embed = hax.Axis("embed", 512)  # 嵌入大小

# 计算注意力分数
def attention_scores(Key, KPos, query, key, mask):
    scores = hax.dot(query, key, axis=Key) / jnp.sqrt(Key.size)
    if mask is not None:
        scores -= 1E9 * (1.0 - mask)
    scores = hnn.softmax(scores, KPos)
    return scores

# 注意力机制
def attention(Key, KPos, query, key, value, mask):
    scores = attention_scores(Key, KPos, query, key, mask)
    answers = hax.dot(scores, value, axis=KPos)
    return answers

# 因果掩码
causal_mask = hax.arange(Pos).broadcast_axis(KPos) >= hax.arange(KPos)

# 注意力模块
class Attention(eqx.Module):
    proj_q: hnn.Linear  # [Embed] -> [Head, Key]
    proj_k: hnn.Linear  # [Embed] -> [Head, Key]
    proj_v: hnn.Linear  # [Embed] -> [Head, Key]
    proj_answer: hnn.Linear  # 输出投影 [Head, Key] -> [Embed]

    @staticmethod
    def init(Embed, Head, Key, *, key):
        k_q, k_k, k_v, k_ans = jax.random.split(key, 4)
        proj_q = hnn.Linear.init(In=Embed, Out=(Head, Key), key=k_q)
        proj_k = hnn.Linear.init(In=Embed, Out=(Head, Key), key=k_k)
        proj_v = hnn.Linear.init(In=Embed, Out=(Head, Key), key=k_v)
        proj_answer = hnn.Linear.init(In=(Head, Key), Out=Embed, key=k_ans)
        return Attention(proj_q, proj_k, proj_v, proj_answer)

    def __call__(self, x, mask=None):
        q = self.proj_q(x)
        k = self.proj_k(x).rename(["position": "key_position"])
        v = self.proj_v(x).rename(["position": "key_position"])
        answers = attention(Key, KPos, q, k, v, causal_mask)
        x = self.proj_answer(answers)
        return x

3. 应用案例和最佳实践

应用案例：大规模语言模型训练

Haliax 被用于训练大规模语言模型，如 GPT-3 等。通过 Haliax 的命名张量和分布式训练支持，研究人员能够更高效地管理和训练这些模型。

最佳实践：分布式训练

在分布式训练中，Haliax 支持全分片数据并行（FSDP）和张量并行。以下是一个简单的分布式训练设置示例：

import haliax.distributed as hdist

# 初始化分布式环境
hdist.initialize()

# 定义模型和数据
model = Attention.init(Embed, Head, Key, key=jax.random.PRNGKey(0))
data = ...

# 分布式训练循环
for batch in data:
    model = hdist.pmap(model)(batch)