Skip to content

Commit c3c8d5d

Browse files
jeffratjruwasejithunnair-amdrraminenjeffdaily
authored
AMD support (deepspeedai#1430)
Co-authored-by: Olatunji Ruwase <[email protected]> Co-authored-by: Jithun Nair <[email protected]> Co-authored-by: rraminen <[email protected]> Co-authored-by: Jeff Daily <[email protected]> Co-authored-by: okakarpa <[email protected]> Co-authored-by: rraminen <[email protected]> Co-authored-by: Jithun Nair <[email protected]> Co-authored-by: Jeff Daily <[email protected]> Co-authored-by: okakarpa <[email protected]> Co-authored-by: Ramya Ramineni <[email protected]>
1 parent f0304bd commit c3c8d5d

44 files changed

Lines changed: 1471 additions & 134 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/main.yml‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
name: Build
1+
name: unit-tests
22

33
on:
44
push:
@@ -14,7 +14,7 @@ on:
1414
jobs:
1515
# unit tests running on nvidia gpus
1616
nv-torch12-p40:
17-
runs-on: [self-hosted, nvidia, torch12]
17+
runs-on: [self-hosted, nvidia, torch12, p40]
1818

1919
steps:
2020
- uses: actions/checkout@v2
@@ -102,6 +102,43 @@ jobs:
102102
find examples/pytorch -regextype posix-egrep -regex '.*(language-modeling|question-answering|summarization|image-classification|text-classification|translation).*/requirements.txt' -exec pip install -r {} \;
103103
TORCH_EXTENSIONS_DIR=./torch-extensions RUN_SLOW=1 pytest --color=yes --durations=0 --verbose tests/deepspeed
104104
105+
# unit tests running on amd gpus
106+
amd:
107+
# The type of runner that the job will run on
108+
runs-on: [self-hosted, amd]
109+
110+
# Steps represent a sequence of tasks that will be executed as part of the job
111+
steps:
112+
# Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
113+
- uses: actions/checkout@v2
114+
115+
# Runs a single command using the runners shell
116+
- name: environment
117+
run: |
118+
rocm-smi --showhw
119+
which python
120+
python --version
121+
which hipcc
122+
hipcc --version
123+
python -c "import torch; print('torch:', torch.__version__, torch)"
124+
python -c "import torch; print('CUDA available:', torch.cuda.is_available())"
125+
sudo apt-get update
126+
sudo apt-get install -y libaio-dev
127+
# Runs a set of commands using the runners shell
128+
- name: Install deepspeed
129+
run: |
130+
pip install .[dev,1bit,autotuning]
131+
python -c "from deepspeed.env_report import cli_main; cli_main()"
132+
#ds_report
133+
# Runs a set of commands using the runners shell
134+
- name: Unit tests
135+
run: |
136+
if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
137+
cd tests
138+
#TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose unit/
139+
TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 -m 'not sequential' unit/
140+
TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/
141+
105142
nv-lightning-v100:
106143
runs-on: [self-hosted, nvidia, torch18, v100]
107144

‎csrc/includes/cublas_wrappers.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@
55
#include <cuda.h>
66
#include <cuda_fp16.h>
77
#include <cuda_runtime.h>
8+
#ifndef __HIP_PLATFORM_HCC__
89
#include <mma.h>
10+
#endif
911
#include <stdio.h>
1012

1113
int cublas_gemm_ex(cublasHandle_t handle,
@@ -19,7 +21,11 @@ int cublas_gemm_ex(cublasHandle_t handle,
1921
const float* A,
2022
const float* B,
2123
float* C,
24+
#ifdef __HIP_PLATFORM_HCC__
25+
rocblas_gemm_algo algo = rocblas_gemm_algo_standard);
26+
#else
2227
cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT);
28+
#endif
2329

2430
int cublas_gemm_ex(cublasHandle_t handle,
2531
cublasOperation_t transa,
@@ -32,7 +38,11 @@ int cublas_gemm_ex(cublasHandle_t handle,
3238
const __half* A,
3339
const __half* B,
3440
__half* C,
41+
#ifdef __HIP_PLATFORM_HCC__
42+
rocblas_gemm_algo algo = rocblas_gemm_algo_standard);
43+
#else
3544
cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP);
45+
#endif
3646

3747
int cublas_strided_batched_gemm(cublasHandle_t handle,
3848
int m,
@@ -49,7 +59,11 @@ int cublas_strided_batched_gemm(cublasHandle_t handle,
4959
int stride_B,
5060
int stride_C,
5161
int batch,
62+
#ifdef __HIP_PLATFORM_HCC__
63+
rocblas_gemm_algo algo = rocblas_gemm_algo_standard);
64+
#else
5265
cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT);
66+
#endif
5367

5468
int cublas_strided_batched_gemm(cublasHandle_t handle,
5569
int m,
@@ -66,4 +80,8 @@ int cublas_strided_batched_gemm(cublasHandle_t handle,
6680
int stride_B,
6781
int stride_C,
6882
int batch,
83+
#ifdef __HIP_PLATFORM_HCC__
84+
rocblas_gemm_algo algo = rocblas_gemm_algo_standard);
85+
#else
6986
cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP);
87+
#endif

‎csrc/includes/custom_cuda_layers.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,15 @@
55
#include <stdio.h>
66
#include <stdlib.h>
77

8+
#ifdef __HIP_PLATFORM_HCC__
9+
#define HALF_PRECISION_AVAILABLE = 1
10+
#include <hip/hip_cooperative_groups.h>
11+
#else
12+
#if __CUDA_ARCH__ >= 700
13+
#define HALF_PRECISION_AVAILABLE = 1
14+
#endif
815
#include <cooperative_groups.h>
16+
#endif
917
#include <curand_kernel.h>
1018

1119
#include "context.h"

‎csrc/includes/feed_forward.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,11 @@ class FeedForward {
4343
weights,
4444
input_ptr,
4545
out,
46+
#ifdef __HIP_PLATFORM_HCC__
47+
rocblas_gemm_algo(config_.gemm_algos[0]));
48+
#else
4649
cublasGemmAlgo_t(config_.gemm_algos[0]));
50+
#endif
4751
}
4852
void Backward(int bsz,
4953
const T* out_grad,
@@ -68,7 +72,11 @@ class FeedForward {
6872
input_ptr,
6973
out_grad,
7074
weights_grad,
75+
#ifdef __HIP_PLATFORM_HCC__
76+
rocblas_gemm_algo(config_.gemm_algos[1]));
77+
#else
7178
cublasGemmAlgo_t(config_.gemm_algos[1]));
79+
#endif
7280

7381
cublas_gemm_ex(_cublasHandle,
7482
CUBLAS_OP_N,
@@ -81,7 +89,11 @@ class FeedForward {
8189
weights,
8290
out_grad,
8391
inp_grad_out,
92+
#ifdef __HIP_PLATFORM_HCC__
93+
rocblas_gemm_algo(config_.gemm_algos[2]));
94+
#else
8495
cublasGemmAlgo_t(config_.gemm_algos[2]));
96+
#endif
8597

8698
launch_fuse_transpose_bias_kernel<T>(out_grad, bias_grad, bsz, config_.outputSize, stream);
8799
}

‎csrc/includes/gemm_test.h‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22
#pragma once
33

44
#include <cuda_fp16.h>
5+
#ifndef __HIP_PLATFORM_HCC__
56
#include <cuda_profiler_api.h>
7+
#endif
68
#include <array>
79
#include <cstdio>
810
#include <cstdlib>
@@ -58,7 +60,11 @@ class GemmTest {
5860
B,
5961
A,
6062
C,
63+
#ifdef __HIP_PLATFORM_HCC__
64+
static_cast<rocblas_gemm_algo>(algo));
65+
#else
6166
static_cast<cublasGemmAlgo_t>(algo));
67+
#endif
6268
});
6369

6470
int algo_bw1 = Run(loops, [=](int algo) {
@@ -73,7 +79,11 @@ class GemmTest {
7379
A,
7480
C,
7581
B,
82+
#ifdef __HIP_PLATFORM_HCC__
83+
static_cast<rocblas_gemm_algo>(algo));
84+
#else
7685
static_cast<cublasGemmAlgo_t>(algo));
86+
#endif
7787
});
7888

7989
int algo_bw2 = Run(loops, [=](int algo) {
@@ -88,7 +98,11 @@ class GemmTest {
8898
B,
8999
C,
90100
A,
101+
#ifdef __HIP_PLATFORM_HCC__
102+
static_cast<rocblas_gemm_algo>(algo));
103+
#else
91104
static_cast<cublasGemmAlgo_t>(algo));
105+
#endif
92106
});
93107

94108
return std::array<int, 3>({algo_fw, algo_bw1, algo_bw2});
@@ -100,8 +114,12 @@ class GemmTest {
100114
float fast_latency = (std::numeric_limits<float>::max)();
101115
int fast_algo = 0;
102116

117+
#ifdef __HIP_PLATFORM_HCC__
118+
for (int algo = (int)rocblas_gemm_algo_standard; algo <= (int)rocblas_gemm_algo_standard;
119+
#else
103120
for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP;
104121
algo <= (int)CUBLAS_GEMM_ALGO15_TENSOR_OP;
122+
#endif
105123
algo++) {
106124
int warm_up = 5;
107125
for (int i = 0; i < warm_up; ++i) f(algo);
@@ -186,7 +204,11 @@ class StridedGemmTest {
186204
stride_b,
187205
stride_c,
188206
bsz,
207+
#ifdef __HIP_PLATFORM_HCC__
208+
static_cast<rocblas_gemm_algo>(algo));
209+
#else
189210
static_cast<cublasGemmAlgo_t>(algo));
211+
#endif
190212
});
191213

192214
int algo_bw1 = Run(loops, [=](int algo) {
@@ -216,7 +238,11 @@ class StridedGemmTest {
216238
stride_b,
217239
stride_c,
218240
bsz,
241+
#ifdef __HIP_PLATFORM_HCC__
242+
static_cast<rocblas_gemm_algo>(algo));
243+
#else
219244
static_cast<cublasGemmAlgo_t>(algo));
245+
#endif
220246
});
221247

222248
int algo_bw2 = Run(loops, [=](int algo) {
@@ -243,7 +269,11 @@ class StridedGemmTest {
243269
stride_b,
244270
stride_c,
245271
bsz,
272+
#ifdef __HIP_PLATFORM_HCC__
273+
static_cast<rocblas_gemm_algo>(algo));
274+
#else
246275
static_cast<cublasGemmAlgo_t>(algo));
276+
#endif
247277
});
248278

249279
return std::array<int, 3>({algo_fw, algo_bw1, algo_bw2});
@@ -255,8 +285,12 @@ class StridedGemmTest {
255285
float fast_latency = (std::numeric_limits<float>::max)();
256286
int fast_algo = 0;
257287

288+
#ifdef __HIP_PLATFORM_HCC__
289+
for (int algo = (int)rocblas_gemm_algo_standard; algo <= (int)rocblas_gemm_algo_standard;
290+
#else
258291
for (int algo = (int)CUBLAS_GEMM_DEFAULT_TENSOR_OP;
259292
algo <= (int)CUBLAS_GEMM_ALGO15_TENSOR_OP;
293+
#endif
260294
algo++) {
261295
int warm_up = 5;
262296
for (int i = 0; i < warm_up; ++i) f(algo);

‎csrc/includes/general_kernels.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,11 @@
33
#include <stdio.h>
44
#include <stdlib.h>
55

6+
#ifdef __HIP_PLATFORM_HCC__
7+
#include <hip/hip_cooperative_groups.h>
8+
#else
69
#include <cooperative_groups.h>
10+
#endif
711
#include <curand_kernel.h>
812

913
#include "context.h"

0 commit comments

Comments
 (0)