Alauda AI
  • 简体中文

    • Kubeflow Trainer 快速入门

    在 Alauda AI 上使用 Kubeflow Trainer v2 进行最小化分布式 PyTorch 训练的配置：自定义运行时镜像、ClusterTrainingRuntime 和一个 MNIST 示例 notebook。

    目录

    运行时镜像ClusterTrainingRuntime运行示例 notebook

    运行时镜像

    使用预构建镜像 alaudadockerhub/torch-distributed:v2.9.1-aml2，或者基于此 torch_distributed.Containerfile 构建你自己的镜像：

    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 && \
    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。该 Pod spec 针对 Alauda AI 的默认 PSA 进行了收紧。

    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:
                    - name: dshm
                      emptyDir:
                        medium: Memory
                        sizeLimit: 2Gi
                    - 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 包并下载 MNIST，因此工作台需要具备出站网络访问能力。

    下载 kubeflow-trainer-mnist.ipynb 并上传到你的工作台，然后按照其中的步骤提交 TrainJob。

    有关 Trainer v2 功能的背景信息，请参阅上游 Kubeflow Trainer 文档。