Rocky9.2部署llama3.1-8b训练环境-优快云博客

总体流程总览

【无 GPU Linux / 跳板机】
  1. docker run 基础镜像（CPU模式）
  2. 容器内安装 Megatron-LM + 依赖
  3. 拷 tokenizer / 数据 / 脚本
  4. 验证“CPU 可运行”
  5. docker commit 固化镜像
  6. docker save 导出镜像

【有 GPU 服务器】
  7. docker load
  8. docker run --gpus all
  9. 验证 CUDA
  10. 启动训练

一、无 GPU 机器：启动基础容器（关键点）

⚠️ 注意：此阶段不要执行任何 GPU 相关命令

拉取官方镜像

docker pull nvcr.io/nvidia/pytorch:25.04-py3

启动容器

#/bin/bash

export MY_CONTAINER="llama3_megatron"
num=`docker ps -a|grep -w "$MY_CONTAINER$"|wc -l`
echo $num $MY_CONTAINER
if [ 0 -eq $num ];then
docker run --runtime=runc -it \
  -v /home/workspace/megatron:/workspace/megatron \
  -v /home/workspace/dataset:/workspace/dataset \
  -v /home/workspace/checkpoints:/workspace/checkpoints \
  nvcr.io/nvidia/pytorch:25.04-py3
else
  docker start $MY_CONTAINER
  docker exec -ti $MY_CONTAINER /bin/bash
fi

✔ 不加 --gpus
✔ 没 GPU 也没问题
✔ 这是 CPU 模式容器

注意：因为docker默认参数/etc/docker/daemon.json中我设置的是"default-runtime": "nvidia"，所以运行CPU容器时需要指定--runtime为runc,不然会报错

二、容器内：安装 Megatron-LM（一次性）

安装一些常用指令

apt update 
apt install -y \
    vim less wget curl git unzip htop tmux bash-completion \
    ca-certificates numactl hwloc net-tools \
    pciutils iproute2 iputils-ping procps \
    build-essential python3-dev 
# 清空缓存包
apt clean

cd /workspace
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM

安装依赖（CPU 环境完全可装）：

pip install --no-cache-dir \
  sentencepiece \
  regex \
  einops \
  transformers \
  datasets

✔ 不会触发 CUDA
✔ 不依赖 GPU

三、准备 tokenizer（完全离线）

1️⃣ Windows 上准备（一次性）

下载：

git clone https://gitee.com/hf-models/Llama-3.1-8B.git

2️⃣ 拷贝到无 GPU Linux

mkdir -p /workspace/checkpoints/Llama-3.1-8B_tokenizer
cp special_tokens_map.json /workspace/checkpoints/Llama-3.1-8B_tokenizer/
cp tokenizer_config.json /workspace/checkpoints/Llama-3.1-8B_tokenizer/
cp tokenizer.json /workspace/checkpoints/Llama-3.1-8B_tokenizer/

四、准备 fineweb小样本文本（不在 Linux 上联网）

✔ 推荐方案

👉 Windows 上生成文本 → 拷贝

download_fineweb.py

import os
from datasets import load_dataset
import json

# 启用国内镜像（可选）
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 加载 FineWeb 10B 样本（公开，无需 token）
dataset = load_dataset(
    "HuggingFaceFW/fineweb",
    name="sample-10BT",      # 小样本，适合测试
    split="train",
    streaming=True
)

samples = []
count = 0
max_samples = 1000_000

print("正在下载 FineWeb 前 100 万条...")
for sample in dataset:
    # FineWeb 的字段是 "text"
    samples.append({
        "text": sample["text"],
        "meta": {"source": "fineweb-sample-10BT"}
    })
    count += 1
    if count >= max_samples:
        break
    if count % 10_000 == 0:
        print(f"✅ 已获取 {count} 条")

# 保存为 JSONL（兼容 Pile 格式）
with open("fineweb_1000k.jsonl", "w", encoding="utf-8") as f:
    for s in samples:
        f.write(json.dumps(s, ensure_ascii=False) + "\n")

print("🎉 完成！数据已保存到 fineweb_1000k.jsonl")

五、容器内：生成 Megatron indexed dataset（CPU 即可）

cd /workspace/Megatron-LM
mkdir -p /workspace/datasets/

python tools/preprocess_data.py \
  --input /data/pile_raw/pile_sample.txt \
  --output-prefix /data/pile_indexed/pile_llama \
  --tokenizer-type SentencePieceTokenizer \
  --tokenizer-model /data/tokenizer/tokenizer.model \
  --append-eod \
  --workers 4

✔ 不用 GPU
✔ 这是最关键的一步

生成数据集如下

六、容器内：CPU 级自检（不要跑训练）

python - << 'EOF'
import sentencepiece
import torch
print("torch:", torch.__version__)
print("cuda available:", torch.cuda.is_available())
EOF

预期输出：

cuda available: False   ← 正常！

七、无 GPU 机器：固化镜像（核心步骤）

1️⃣ 退出容器

exit

2️⃣ docker commit

docker commit llama3-megatron megatron_blackwell:pytorch_25.04-py3

⚠️ 这个镜像 已经是 GPU-ready 的

GPU 是否可用只取决于运行环境

八、导出镜像（传给 GPU 服务器）

docker save megatron_blackwell:pytorch_25.04-py3 -o megatron_blackwell.tar

九、有 GPU 服务器：加载并运行

docker load -i megatron_blackwwell.tar

启动容器：

docker run -it --rm \
  --gpus all \
  --ipc=host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -v /data:/data \
  megatron_blackwell:pytorch_25.04-py \
  bash

十、有 GPU 服务器：GPU 验证（第一次才做）

nvidia-smi

python - << 'EOF'
import torch
print(torch.cuda.device_count())
EOF

十一、启动 LLaMA-3 训练（示例）

cd /workspace/Megatron-LM

torchrun --nproc_per_node=1 pretrain_gpt.py \
  --num-layers 8 \
  --hidden-size 1024 \
  --num-attention-heads 8 \
  --ffn-hidden-size 4096 \
  --seq-length 1024 \
  --max-position-embeddings 1024 \
  --micro-batch-size 1 \
  --global-batch-size 4 \
  --train-iters 20 \
  --lr 3e-4 \
  --data-path /data/pile_indexed/pile_llama \
  --tokenizer-type SentencePieceTokenizer \
  --tokenizer-model /data/tokenizer/tokenizer.model \
  --save /data/checkpoints \
  --fp16