Skip to content

在开发机中构建 LLaMA-Factory 训练镜像

提示

本文中的平台预置镜像地址是示例。使用前,请在镜像中心或当前实例的预置镜像列表中确认当前可用区提供了对应仓库及 tag;如果列表中不存在,请改用页面显示的可用镜像地址。替换 Registry 域名并不能保证相同仓库或 tag 在其他环境或可用区中存在。

本页介绍如何在开启 Docker 容器功能的开发机中构建 LLaMA-Factory 训练镜像。推荐先在开发机中克隆源码、固定版本、调整配置并完成最小验证,再用 Dockerfile 将已经确认可用的源码和依赖固化成镜像。

该方式适合正式训练、PyTorch DDP、DeepSpeed、FlashAttention,以及需要编译 CUDA 扩展的 LLaMA-Factory 环境。平台镜像中心也支持粘贴 Dockerfile 构建镜像,但无法访问开发机中的源码目录,不适合作为 LLaMA-Factory 这类源码型训练环境的主路径。

准备开发机和构建目录

本页的目标是在开发机中完成 LLaMA-Factory 源码准备、Dockerfile 构建、镜像检查和推送。创建或改配开发机时,先确认以下条件:

  • 开启「Docker 容器」功能。未开启时,开发机内无法执行 docker buildx builddocker rundocker push
  • 准备 BuildKit/Buildx。本文命令使用 docker buildx build,首次使用前按在开发机中使用 Dockerfile 构建镜像准备 BuildKit。
  • 预留足够存储空间。CUDA、PyTorch、LLaMA-Factory 依赖和镜像层会占用较多空间;空间不足时,构建可能在下载或安装依赖阶段失败。
  • 如需执行 docker run --gpus all 检查镜像,开发机需要挂载 GPU。只构建和推送镜像本身不需要使用 GPU。

完整 LLaMA-Factory 训练镜像建议包含系统级 CUDA 12.8.1 与 cuDNN 9。系统级 CUDA 提供 nvcc、CUDA 头文件和 CUDA 动态库,用于支持 DeepSpeed CUDA op、FlashAttention、Apex 或自定义 CUDA kernel 的编译。

准备这层 CUDA 环境有两种方式:

  • 使用镜像中心构建 CUDA 基础镜像。按系统级 CUDA 12.8.1 与 cuDNN 9 镜像操作即可;这种方式可能耗时较长,但不需要创建开发机,也不会占用 GPU。
  • 在开发机中完成全部构建。继续按本页操作即可;本页 Dockerfile 支持复用已有 CUDA 基础镜像,也支持在开发机中从零构建 cuda-base target,再构建 LLaMA-Factory 训练镜像。

拉取并确认 LLaMA-Factory 源码

在开发机中创建构建目录,并拉取 LLaMA-Factory 源码:

language-shell
mkdir -p ~/build-llamafactory
cd ~/build-llamafactory

git clone https://github.com/hiyouga/LLaMA-Factory.git
cd LLaMA-Factory
git checkout main

如果 git clone 速度很慢或无法访问,通常是中国大陆网络环境访问 GitHub 受限。平台不提供网络代理服务,请自行处理网络访问问题。可选做法之一是通过 VS Code Remote - SSH 将本地代理共享到开发机,见共享本地代理服务至远程主机

正式镜像建议 checkout 到已经验证过的 tag 或 commit,而不是长期使用 main

language-shell
git checkout YOUR_VERIFIED_TAG_OR_COMMIT

如果训练需要自定义数据集、插件、模板或内部补丁,在开发机中完成修改后再构建镜像。常见修改包括:

  • 修改 data/dataset_info.json
  • 添加团队内部数据处理或模型插件。
  • 固定 LLaMA-Factory 代码版本。
  • 准备可复用的训练配置文件。

选择 CUDA 基础镜像来源

先判断团队是否已经有验证过的 CUDA 12.8.1 与 cuDNN 9 基础镜像。已有基础镜像时,后续构建只需要安装 LLaMA-Factory 训练环境;没有基础镜像时,再使用 Dockerfile 内置的 cuda-base target 从零构建。

已有 CUDA 基础镜像时直接复用

如果已有基础镜像,例如 cr.infini-ai.com/YOUR_TENANT/cuda-12.8.1-cudnn-9x:v1,后续构建时通过 CUDA_BASE_IMAGE 指定它。这样可以跳过 CUDA 和 cuDNN 安装步骤,减少构建时间,也能复用团队已经验证过的系统 CUDA 环境。

复用前建议确认基础镜像满足以下条件:

  • 包含 nvcc、CUDA 头文件和 CUDA 动态库。
  • 包含 cuDNN 9。
  • 已推送到训练任务能访问的租户镜像仓库,并且在训练任务创建页可以选到。
  • CUDA Toolkit 版本不高于当前 GPU 驱动支持上限。

信息

如果团队已经验证过 CUDA 12.x 兼容行为,也可以在较低 nvidia-smi CUDA 版本的驱动上运行 CUDA 12.8 镜像。使用这种方式前,先按后文检查 PyTorch CUDA、GPU 矩阵计算和可选 FlashAttention kernel。

没有 CUDA 基础镜像时使用内置 target

如果还没有基础镜像,直接使用本页 Dockerfile 中的 cuda-base target。该 target 会基于 Ubuntu 24.04 安装 CUDA 12.8.1 与 cuDNN 9,后续 llamafactory-train target 会在它之上继续安装 LLaMA-Factory 训练环境。

编写多 target 训练镜像 Dockerfile

回到构建目录,创建 Dockerfile:

language-shell
cd ~/build-llamafactory
touch llamafactory.dockerfile

写入以下内容。该 Dockerfile 包含两个 target,并用 CUDA_BASE_IMAGE 决定 llamafactory-train 从哪里继承 CUDA 环境:

  • cuda-base:安装系统级 CUDA 12.8.1 与 cuDNN 9。
  • llamafactory-train:在 cuda-base 上安装 uv、PyTorch、LLaMA-Factory、DeepSpeed 和可选 FlashAttention。

该 Dockerfile 将 LLaMA-Factory 源码放在 /app,与官方 LLaMA-Factory 镜像和使用 LLaMA-Factory 进行训练中的启动脚本保持一致。

language-dockerfile
ARG CUDA_BASE_IMAGE=cuda-base

FROM cr.infini-ai.com/infini-ai/ubuntu:24.04 AS cuda-base

ENV PATH=/usr/local/cuda/bin:/usr/local/cuda-12.8/bin${PATH:+:${PATH}} \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.8/lib64:/usr/lib/x86_64-linux-gnu${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} \
    CUDA_HOME=/usr/local/cuda \
    CPATH=/usr/include${CPATH:+:${CPATH}} \
    LIBRARY_PATH=/usr/lib/x86_64-linux-gnu${LIBRARY_PATH:+:${LIBRARY_PATH}}

RUN apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y wget build-essential \
    && rm -rf /var/lib/apt/lists/* \
    && wget https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda_12.8.1_570.124.06_linux.run \
    && sh cuda_12.8.1_570.124.06_linux.run --toolkit --silent --override \
    && echo '/usr/local/cuda/lib64' >> /etc/ld.so.conf \
    && echo '/usr/local/cuda-12.8/lib64' >> /etc/ld.so.conf \
    && ldconfig \
    && rm cuda_12.8.1_570.124.06_linux.run

RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb \
    && dpkg -i cuda-keyring_1.1-1_all.deb \
    && apt-get update \
    && apt-get install -y cudnn-cuda-12 \
    && echo '/usr/lib/x86_64-linux-gnu' >> /etc/ld.so.conf \
    && ldconfig \
    && rm cuda-keyring_1.1-1_all.deb \
    && rm -rf /var/lib/apt/lists/*

FROM ghcr.io/astral-sh/uv:0.5.31 AS uv

FROM ${CUDA_BASE_IMAGE} AS llamafactory-train

COPY --from=uv /uv /uvx /bin/

ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
ARG PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple
ARG INSTALL_FLASH_ATTN=false
ARG FLASH_ATTN_VERSION=2.8.3
ARG MAX_JOBS=4

ENV DEBIAN_FRONTEND=noninteractive \
    VIRTUAL_ENV=/opt/llamafactory-env \
    PATH=/opt/llamafactory-env/bin:$PATH \
    UV_INDEX_URL=${PIP_INDEX_URL} \
    UV_LINK_MODE=copy \
    UV_COMPILE_BYTECODE=1 \
    PIP_ROOT_USER_ACTION=ignore \
    MAX_JOBS=${MAX_JOBS}

WORKDIR /app

RUN apt-get update && apt-get install -y \
    python3 python3-venv python3-dev git ca-certificates ninja-build \
    && rm -rf /var/lib/apt/lists/*

RUN printf '%s\n' \
    'torch==2.8.0' \
    'torchvision==0.23.0' \
    'torchaudio==2.8.0' \
    > /tmp/llamafactory-constraints.txt

RUN --mount=type=cache,target=/root/.cache/uv \
    uv venv "$VIRTUAL_ENV" && \
    uv pip install \
        torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 \
        --index-url "$PYTORCH_INDEX_URL"

COPY LLaMA-Factory/ /app/

RUN --mount=type=cache,target=/root/.cache/uv \
    cd /app && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt -e . && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt tensorboard && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt -r requirements/metrics.txt && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt -r requirements/deepspeed.txt && \
    if [ "$INSTALL_FLASH_ATTN" = "true" ]; then \
        uv pip install --constraint /tmp/llamafactory-constraints.txt packaging psutil ninja setuptools wheel && \
        uv pip install --constraint /tmp/llamafactory-constraints.txt --no-build-isolation flash-attn=="$FLASH_ATTN_VERSION"; \
    fi && \
    uv pip check

RUN llamafactory-cli version && \
    python - <<'PY'
import torch

print("torch:", torch.__version__)
print("cuda:", torch.version.cuda)
print("nccl:", torch.cuda.nccl.version())
print("cuda available:", torch.cuda.is_available())
PY

WORKDIR /app

使用已有 CUDA 基础镜像构建训练镜像

如果已经有验证过的 CUDA 基础镜像,直接构建 llamafactory-train target,并用 CUDA_BASE_IMAGE 指向该基础镜像:

language-shell
cd ~/build-llamafactory

docker buildx build \
  -f llamafactory.dockerfile \
  --target llamafactory-train \
  --build-arg CUDA_BASE_IMAGE=cr.infini-ai.com/YOUR_TENANT/cuda-12.8.1-cudnn-9x:v1 \
  -t llamafactory-train:cuda12.8-torch2.8 \
  --load .

构建时,Docker 会从 CUDA_BASE_IMAGE 指定的镜像开始构建训练环境,不会重新执行 cuda-base target 中的 CUDA 和 cuDNN 安装步骤。

从零构建 CUDA target 后再构建训练镜像

如果没有可复用的 CUDA 基础镜像,先构建 cuda-base target,确认系统级 CUDA 和 cuDNN 安装成功:

language-shell
cd ~/build-llamafactory

docker buildx build \
  -f llamafactory.dockerfile \
  --target cuda-base \
  -t llamafactory-cuda-base:12.8.1-cudnn9 \
  --load .

在开发机中检查 CUDA 编译器和 cuDNN 库。以下命令使用具名容器和 docker logs,避免部分开发机中 docker run 实时输出不显示的问题:

language-shell
docker rm -f llamafactory-cuda-check 2>/dev/null || true

docker run --name llamafactory-cuda-check --gpus all \
  llamafactory-cuda-base:12.8.1-cudnn9 \
  bash -lc '
    set -e
    nvcc --version
    ldconfig -p | grep libcudnn
  '

docker logs llamafactory-cuda-check
docker inspect llamafactory-cuda-check --format 'ExitCode={{.State.ExitCode}} Error={{.State.Error}}'
docker rm llamafactory-cuda-check

CUDA target 通过后,再构建完整训练镜像。未指定 CUDA_BASE_IMAGE 时,llamafactory-train 会继承 Dockerfile 内置的 cuda-base target:

language-shell
docker buildx build \
  -f llamafactory.dockerfile \
  --target llamafactory-train \
  -t llamafactory-train:cuda12.8-torch2.8 \
  --load .

按需打开 FlashAttention

如果训练配置会使用 FlashAttention,可在构建训练镜像时打开 INSTALL_FLASH_ATTN

language-shell
docker buildx build \
  -f llamafactory.dockerfile \
  --target llamafactory-train \
  --build-arg CUDA_BASE_IMAGE=cr.infini-ai.com/YOUR_TENANT/cuda-12.8.1-cudnn-9x:v1 \
  --build-arg INSTALL_FLASH_ATTN=true \
  --build-arg FLASH_ATTN_VERSION=2.8.3 \
  -t llamafactory-train:cuda12.8-torch2.8-fa \
  --load .

如果使用 Dockerfile 内置的 cuda-base target,去掉命令中的 --build-arg CUDA_BASE_IMAGE=...

FlashAttention 是 CUDA/C++ 编译型扩展。打开 INSTALL_FLASH_ATTN 后,镜像构建会安装或编译 flash-attn;如果 Python、PyTorch、CUDA Toolkit、GPU 架构或编译环境不匹配,构建可能失败。本文 Dockerfile 已固定 PyTorch 2.8.0、CUDA 12.8.1,并显式安装 packagingpsutilninjasetuptoolswheel,以提高构建成功率。构建完成后,按后文运行 FlashAttention kernel 检查,不要只检查 import flash_attn

FlashAttention-2 通常用于 Ampere、Ada、Hopper 等 GPU 架构;Hopper 上的 FlashAttention-3 对 CUDA 版本要求更高。使用其他 GPU 架构或更换 FlashAttention 版本前,先查看 FlashAttention 官方安装要求

默认 MAX_JOBS=4,用于降低编译时内存不足的概率。如果开发机内存充足,可在构建命令中追加 --build-arg MAX_JOBS=8 或更高值加快编译。只有训练配置明确会使用 FlashAttention 时,才建议打开 INSTALL_FLASH_ATTN。如果团队已经有验证过的 FlashAttention wheel,也可以把 Dockerfile 中的安装命令替换为内部制品库地址。

检查 LLaMA-Factory 训练镜像

构建完成后,在开发机中检查镜像。以下命令使用具名容器和 docker logs;如果直接执行 docker run --rm ... 没有显示输出,见操作容器

language-shell
docker rm -f llamafactory-smoke-test 2>/dev/null || true

docker run --name llamafactory-smoke-test --gpus all \
  llamafactory-train:cuda12.8-torch2.8 \
  bash -lc '
    set -e
    python - <<'"'"'PY'"'"'
import torch

print("torch:", torch.__version__)
print("cuda:", torch.version.cuda)
print("nccl:", torch.cuda.nccl.version())
print("cuda available:", torch.cuda.is_available())
print("device:", torch.cuda.get_device_name(0))
print("capability:", torch.cuda.get_device_capability(0))

x = torch.randn(2048, 2048, device="cuda", dtype=torch.float16)
y = x @ x
torch.cuda.synchronize()
print("matmul ok:", y.float().mean().item())
PY
    uv pip show nvidia-nccl-cu12
    uv pip check
    llamafactory-cli version
    nvcc --version
  '

docker logs llamafactory-smoke-test
docker inspect llamafactory-smoke-test --format 'ExitCode={{.State.ExitCode}} Error={{.State.Error}}'
docker rm llamafactory-smoke-test

如果 torch.version.cuda 显示为 12.8cuda availableTrue,且 GPU 矩阵计算通过,说明 PyTorch CUDA 运行时可以在当前驱动和 GPU 上工作。若 nvidia-nccl-cu12 版本不低于 2.24,镜像也满足平台 NCCL RAS 诊断的版本前提。NCCL 状态页仍需任务实际初始化 NCCL communicator 后才有可用信息。

如果构建时打开了 INSTALL_FLASH_ATTN=true,继续运行 FlashAttention kernel 检查:

language-shell
docker rm -f llamafactory-fa-kernel-test 2>/dev/null || true

docker run --name llamafactory-fa-kernel-test --gpus all \
  llamafactory-train:cuda12.8-torch2.8-fa \
  bash -lc '
    set -e
    python - <<'"'"'PY'"'"'
import torch
from flash_attn import flash_attn_func

print("torch:", torch.__version__)
print("cuda:", torch.version.cuda)
print("device:", torch.cuda.get_device_name(0))
print("capability:", torch.cuda.get_device_capability(0))

batch = 2
seqlen = 128
nheads = 8
headdim = 64

q = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.float16)
k = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.float16)
v = torch.randn(batch, seqlen, nheads, headdim, device="cuda", dtype=torch.float16)

out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)
torch.cuda.synchronize()

print("output shape:", tuple(out.shape))
print("output dtype:", out.dtype)
print("FlashAttention kernel OK")
PY
  '

docker logs llamafactory-fa-kernel-test
docker inspect llamafactory-fa-kernel-test --format 'ExitCode={{.State.ExitCode}} Error={{.State.Error}}'
docker rm llamafactory-fa-kernel-test

推送镜像到租户镜像仓库

确认镜像可用后,将镜像推送到开发机所在可用区的租户镜像仓库。以下命令中的仓库地址和镜像名称请替换为实际值:

警告

以下命令会把镜像写入当前实例所在可用区的 Registry。如果目标是其他可用区,请不要使用平台内部 Registry 域名;请改用镜像中心提供的目标可用区 Registry 公网访问地址和登录凭证。

如需确认平台内部 Registry 域名的登录方式并获取当前可用区的租户镜像仓库根路径,参见在开发机中访问镜像仓库

language-shell
export TENANT_ID="te-xxxxxxxxxxxxxxxx"
export CURRENT_ZONE_REPO_ROOT="cr.infini-ai.com/$TENANT_ID"
docker tag llamafactory-train:cuda12.8-torch2.8 \
  $CURRENT_ZONE_REPO_ROOT/llamafactory-train:cuda12.8-torch2.8

docker push $CURRENT_ZONE_REPO_ROOT/llamafactory-train:cuda12.8-torch2.8

推送完成后,在创建训练任务时选择该镜像。创建任务和填写启动命令的步骤见使用 LLaMA-Factory 进行训练

只运行基础训练时构建轻量镜像

轻量镜像适合基础 SFT、LoRA、QLoRA 或 PyTorch DDP 任务。它依赖 PyTorch CUDA wheel 提供运行时库,不包含 nvcc 和 CUDA 头文件。

如果任务会使用 DeepSpeed CUDA op、FlashAttention、Apex 或自定义 CUDA kernel,不要使用轻量镜像。

轻量镜像的 Dockerfile 与完整镜像类似,但基础镜像可以改为 Ubuntu 24.04,并且不安装 requirements/deepspeed.txt

language-dockerfile
FROM ghcr.io/astral-sh/uv:0.5.31 AS uv

FROM cr.infini-ai.com/infini-ai/ubuntu:24.04

COPY --from=uv /uv /uvx /bin/

ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
ARG PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple

ENV DEBIAN_FRONTEND=noninteractive \
    VIRTUAL_ENV=/opt/llamafactory-env \
    PATH=/opt/llamafactory-env/bin:$PATH \
    UV_INDEX_URL=${PIP_INDEX_URL} \
    UV_LINK_MODE=copy \
    UV_COMPILE_BYTECODE=1 \
    PIP_ROOT_USER_ACTION=ignore

WORKDIR /app

RUN apt-get update && apt-get install -y \
    python3 python3-venv python3-dev git build-essential ca-certificates \
    && rm -rf /var/lib/apt/lists/*

RUN printf '%s\n' \
    'torch==2.8.0' \
    'torchvision==0.23.0' \
    'torchaudio==2.8.0' \
    > /tmp/llamafactory-constraints.txt

RUN --mount=type=cache,target=/root/.cache/uv \
    uv venv "$VIRTUAL_ENV" && \
    uv pip install \
        torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 \
        --index-url "$PYTORCH_INDEX_URL"

COPY LLaMA-Factory/ /app/

RUN --mount=type=cache,target=/root/.cache/uv \
    cd /app && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt -e . && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt tensorboard && \
    uv pip install --constraint /tmp/llamafactory-constraints.txt -r requirements/metrics.txt && \
    uv pip check

WORKDIR /app

在镜像中心构建时使用源码仓库地址

如果只能使用镜像中心粘贴 Dockerfile 构建镜像,可以在 Dockerfile 中执行 git clone 拉取 LLaMA-Factory 源码。但这种方式不适合需要修改源码、添加数据集配置、反复验证依赖或运行最小训练的场景。

镜像中心构建时也需要能访问 github.com。如果访问 GitHub 速度很慢或无法访问,通常是中国大陆网络环境访问 GitHub 受限;请自行将 LLAMA_FACTORY_REPO 改为可访问的源码仓库地址,或回到开发机构建路径并自行配置网络访问。

使用镜像中心构建时,至少固定源码版本:

language-dockerfile
ARG LLAMA_FACTORY_REPO=https://github.com/hiyouga/LLaMA-Factory.git
ARG LLAMA_FACTORY_REF=YOUR_VERIFIED_TAG_OR_COMMIT

RUN git clone "$LLAMA_FACTORY_REPO" /app && \
    cd /app && \
    git checkout "$LLAMA_FACTORY_REF"

长期训练环境仍建议回到开发机中构建镜像。这样可以在构建前直接检查源码、依赖、数据集配置和最小训练路径,把问题提前暴露在正式训练任务之前。