Alauda AI
  • 简体中文

    • Kubeflow Trainer 快速入门

    目录

    背景准备运行时镜像准备 ClusterTrainingRuntime运行示例 Notebook总结

    背景

    Kubeflow Trainer v2 是 Kubeflow 的一个组件，简化了在 Kubernetes 上运行分布式机器学习训练作业的流程。它提供了一种标准化的方式来定义训练运行时和作业，支持多种框架，如 PyTorch、Transformers、TensorFlow 等。在 Alauda AI 中，Kubeflow Trainer v2 与平台的 notebook 环境无缝集成，允许用户直接从开发工作区提交和管理训练作业。

    本快速入门指南演示如何使用 Kubeflow Trainer v2 搭建分布式 PyTorch 训练环境。您将学习如何构建自定义运行时镜像、配置 ClusterTrainingRuntime，并运行一个用于 MNIST 分类的示例训练作业。该设置支持在 GPU 集群上高效进行分布式训练，充分利用 Alauda AI 的资源管理和安全特性。

    准备运行时镜像

    创建一个名为 torch_distributed.Containerfile 的文件，内容如下，并构建镜像。或者您也可以使用预构建镜像 alaudadockerhub/torch-distributed:v2.9.1-aml2。

    FROM python:3.13-trixie
ARG USERNAME=appuser
ARG USER_UID=1000
ARG USER_GID=$USER_UID

RUN sed -i 's/deb.debian.org/mirrors.ustc.edu.cn/g' /etc/apt/sources.list.d/debian.sources
RUN apt-get update && \
apt-get install -y build-essential

RUN pip install --no-cache-dir -i https://pypi.tuna.tsinghua.edu.cn/simple -U pip && \
pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cu126 \
"torch==2.9.1" \
"torchvision==0.24.1"

RUN groupadd --gid $USER_GID $USERNAME \
    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME

WORKDIR /workspace
RUN chown $USERNAME:$USER_GID /workspace

    准备 ClusterTrainingRuntime

    创建一个 kf-torch-distributed.yaml 文件，添加一个 ClusterTrainingRuntime 配置，以便在 Alauda AI 上启动分布式 PyTorch TrainJob。然后以管理员身份运行 kubectl apply -f kf-torch-distributed.yaml 进行创建。

    注意：默认的 ClusterTrainingRuntime 已根据 Alauda AI 的默认安全设置进行了修改。

    apiVersion: trainer.kubeflow.org/v1alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-distributed
  labels:
    trainer.kubeflow.org/framework: torch
spec:
  mlPolicy:
    numNodes: 1
    torch:
      numProcPerNode: auto
  template:
    spec:
      replicatedJobs:
        - name: node
          template:
            metadata:
              labels:
                trainer.kubeflow.org/trainjob-ancestor-step: trainer
            spec:
              template:
                spec:
                  securityContext:
                    runAsNonRoot: true
                    runAsUser: 1000
                    runAsGroup: 1000
                    fsGroup: 1000
                  volumes:
                  - emptyDir:
                      medium: Memory
                      # 这里可以调整共享内存的大小
                      sizeLimit: 2Gi
                    name: dshm
                  # EmptyDir 作为训练的工作空间，并共享输出模型到最终步骤。
                  - name: workspace
                    emptyDir: {}
                  containers:
                    - name: node
                      # 在此处使用上一步构建的镜像或我们的预构建镜像。
                      image: alaudadockerhub/torch-distributed:v2.9.1-aml2
                      env:
                      - name: TORCH_HOME
                        value: /tmp/torch_cache
                      - name: TORCH_EXTENSIONS_DIR
                        value: /tmp/torch_extensions
                      - name: TRITON_CACHE_DIR
                        value: /tmp/triton_cache
                      volumeMounts:
                        - name: workspace
                          mountPath: /workspace
                        - name: dshm
                          mountPath: /dev/shm
                      securityContext:
                        allowPrivilegeEscalation: false
                        capabilities:
                          drop:
                            - ALL
                        runAsNonRoot: true
                        seccompProfile:
                          type: RuntimeDefault

    运行示例 Notebook

    注意：运行以下示例 notebook 需要联网，因为需要安装 python 包和下载数据集。

    从 github workbench howtos 下载 kubeflow_trainer_mnist.ipynb，并将文件拖拽到您的 notebook 实例中。按照 notebook 中的指导，使用 PyTorch 启动一个 TrainJob。

    有关如何使用 Kubeflow Trainer v2 的更多信息，请参阅 Kubeflow Document。

    总结

    通过本快速入门指南，您已成功在 Alauda AI 环境中搭建了 Kubeflow Trainer v2，并运行了一个分布式 PyTorch 训练作业。此基础使您能够高效地在多个节点和 GPU 上扩展机器学习工作负载。

    后续步骤：

    • 通过修改示例 notebook，尝试不同的模型和数据集。
    • 探索自定义指标、超参数调优以及与 MLflow 集成进行实验跟踪等高级功能。
    • 针对 TensorFlow 或自定义训练脚本调整 ClusterTrainingRuntime。

    有关更详细的文档和高级配置，请参阅 Kubeflow Trainer v2 文档。