Llama 2 JAX

这个项目是Llama 2的JAX实现。该项目有三个目标：

使用JAX实现Llama 2模型，以实现在Google Cloud TPU上的高效训练和推理；
开发一个高质量的代码库，作为使用JAX实现Transformer模型的范例；
通过实现高质量代码库，促进识别各种transformer模型中的常见错误和不一致之处，从而为自然语言处理社区提供有价值的见解。

特性

本项目库的文档发布在GitHub Pages上。

环境设置

本项目至少需要Python 3.11、JAX 0.4.19、PyTorch 2.1.0、Optax 0.1.8.dev0和Transformers 4.35.0.dev0。

PyTorch和Transformers用于测试目的。此外，数据加载器依赖于PyTorch的DataLoader，而性能分析功能需要TensorFlow。

安装Python 3.11

对于Ubuntu用户，你可以参照如何在Ubuntu 22.04上安装Python 3.11来安装Python 3.11。该教程同样适用于Ubuntu 20.04。

创建`venv`

python3.11 -m venv venv
. venv/bin/activate
pip install -U pip
pip install -U wheel

TPU Pod的特殊配置

如果你在TPU pod上运行，你需要将所有其他主机的IP地址放在~/podips.txt中（每行一个IP地址）。此外，你应该确保本地主机可以SSH到自身和文件中列出的所有其他主机。

安装适当版本的JAX

你需要按照JAX的官方GitHub页面上的安装说明进行操作。

安装适当版本的PyTorch

通常，你只需要安装CPU版本的PyTorch，因为我们使用JAX进行大部分计算。然而，值得注意的是，当前代码库的生成过程还没有完全优化。为了加速推理，一个有效的方法是将模型转换回Hugging Face格式，并在PyTorch中运行推理。

要安装PyTorch，你可以参照官方安装指南。

在TPU VM上，通常是：

pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

在TPU Pod上：

./podrun -i -- ~/venv/bin/pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

安装其他依赖

pip install git+https://github.com/huggingface/transformers.git
pip install git+https://github.com/deepmind/optax.git  # https://github.com/google-deepmind/optax/issues/472
pip install -r requirements.txt

在TPU Pod上：

./podrun -i -- ~/venv/bin/pip install git+https://github.com/huggingface/transformers.git
./podrun -i -- ~/venv/bin/pip install git+https://github.com/deepmind/optax.git
./podrun -iw -- ~/venv/bin/pip install -r requirements.txt

下载LLaMA权重

LLaMA 1:

如果你无法获得LLaMA权重，可以使用shawwn/llama-dl下载它们。

mkdir ../llama-weights-original && cd ../llama-weights-original
curl -o- https://raw.githubusercontent.com/shawwn/llama-dl/56f50b96072f42fb2520b1ad5a1d6ef30351f23c/llama.sh | bash
python ../llama-2-jax/venv/lib/python3.11/site-packages/transformers/models/llama/convert_llama_weights_to_hf.py --input_dir ../llama-weights-original --model_size 7B --output_dir ../llama-weights/7B

Llama 2:

你可以从官方网站申请访问Llama权重。申请被批准后，你将自动获得访问Hugging Face Llama 2模型的权限。你可以通过尝试访问Llama 2 7B版本来验证模型是否可访问。

登录Hugging Face CLI

如果你需要使用Llama 2模型，你需要登录Hugging Face CLI：

huggingface-cli login

或者，如果你不使用交互式shell，可以以非交互模式登录：

python -c "from huggingface_hub.hf_api import HfFolder; HfFolder.save_token('<YOUR_HUGGING_FACE_TOKEN>')"

在TPU Pod上：

./podrun -i -- ~/venv/bin/python -c "from huggingface_hub.hf_api import HfFolder; HfFolder.save_token('<YOUR_HUGGING_FACE_TOKEN>')"

转换参数

python scripts/convert_params_runner.py llama1-7B
python scripts/convert_params_runner.py llama2-7B
python scripts/convert_params_runner.py llama2-70B

生成

python generate.py

在TPU pod上，命令为：

./podrun -icw ~/venv/bin/python generate.py

登录W&B

在训练之前，你需要登录W&B：

wandb login <YOUR_WANDB_API_KEY>

下载GSM数据集

我通过在GSM数据集上微调模型来展示训练流程的简单示例。

cd .. && git clone --depth=1 https://github.com/openai/grade-school-math.git

训练

python train.py

在TPU pod上，命令为：

./podrun -icw ~/venv/bin/python train.py

模型配置

B: batch_size
L: seq_len
S: src_seq_len
D: dst_seq_len
C: vocab_size
N: n_layers
K: d_k
V: d_v
H: n_heads_kv
R: n_rep_kv
M: d_model
F: d_ff

名称	参数	C	N	K/V	H	R	M	F
LLaMA 1 7B	6738415616	32000	32	128	32	1	4096	11008
Llama 2 7B	6738415616	32000	32	128	32	1	4096	11008
LLaMA 1 13B		32000	40		40	1	5120
Llama 2 13B	13015864320	32000	40	128	40	1	5120	13824
LLaMA 1 33B		32000	60		52	1	6656
LLaMA 1 65B		32000	80		64	1	8192
Llama 2 70B	68976648192	32000	80	128	8	8	8192	28672

n_params = 2CM + (2N + 1)M + 2NMRHK + 2NMHK + 3NMF

模型架构

LLaMA 1 (7B)

Hugging Face格式：

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=0)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

本项目使用的格式：

model
  embedding: (32000, 4096)
  decoder: decoder_block
    input_norm: (32, 4096)
    attention
      q_proj: (32, 4096, 1, 32, 128)
      k_proj: (32, 4096, 32, 128)
      v_proj: (32, 4096, 32, 128)
      out_proj: (32, 1, 32, 128, 4096)
    post_attn_norm: (32, 4096)
    gate_proj: (32, 4096, 11008)
    up_proj: (32, 4096, 11008)
    down_proj: (32, 11008, 4096)
  norm: (4096)
lm_head: (4096, 32000)

Llama 2 (70B)

Hugging Face 格式：

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 8192, padding_idx=0)
    (layers): ModuleList(
      (0-79): 80 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=8192, out_features=8192, bias=False)
          (k_proj): Linear(in_features=8192, out_features=1024, bias=False)
          (v_proj): Linear(in_features=8192, out_features=1024, bias=False)
          (o_proj): Linear(in_features=8192, out_features=8192, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=8192, out_features=28672, bias=False)
          (up_proj): Linear(in_features=8192, out_features=28672, bias=False)
          (down_proj): Linear(in_features=28672, out_features=8192, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=8192, out_features=32000, bias=False)
)

本项目使用的格式：

model
  embedding: (32000, 8192)
  decoder: decoder_block
    input_norm: (80, 8192)
    attention
      q_proj: (80, 8192, 8, 8, 128)
      k_proj: (80, 8192, 8, 128)
      v_proj: (80, 8192, 8, 128)
      out_proj: (80, 8, 8, 128, 8192)
    post_attn_norm: (80, 8192)
    gate_proj: (80, 8192, 28672)
    up_proj: (80, 8192, 28672)
    down_proj: (80, 28672, 8192)
  norm: (8192)
lm_head: (8192, 32000)

发现

LLaMA 使用旋转位置嵌入。
Q、K、V矩阵和前馈网络中的线性投影没有偏置项，这与原始 Transformer 相同，但与 BERT 和 BART 不同。
在 Llama 模型中，每个前馈网络有 3 个线性投影，而 BART 只有 2 个。
原始 LLaMA 实现中没有 dropout。
原始实现中没有左侧填充，详见 huggingface/transformers#26072。
Llama 2 70B 使用分组查询注意力（GQA）。
许多人以 16 位精度（float16 或 bfloat16）微调 Llama，但这会影响性能，因此与其他以 32 位精度训练的模型进行比较会不公平。另一个值得注意的是，旋转嵌入的参数应始终保持 32 位精度以避免冲突。