You need to enable JavaScript to run this app.
导航
使用 AI 容器镜像
最近更新时间:2024.07.08 14:30:33首次发布时间:2023.10.11 10:46:58

AI 容器镜像是火山引擎提供的一系列 Docker 镜像,其中预装了常见的 AI 框架、工具以及相关依赖,为用户提供快速启动且高度一致的 AI 开发环境。本文为您介绍 AI 容器镜像的使用方式。

背景信息

在 AI 模型训练和推理的前置流程中,用户通常基于特定的AI框架和工具的组合开展工作(例如:PyTorch 框架、deepspeed 库),在基础环境中进行算法调优、模型注入或者脚本开发后,将重新打包好的镜像提交到训练或者推理平台的任务队列中。

AI 容器镜像用于加速用户 AI 开发、训练和推理环境的搭建,并与 AI 套件中其他工具和产品能力共同为用户提供覆盖 AI 应用完整生命周期的服务。

前提条件

已经基于 VKE 搭建 AI 训练或推理任务的基础环境,例如:集群、VPC、GPU 节点等。

使用限制

VKE AI 容器镜像限定只能在 VPC 内拉取镜像,暂不支持公网拉取。请在火山引擎对应地域的 VKE 集群或者 ECS 实例中拉取。

操作步骤

步骤一:获取 AI 预置镜像地址

火山引擎预置提供一系列 AI 容器镜像,为用户提供快速启动且高度一致的AI开发环境。以下是 AI 容器镜像支持的版本、支持状态以及镜像地址。

说明

  • 镜像地址中的 {地域ID} 需要替换为真实 RegionID,支持使用 AI 容器镜像的地域和 RegionID 参见:地域和可用区
  • AI 预置镜像支持的策略如下:
    • AI 预置镜像中的软件和依赖软件仅限于 NVIDIA 官方已编译好且已正式发布的版本。
    • AI 预置镜像的更新,跟随 NVIDIA 官方的软件版本更新节奏,每个季度将会更新一次镜像,更新后的镜像将使用 NVIDIA 官方正式发布的最新软件版本。
机器学习框架版本其它软件&依赖版本操作系统版本是否仍在支持中镜像地址

-

  • CUDA 11.8.0
  • Python 3.10

Ubuntu 22.04

ai-containers-{地域ID}.cr.volces.com/deeplearning/cuda11.8.0-py3.10-ubuntu22.04:latest

Pytorch 2.0.1

  • CUDA 11.8.0
  • Python 3.10

Ubuntu 22.04

ai-containers-{地域ID}.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest

Tensorflow 2.13.0

  • CUDA 11.8.0
  • Python 3.10

Ubuntu 22.04

ai-containers-{地域ID}.cr.volces.com/deeplearning/tensorflow2.13.0-cuda11.8.0-py3.10-ubuntu22.04:latest

Pytorch 2.0.1

  • Megatron 23.05
  • CUDA 11.8.0
  • Python 3.10

Ubuntu 22.04

ai-containers-{地域ID}.cr.volces.com/deeplearning/megatron23.05-torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest

Pytorch 2.0.1

  • Deepspeed0.10.0
  • CUDA 11.8.0
  • Python 3.10

Ubuntu 22.04

ai-containers-{地域ID}.cr.volces.com/deeplearning/deepspeed0.10.0-torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest

步骤二:拉取 AI 预置镜像启动容器

按需获取 AI 预置容器镜像,并拉取到 VKE 中启动容器镜像,使用此 AI 容器镜像创建基础训练环境。

  1. 登录 容器服务控制台
  2. 在左侧导航栏选择 集群,单击可用于 AI 训练或推理的目标集群名称,进入集群管理页面。
  3. 在左侧导航栏选择 工作负载 > 容器组,进入容器组管理页面。
  4. 单击 使用 YAML创建,使用 Yaml 方式拉取 AI 预置镜像启动容器。Yaml 示例如下:
apiVersion: v1
kind: Pod
metadata:
  name: ai-env         # 自定义 Pod 名称。
  namespace: default   # 按需配置 Pod 所属的命名空间。
spec:
  containers:
  - name: pytorch      # 自定义容器名称。
    image: ai-containers-cn-beijing.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest  # AI 预置镜像及其版本。从步骤一中按需获取预置镜像地址,镜像版本使用 latest 即可。
    command: [ "/bin/bash", "-c", "--" ]
    args: [ "while true; do sleep 3600; done;" ]
    resources:
      limits:
        nvidia.com/gpu: "1" # 按需配置 GPU 显卡数量。

步骤三:开发训练环境

Pod 进入 running 状态后,即可通过 VKE 登录 Pod 控制台,进一步进行环境准备和开发工作,具体内容可根据实际使用场景进行开发。例如:通过开发工具进行算法、脚本开发/置入模型文件。

alt

步骤四:打包镜像并推送镜像仓库

  1. 以 AI 容器镜像为基础镜像,通过 dockerfile 构建定制化的 AI 开发环境。dockerfile 示例如下:
# 所使用预置镜像的地址和版本,与步骤二中拉取的 AI 预置镜像保持一致。
FROM ai-containers-cn-beijing.cr.volces.com/deeplearning/torch2.0.1-cuda11.8.0-py3.10-ubuntu22.04:latest
# 为 Dockerfile 中后续命令设置工作目录。
WORKDIR /home
# 从本地文件系统中导入文件到镜像,例如:模型文件。
COPY chatglm2-6b.tar /chatglm2-6b.tar
# 为镜像预装其他所需的 AI 工具软件,例如:jupyterlab。
RUN pip install "jupyterlab"
# 为镜像环境添加其他所需的环境变量。
ARG VARNAME = "value"
ENV PATH=$VARNAME
  1. 将构建完成的派生镜像推送到火山引擎镜像仓库服务中,详细介绍参见:使用标准版实例拉取和推送镜像

后续操作

至此,以 AI 容器镜像为基础打包的派生镜像已经推送到火山引擎镜像仓库,您可通过 VKE 部署模型训练或推理任务。