saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Enable TF32 for training on A100 (#4914)

Browse source 
* enable TF32 for training on A100

it can be disabled by env: NVIDIA_TF32_OVERRIDE = 0
This commit is contained in:
Weixing Zhang 2020-09-02 19:21:54 -07:00 committed by GitHub
parent a9db287bd7
commit 3268717615
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 46 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
onnxruntime/core/providers/cuda/cuda_common.h
View file

@ -226,15 +226,21 @@ inline bool CalculateFdmStrides(gsl::span<fast_divmod> p, const std::vector<int6
class CublasMathModeSetter {
public:
CublasMathModeSetter(const cudaDeviceProp& prop,cublasHandle_t handle, cublasMath_t mode) : prop_(prop), handle_(handle) {
if (prop_.major >= 7) {
cublasGetMathMode(handle, &mode_);
cublasGetMathMode(handle, &mode_);
#if defined(CUDA_VERSION) && CUDA_VERSION < 11000
if (prop.major >= 7 && mode == CUBLAS_TENSOR_OP_MATH) {
cublasSetMathMode(handle, mode);
}
}
~CublasMathModeSetter() {
if (prop_.major >= 7) {
cublasSetMathMode(handle_, mode_);
#endif
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
if (prop.major >= 8 && mode == CUBLAS_TF32_TENSOR_OP_MATH) {
cublasSetMathMode(handle, mode);
}
#endif
}
~CublasMathModeSetter() {
cublasSetMathMode(handle_, mode_);
}
private:

37
onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h
View file

@ -27,7 +27,17 @@ inline cublasStatus_t cublasGemmHelper(cublasHandle_t handle,
const float* B, int ldb,
const float* beta,
float* C, int ldc,
const cudaDeviceProp& /*prop*/) {
const cudaDeviceProp& prop) {
#ifdef ENABLE_TRAINING
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TF32_TENSOR_OP_MATH);
#else
ORT_UNUSED_PARAMETER(prop);
#endif
#else
ORT_UNUSED_PARAMETER(prop);
#endif
return cublasSgemm(handle,
transa,
transb,
@ -111,7 +121,16 @@ inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle,
const float* beta,
float* Carray[], int ldc,
int batch_count,
const cudaDeviceProp& /*prop*/) {
const cudaDeviceProp& prop) {
#ifdef ENABLE_TRAINING
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TF32_TENSOR_OP_MATH);
#else
ORT_UNUSED_PARAMETER(prop);
#endif
#else
ORT_UNUSED_PARAMETER(prop);
#endif
return cublasSgemmBatched(handle,
transa,
transb,
@ -157,6 +176,7 @@ inline cublasStatus_t cublasGemmBatchedHelper(cublasHandle_t handle,
int batch_count,
const cudaDeviceProp& prop) {
onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TENSOR_OP_MATH);
#ifdef ENABLE_TRAINING
float h_a = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(alpha));
float h_b = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(beta));
@ -203,7 +223,17 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle,
float* C, int ldc,
long long int strideC,
int batch_count,
const cudaDeviceProp& /*prop*/) {
const cudaDeviceProp& prop) {
#ifdef ENABLE_TRAINING
#if defined(CUDA_VERSION) && CUDA_VERSION >= 11000
onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TF32_TENSOR_OP_MATH);
#else
ORT_UNUSED_PARAMETER(prop);
#endif
#else
ORT_UNUSED_PARAMETER(prop);
#endif
return cublasSgemmStridedBatched(handle,
transa,
transb,
@ -257,6 +287,7 @@ inline cublasStatus_t cublasGemmStridedBatchedHelper(cublasHandle_t handle,
int batch_count,
const cudaDeviceProp& prop) {
onnxruntime::cuda::CublasMathModeSetter math_mode_setter(prop, handle, CUBLAS_TENSOR_OP_MATH);
#ifdef ENABLE_TRAINING
float h_a = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(alpha));
float h_b = onnxruntime::math::halfToFloat(*reinterpret_cast<const uint16_t*>(beta));