Skip to content

[ROCm] add rocm wheels #1477

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
May 28, 2021
Merged

[ROCm] add rocm wheels #1477

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
23cf3fb
add rocm binary
micmelesse Apr 26, 2021
df1e434
update config.yml
micmelesse May 27, 2021
8f3231c
remove conda builds
micmelesse May 27, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
172 changes: 172 additions & 0 deletions .circleci/config.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -650,6 +650,13 @@ workflows:
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-cuda111
- binary_linux_wheel:
cuda_version: rocm4.1
name: binary_linux_wheel_py3.6_rocm4.1
python_version: '3.6'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_linux_wheel:
cuda_version: cpu
name: binary_linux_wheel_py3.7_cpu
Expand All @@ -670,6 +677,13 @@ workflows:
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-cuda111
- binary_linux_wheel:
cuda_version: rocm4.1
name: binary_linux_wheel_py3.7_rocm4.1
python_version: '3.7'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_linux_wheel:
cuda_version: cpu
name: binary_linux_wheel_py3.8_cpu
Expand All @@ -690,6 +704,13 @@ workflows:
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-cuda111
- binary_linux_wheel:
cuda_version: rocm4.1
name: binary_linux_wheel_py3.8_rocm4.1
python_version: '3.8'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_linux_wheel:
cuda_version: cpu
name: binary_linux_wheel_py3.9_cpu
Expand All @@ -710,6 +731,13 @@ workflows:
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-cuda111
- binary_linux_wheel:
cuda_version: rocm4.1
name: binary_linux_wheel_py3.9_rocm4.1
python_version: '3.9'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_macos_wheel:
cuda_version: cpu
name: binary_macos_wheel_py3.6_cpu
Expand Down Expand Up @@ -1217,6 +1245,42 @@ workflows:
python_version: '3.6'
requires:
- nightly_binary_linux_wheel_py3.6_cu111_upload
- binary_linux_wheel:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.6_rocm4.1
python_version: '3.6'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_wheel_upload:
context: org-member
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.6_rocm4.1_upload
requires:
- nightly_binary_linux_wheel_py3.6_rocm4.1
- smoke_test_linux_pip:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.6_rocm4.1_smoke_test_pip
python_version: '3.6'
requires:
- nightly_binary_linux_wheel_py3.6_rocm4.1_upload
- binary_linux_wheel:
cuda_version: cpu
filters:
Expand Down Expand Up @@ -1324,6 +1388,42 @@ workflows:
python_version: '3.7'
requires:
- nightly_binary_linux_wheel_py3.7_cu111_upload
- binary_linux_wheel:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.7_rocm4.1
python_version: '3.7'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_wheel_upload:
context: org-member
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.7_rocm4.1_upload
requires:
- nightly_binary_linux_wheel_py3.7_rocm4.1
- smoke_test_linux_pip:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.7_rocm4.1_smoke_test_pip
python_version: '3.7'
requires:
- nightly_binary_linux_wheel_py3.7_rocm4.1_upload
- binary_linux_wheel:
cuda_version: cpu
filters:
Expand Down Expand Up @@ -1431,6 +1531,42 @@ workflows:
python_version: '3.8'
requires:
- nightly_binary_linux_wheel_py3.8_cu111_upload
- binary_linux_wheel:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.8_rocm4.1
python_version: '3.8'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_wheel_upload:
context: org-member
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.8_rocm4.1_upload
requires:
- nightly_binary_linux_wheel_py3.8_rocm4.1
- smoke_test_linux_pip:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.8_rocm4.1_smoke_test_pip
python_version: '3.8'
requires:
- nightly_binary_linux_wheel_py3.8_rocm4.1_upload
- binary_linux_wheel:
cuda_version: cpu
filters:
Expand Down Expand Up @@ -1538,6 +1674,42 @@ workflows:
python_version: '3.9'
requires:
- nightly_binary_linux_wheel_py3.9_cu111_upload
- binary_linux_wheel:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.9_rocm4.1
python_version: '3.9'
requires:
- download_third_parties_nix
wheel_docker_image: pytorch/manylinux-rocm:4.1
- binary_wheel_upload:
context: org-member
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.9_rocm4.1_upload
requires:
- nightly_binary_linux_wheel_py3.9_rocm4.1
- smoke_test_linux_pip:
cuda_version: rocm4.1
filters:
branches:
only:
- nightly
tags:
only: /v[0-9]+(\.[0-9]+)*-rc[0-9]+/
name: nightly_binary_linux_wheel_py3.9_rocm4.1_smoke_test_pip
python_version: '3.9'
requires:
- nightly_binary_linux_wheel_py3.9_rocm4.1_upload
- binary_macos_wheel:
cuda_version: cpu
filters:
Expand Down
11 changes: 8 additions & 3 deletions .circleci/regenerate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@


PYTHON_VERSIONS = ["3.6", "3.7", "3.8", "3.9"]
CU_VERSIONS_DICT = {"linux": ["cpu", "cu102", "cu111"],
CU_VERSIONS_DICT = {"linux": ["cpu", "cu102", "cu111","rocm4.1"],
"windows": ["cpu", "cu102", "cu111"],
"macos": ["cpu"]}

Expand All @@ -36,7 +36,10 @@ def build_workflows(prefix='', upload=False, filter_branch=None, indentation=6):
for os_type in ["linux", "macos", "windows"]:
for python_version in PYTHON_VERSIONS:
for cu_version in CU_VERSIONS_DICT[os_type]:
w += build_workflow_pair(btype, os_type, python_version, cu_version, filter_branch, prefix, upload)
if cu_version.startswith("rocm") and btype=="conda":
pass
else:
w += build_workflow_pair(btype, os_type, python_version, cu_version, filter_branch, prefix, upload)

if not filter_branch:
# Build on every pull request, but upload only on nightly and tags
Expand Down Expand Up @@ -124,8 +127,10 @@ def generate_base_workflow(base_workflow_name, python_version, cu_version, filte
d['requires'] = ['download_third_parties_nix']
if btype == 'conda':
d['conda_docker_image'] = f'pytorch/conda-builder:{cu_version.replace("cu1","cuda1")}'
elif cu_version != 'cpu':
elif cu_version.startswith('cu'):
d['wheel_docker_image'] = f'pytorch/manylinux-{cu_version.replace("cu1","cuda1")}'
elif cu_version.startswith('rocm'):
d["wheel_docker_image"] = f"pytorch/manylinux-rocm:{cu_version[len('rocm'):]}"

if filter_branch:
d["filters"] = gen_filter_branch_tree(filter_branch)
Expand Down
4 changes: 4 additions & 0 deletions packaging/pkg_helpers.bash
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,10 @@ setup_cuda() {
export FORCE_CUDA=1
export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0"
;;
rocm*)
export FORCE_CUDA=1
export USE_ROCM=1
;;
cpu)
;;
*)
Expand Down