在多服务器环境中部署LLaMA 3.1 405B模型

# LLaMA 3.1 405B 模型部署指南

本指南提供了在三台服务器(100.10.128.1、100.10.128.2、100.10.128.3)上部署LLaMA 3.1 405B模型的详细步骤。

## 1. 安装NVIDIA驱动和CUDA

在所有三台服务器上执行以下步骤:

```bash
# 更新系统
sudo apt update && sudo apt upgrade -y

# 安装必要的依赖
sudo apt install gcc make dkms linux-headers-$(uname -r)

# 下载并安装NVIDIA驱动（适用于V100的470版本）
wget https://us.download.nvidia.com/tesla/470.103.01/NVIDIA-Linux-x86_64-470.103.01.run
sudo sh NVIDIA-Linux-x86_64-470.103.01.run

# 安装CUDA 11.3（与PyTorch 1.12兼容）
wget https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda_11.3.0_465.19.01_linux.run
sudo sh cuda_11.3.0_465.19.01_linux.run

# 设置环境变量
echo 'export PATH=/usr/local/cuda-11.3/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

2. 安装Anaconda和创建虚拟环境

在所有服务器上执行:

# 下载并安装Anaconda
wget https://repo.anaconda.com/archive/Anaconda3-2023.03-Linux-x86_64.sh
bash Anaconda3-2023.03-Linux-x86_64.sh

# 创建虚拟环境
conda create -n llama python=3.9
conda activate llama

3. 安装PyTorch和其他必要的库

在所有服务器上执行:

# 安装PyTorch（确保与CUDA 11.3兼容）
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

# 安装其他必要的库
pip install transformers deepspeed ray

4. 安装NCCL（用于多GPU通信）

在所有服务器上执行:

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.3.0/local_installers/cuda-repo-ubuntu2004-11-3-local_11.3.0-465.19.01-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-3-local_11.3.0-465.19.01-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-3-local/7fa2af80.pub
sudo apt-get update
sudo apt-get install libnccl2 libnccl-dev

5. 下载LLaMA 3.1 405B模型

创建以下download.sh脚本:

#!/bin/bash

# 设置变量
MODEL_SIZE="405B"
BASE_URL="https://example.com/llama-3.1-405b"  # 替换为实际的下载URL
OUTPUT_DIR="/path/to/llama-3.1-405b"  # 替换为您想要保存模型的路径

# 创建输出目录
mkdir -p $OUTPUT_DIR

# 下载模型文件
echo "Downloading LLaMA 3.1 ${MODEL_SIZE} model files..."

# 假设模型被分成了多个分片
for i in {0..7}  # 假设有8个分片,根据实际情况调整
do
    wget -c "${BASE_URL}/model_part_${i}.bin" -O "${OUTPUT_DIR}/model_part_${i}.bin"
done

# 下载tokenizer文件
wget -c "${BASE_URL}/tokenizer.model" -O "${OUTPUT_DIR}/tokenizer.model"

# 下载配置文件
wget -c "${BASE_URL}/config.json" -O "${OUTPUT_DIR}/config.json"

echo "Download completed. Files saved in ${OUTPUT_DIR}"

# 验证下载
echo "Verifying downloaded files..."
for file in ${OUTPUT_DIR}/*
do
    if [ -f "$file" ]; then
        echo "$(basename "$file") - $(du -h "$file" | cut -f1)"
    fi
done

echo "Download and verification complete."

使用说明:

1. 将此脚本保存为download.sh。
2. 替换BASE_URL和OUTPUT_DIR为实际值。
3. 给脚本添加执行权限: chmod +x download.sh
4. 运行脚本: ./download.sh

6. 启动Ray

在主节点(100.10.128.1)上:

ray start --head --port=6379 --num-gpus=8

在其他两个节点(100.10.128.2 和 100.10.128.3)上:

ray start --address='100.10.128.1:6379' --num-gpus=8

7. 创建部署脚本

创建deploy_llama.py:

import os
import ray
import torch
import deepspeed
from transformers import AutoTokenizer, AutoModelForCausalLM

ray.init(address="auto")

@ray.remote(num_gpus=8)
def load_model_shard(model_path, shard_id, num_shards):
    ds_config = {
        "fp16": {"enabled": True},
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {"device": "cpu", "pin_memory": True},
            "offload_param": {"device": "cpu", "pin_memory": True},
            "overlap_comm": True,
            "contiguous_gradients": True,
            "sub_group_size": 1e9
        },
        "train_batch_size": 1,
        "train_micro_batch_size_per_gpu": 1,
    }

    with deepspeed.OnDevice(dtype=torch.float16, device="cuda"):
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            device_map="auto",
            offload_folder="offload",
        )
    
    model_engine, _, _, _ = deepspeed.initialize(model=model, config=ds_config)
    return model_engine

def main():
    model_path = "/path/to/llama-3.1-405b"  # 替换为实际路径
    num_shards = 3

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model_shards = ray.get([load_model_shard.remote(model_path, i, num_shards) for i in range(num_shards)])

    input_text = "Hello, how are you?"
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to('cuda')
    
    with torch.no_grad():
        output = model_shards[0].generate(input_ids, max_length=50)
    
    print(tokenizer.decode(output[0]))

if __name__ == "__main__":
    main()

8. 运行部署脚本

在主节点(100.10.128.1)上运行:

python deploy_llama.py

注意事项

• 确保所有服务器上的防火墙设置允许Ray所需的端口通信。
• 405B模型非常大，可能需要更复杂的并行策略和内存优化。
• 监控GPU内存使用情况，可能需要进一步优化以充分利用可用资源。
• 确保网络带宽足够，因为会有大量数据在节点间传输。
• 这个设置是基础的，可能需要根据实际情况进行进一步的优化和调整。

这个Markdown文档提供了完整的步骤,从环境设置到模型下载和部署。您可以根据需要进一步调整或扩展这个文档。