快速入门：PyTorch

此示例使用 SkyPilot 在 PyTorch 中通过分布式数据并行 (DDP) 训练一个类似 GPT 的模型（灵感来自 Karpathy 的 minGPT）。

我们定义一个 SkyPilot YAML 文件，其中包含资源需求、设置命令以及要运行的命令

# train.yaml

name: minGPT-ddp

resources:
    cpus: 4+
    accelerators: L4:4  # Or A100:8, H100:8

# Optional: upload a working directory to remote ~/sky_workdir.
# Commands in "setup" and "run" will be executed under it.
#
# workdir: .

# Optional: upload local files.
# Format:
#   /remote/path: /local/path
#
# file_mounts:
#   ~/.vimrc: ~/.vimrc
#   ~/.netrc: ~/.netrc

setup: |
    git clone --depth 1 https://github.com/pytorch/examples || true
    cd examples
    git filter-branch --prune-empty --subdirectory-filter distributed/minGPT-ddp
    # SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
    uv pip install -r requirements.txt "numpy<2" "torch==1.12.1+cu113" --extra-index-url https://download.pytorch.org/whl/cu113

run: |
    cd examples/mingpt
    export LOGLEVEL=INFO

    echo "Starting minGPT-ddp training"

    torchrun \
    --nproc_per_node=$SKYPILOT_NUM_GPUS_PER_NODE \
    main.py

提示

在 YAML 中，workdir 和 file_mounts 字段被注释掉了。要了解如何使用它们将本地目录/文件或对象存储桶 (S3, GCS, R2) 挂载到集群中，请参阅同步代码和文件。

提示

SkyPilot 会自动将 SKYPILOT_NUM_GPUS_PER_NODE 环境变量设置为每个节点的 GPU 数量。更多信息请参阅秘钥和环境变量。

然后，启动训练

$ sky launch -c mingpt train.yaml

这将为您调配具有所需资源的最便宜的集群，执行设置命令，然后执行运行命令。

训练作业开始运行后，您可以安全地按下 Ctrl-C 来分离日志输出，作业将继续在远程集群上运行。要停止作业，请使用 sky cancel <cluster_name> <job_id> 命令（请参阅CLI 参考）。

训练完成后，使用常用工具传输文件，例如日志和检查点。

提示

您可以自由地复制粘贴上面的 YAML 并根据您自己的项目进行定制。