rocm-pypi安装加速

ROCm PyPI 安装加速实践

在 AMD GPU 上跑深度学习时，安装带 ROCm 支持的 PyTorch 往往是第一步，也是最容易卡住的一步。一个 `torch` + `torchvision` + `torchaudio` 的 ROCm 组合动辄数 GB，默认从境外源拉取，速度可能只有几百 KB/s，中途断流、校验失败也时有发生。本文梳理几条可落地的加速思路，帮助把安装时间从“泡杯咖啡等半天”压到可接受范围。

一、先搞清楚“慢”在哪里

ROCm 版 PyTorch 并不在默认的 PyPI 上，而是分布在几个专用索引里：

• PyTorch 官方 ROCm 索引：`https://download.pytorch.org/whl/rocm<版本>`，例如 `rocm6.1`、`rocm6.2`。
• AMD 官方 PyPI：AMD 为 ROCm 生态提供的 wheel 索引（不同发行版路径会调整，务必以官方文档为准）。
• 基础依赖：`numpy`、`filelock`、`sympy` 等仍来自 PyPI 或镜像站。

慢的原因通常是三点：服务器在境外、单文件体积大、解析过程反复探测多版本。加速就要分别针对这三点下手。

二、镜像站加速

国内多个镜像站对 PyTorch 及常用包做了同步，PyPI 主源可换成： bash pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple pip config set global.trusted-host pypi.tuna.tsinghua.edu.cn 常见可选：清华 TUNA、中科大 USTC、阿里云、上海交大。它们对 `torch` 的 CPU/CUDA 版同步较好，但 ROCm wheel 通常不在镜像范围内，仍需回源 `download.pytorch.org`。

因此更实用的组合是： bash pip install torch torchvision torchaudio \ --index-url https://pypi.tuna.tsinghua.edu.cn/simple \ --extra-index-url https://download.pytorch.org/whl/rocm6.2 \ --find-links https://download.pytorch.org/whl/rocm6.2/torch 这样基础依赖走国内镜像，只有 ROCm 大头回源，整体带宽占用更集中在关键包上。

三、pip 参数调优

几个直接影响体验的参数： bash pip install torch \ --index-url https://download.pytorch.org/whl/rocm6.2 \ --timeout 120 \ --retries 10 \ --progress-bar on \ -v

• `--timeout`：默认 15 秒，大文件场景容易超时误判，调到 60~120。
• `--retries`：网络抖动时自动重试，避免整条命令从头再来。
• `--no-cache-dir`：仅在排查缓存损坏时使用；正常安装不建议关闭缓存。

如果只是版本试错，先用 `--dry-run`（较新 pip 支持）确认解析结果，避免下完几个 GB 才发现版本不匹配。

四、用 uv 替代 pip

`uv` 是 Rust 写的包管理器，解析和下载并发度远高于 pip，对 ROCm 这种大包尤其明显： bash pip install uv uv pip install torch --index-url https://download.pytorch.org/whl/rocm6.2 `uv` 默认多线程下载、全局缓存去重。同一台机器反复重建虚拟环境时，只要缓存命中，第二次安装几乎是秒级。它同时兼容 `pip` 的索引参数，迁移成本很低。

五、预下载 + 本地缓存复用

对多机部署或需要反复重装的环境，把 wheel 先落到本地最省事： bash

只下载不安装，按平台抓取

pip download torch torchvision torchaudio \ --index-url https://download.pytorch.org/whl/rocm6.2 \ -d ./rocm_wheels 之后离线安装： bash pip install --no-index --find-links ./rocm_wheels torch torchvision torchaudio `pip download` 同样走 `--retries`/`--timeout`。如果单个 wheel 仍慢，可以用 `aria2c -x 16 -s 16 <url>` 多线程单独抓那个文件，再放进 `rocm_wheels` 目录，让 `pip` 直接复用。

共享缓存也是思路之一：在 CI 或多用户机器上把 `PIP_CACHE_DIR` 指向一块共享盘，命中缓存的安装不再走网络。 bash export PIP_CACHE_DIR=/data/pip-cache

六、容器与内网场景

用 Docker 时，镜像层缓存本身就是加速器：把 `pip install torch...` 放在 Dockerfile 靠前的独立层，依赖不变时重建不会重复下载。

若使用 AMD 官方 ROCm 容器镜像，torch 往往已预装，可省去下载环节，只需确认容器内 `torch.version.hip` 非空、`torch.cuda.is_available()` 为真（ROCm 下该接口沿用 CUDA 命名）。

内网环境建议自建索引代理（如 devpi、Nexus、Artifactory），把 `download.pytorch.org` 的 ROCm 目录一次性镜像进来，再让所有节点指向内