Skip to content

GitLab

Commit 6f7a8b39 authored by lcskrishna's avatar lcskrishna
Browse files

Merge remote-tracking branch 'rocm_upstream/master' into ifu_07272020

parents 459de22d 9c80f6d3
Hide whitespace changes
Inline Side-by-side
Showing with 575 additions and 178 deletions
csrc/type_shim.h
View file @ 6f7a8b39
......@@ -105,6 +105,66 @@
AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
}
// TODO: We might have come up with an optimal set of dispatch macros by
// changing the signature to have an integer suffix of number of types
// to dispatch for as defined in upstream (e.g AT_DISPATCH_FLOATING_TYPES_AND2)
// Refactor once all the extension ops are enabled.
#define DISPATCH_FLOAT_AND_HALF_AND_BFLOAT16(TYPE, LEVEL, NAME, ...) \
switch(TYPE) \
{ \
case at::ScalarType::Float: \
{ \
using scalar_t_##LEVEL = float; \
__VA_ARGS__; \
break; \
} \
case at::ScalarType::Half: \
{ \
using scalar_t_##LEVEL = at::Half; \
__VA_ARGS__; \
break; \
} \
case at::ScalarType::BFloat16: \
{ \
using scalar_t_##LEVEL = at::BFloat16; \
__VA_ARGS__; \
break; \
} \
default: \
AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
}
#define DISPATCH_DOUBLE_FLOAT_AND_HALF_AND_BFLOAT16(TYPE, LEVEL, NAME, ...) \
switch(TYPE) \
{ \
case at::ScalarType::Double: \
{ \
using scalar_t_##LEVEL = double; \
__VA_ARGS__; \
break; \
} \
case at::ScalarType::Float: \
{ \
using scalar_t_##LEVEL = float; \
__VA_ARGS__; \
break; \
} \
case at::ScalarType::Half: \
{ \
using scalar_t_##LEVEL = at::Half; \
__VA_ARGS__; \
break; \
} \
case at::ScalarType::BFloat16: \
{ \
using scalar_t_##LEVEL = at::BFloat16; \
__VA_ARGS__; \
break; \
} \
default: \
AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
}
template<typename T>
__device__ __forceinline__ T reduce_block_into_lanes
......@@ -141,8 +201,13 @@ __device__ __forceinline__ T reduce_block_into_lanes
// __SYNCWARP();
#pragma unroll
for(int i = 16; i >= lanes; i >>= 1)
for(int i = 16; i >= lanes; i >>= 1) {
#ifdef __HIP_PLATFORM_HCC__
final = final + __shfl_down(0xffffffff, final, i);
#else
final = final + __shfl_down_sync(0xffffffff, final, i);
#endif
}
}
if(share_result)
......@@ -191,8 +256,13 @@ __device__ __forceinline__ T reduce_block_into_lanes_max_op
// __SYNCWARP();
#pragma unroll
for(int i = 16; i >= lanes; i >>= 1)
for(int i = 16; i >= lanes; i >>= 1) {
#ifdef __HIP_PLATFORM_HCC__
final = fmaxf(fabsf(final), fabsf(__shfl_down(0xffffffff, final, i)));
#else
final = fmaxf(fabsf(final), fabsf(__shfl_down_sync(0xffffffff, final, i)));
#endif
}
}
if(share_result)
......
csrc/welford.cu
View file @ 6f7a8b39
......@@ -11,6 +11,11 @@
#include "type_shim.h"
#include "compat.h"
#if defined __HIP_PLATFORM_HCC__
#define SHFL_DOWN __shfl_down
#else
#define SHFL_DOWN __shfl_down_sync
#endif
__device__ __forceinline__ int lastpow2(int n)
{
......@@ -47,7 +52,7 @@ __device__ __forceinline__ T warp_reduce_sum(T val)
{
#pragma unroll
for(int i = WARP_SIZE/2; i > 0; i >>= 1)
val = val + __shfl_down_sync(0xffffffff, val, i);
val = val + SHFL_DOWN(0xffffffff, val, i);
return val;
}
......@@ -129,10 +134,14 @@ __device__ __forceinline__ void warp_reduce_mean_m2n(T &mean, T &m2n, int &num)
{
#pragma unroll
for(int i = WARP_SIZE/2; i > 0; i >>= 1) {
auto num_new = __shfl_down_sync(0xffffffff, num, i);
auto mean_new = __shfl_down_sync(0xffffffff, mean, i);
auto m2n_new = __shfl_down_sync(0xffffffff, m2n, i);
auto num_new = SHFL_DOWN(0xffffffff, num, i);
auto mean_new = SHFL_DOWN(0xffffffff, mean, i);
auto m2n_new = SHFL_DOWN(0xffffffff, m2n, i);
#if defined __HIP_PLATFORM_HCC__
welford_merge_element<T, int>(num, mean, m2n, num_new, mean_new, m2n_new);
#else
welford_merge_element(num, mean, m2n, num_new, mean_new, m2n_new);
#endif
}
}
......
setup.py
View file @ 6f7a8b39
......@@ -6,6 +6,8 @@ import sys
import warnings
import os
from torch.utils.hipify import hipify_python
# ninja build does not work unless include_dirs are abs path
this_dir = os.path.dirname(os.path.abspath(__file__))
......@@ -124,51 +126,104 @@ if "--cuda_ext" in sys.argv:
from torch.utils.cpp_extension import CUDAExtension
sys.argv.remove("--cuda_ext")
if torch.utils.cpp_extension.CUDA_HOME is None:
is_rocm_pytorch = False
if torch.__version__ >= '1.5':
from torch.utils.cpp_extension import ROCM_HOME
is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False
if torch.utils.cpp_extension.CUDA_HOME is None and (not is_rocm_pytorch):
raise RuntimeError("--cuda_ext was requested, but nvcc was not found. Are you sure your environment has nvcc available? If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
else:
check_cuda_torch_binary_vs_bare_metal(torch.utils.cpp_extension.CUDA_HOME)
ext_modules.append(
CUDAExtension(name='amp_C',
sources=['csrc/amp_C_frontend.cpp',
'csrc/multi_tensor_sgd_kernel.cu',
'csrc/multi_tensor_scale_kernel.cu',
'csrc/multi_tensor_axpby_kernel.cu',
'csrc/multi_tensor_l2norm_kernel.cu',
'csrc/multi_tensor_lamb_stage_1.cu',
'csrc/multi_tensor_lamb_stage_2.cu',
'csrc/multi_tensor_adam.cu',
'csrc/multi_tensor_adagrad.cu',
'csrc/multi_tensor_novograd.cu',
'csrc/multi_tensor_lamb.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-lineinfo',
'-O3',
# '--resource-usage',
'--use_fast_math'] + version_dependent_macros}))
ext_modules.append(
CUDAExtension(name='syncbn',
sources=['csrc/syncbn.cpp',
'csrc/welford.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-O3'] + version_dependent_macros}))
ext_modules.append(
CUDAExtension(name='fused_layer_norm_cuda',
sources=['csrc/layer_norm_cuda.cpp',
'csrc/layer_norm_cuda_kernel.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-maxrregcount=50',
'-O3',
'--use_fast_math'] + version_dependent_macros}))
ext_modules.append(
CUDAExtension(name='mlp_cuda',
sources=['csrc/mlp.cpp',
'csrc/mlp_cuda.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-O3'] + version_dependent_macros}))
if not is_rocm_pytorch:
check_cuda_torch_binary_vs_bare_metal(torch.utils.cpp_extension.CUDA_HOME)
if is_rocm_pytorch:
import shutil
with hipify_python.GeneratedFileCleaner(keep_intermediates=True) as clean_ctx:
hipify_python.hipify(project_directory=this_dir, output_directory=this_dir, includes="csrc/*",
show_detailed=True, is_pytorch_extension=True, clean_ctx=clean_ctx)
shutil.copy("csrc/compat.h", "csrc/hip/compat.h")
shutil.copy("csrc/type_shim.h", "csrc/hip/type_shim.h")
if not is_rocm_pytorch:
ext_modules.append(
CUDAExtension(name='amp_C',
sources=['csrc/amp_C_frontend.cpp',
'csrc/multi_tensor_sgd_kernel.cu',
'csrc/multi_tensor_scale_kernel.cu',
'csrc/multi_tensor_axpby_kernel.cu',
'csrc/multi_tensor_l2norm_kernel.cu',
'csrc/multi_tensor_lamb_stage_1.cu',
'csrc/multi_tensor_lamb_stage_2.cu',
'csrc/multi_tensor_adam.cu',
'csrc/multi_tensor_adagrad.cu',
'csrc/multi_tensor_novograd.cu',
'csrc/multi_tensor_lamb.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-lineinfo',
'-O3',
# '--resource-usage',
'--use_fast_math'] + version_dependent_macros}))
else:
print ("INFO: Building Multitensor apply extension")
ext_modules.append(
CUDAExtension(name='amp_C',
sources=['csrc/amp_C_frontend.cpp',
'csrc/hip/multi_tensor_sgd_kernel.hip',
'csrc/hip/multi_tensor_scale_kernel.hip',
'csrc/hip/multi_tensor_axpby_kernel.hip',
'csrc/hip/multi_tensor_l2norm_kernel.hip',
'csrc/hip/multi_tensor_lamb_stage_1.hip',
'csrc/hip/multi_tensor_lamb_stage_2.hip',
'csrc/hip/multi_tensor_adam.hip',
'csrc/hip/multi_tensor_adagrad.hip',
'csrc/hip/multi_tensor_novograd.hip',
'csrc/hip/multi_tensor_lamb.hip'],
extra_compile_args=['-O3'] + version_dependent_macros))
if not is_rocm_pytorch:
ext_modules.append(
CUDAExtension(name='syncbn',
sources=['csrc/syncbn.cpp',
'csrc/welford.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-O3'] + version_dependent_macros}))
else:
print ("INFO: Building syncbn extension.")
ext_modules.append(
CUDAExtension(name='syncbn',
sources=['csrc/syncbn.cpp',
'csrc/hip/welford.hip'],
extra_compile_args=['-O3'] + version_dependent_macros))
if not is_rocm_pytorch:
ext_modules.append(
CUDAExtension(name='fused_layer_norm_cuda',
sources=['csrc/layer_norm_cuda.cpp',
'csrc/layer_norm_cuda_kernel.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-maxrregcount=50',
'-O3',
'--use_fast_math'] + version_dependent_macros}))
else:
print ("INFO: Building FusedLayerNorm extension.")
ext_modules.append(
CUDAExtension(name='fused_layer_norm_cuda',
sources=['csrc/layer_norm_cuda.cpp',
'csrc/hip/layer_norm_hip_kernel.hip'],
extra_compile_args={'cxx' : ['-O3'] + version_dependent_macros,
'nvcc' : []}))
if not is_rocm_pytorch:
ext_modules.append(
CUDAExtension(name='mlp_cuda',
sources=['csrc/mlp.cpp',
'csrc/mlp_cuda.cu'],
extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
'nvcc':['-O3'] + version_dependent_macros}))
else:
print ("INFO: Skipping MLP extension")
if "--bnp" in sys.argv:
from torch.utils.cpp_extension import CUDAExtension
......@@ -250,7 +305,7 @@ if "--deprecated_fused_lamb" in sys.argv:
'nvcc':['-O3',
'--use_fast_math'] + version_dependent_macros}))
# Check, if ATen/CUDAGenerator.h is found, otherwise use the new ATen/CUDAGeneratorImpl.h, due to breaking change in https://github.com/pytorch/pytorch/pull/36026
# Check, if ATen/CUDAGenerator.h is found, otherwise use the new ATen/CUDAGeneratorImpl.h, due to breaking change in https://github.com/pytorch/pytorch/pull/36026
generator_flag = []
torch_dir = torch.__path__[0]
if os.path.exists(os.path.join(torch_dir, 'include', 'ATen', 'CUDAGenerator.h')):
......
tests/L0/run_amp/test_add_param_group.py
View file @ 6f7a8b39
......@@ -14,11 +14,11 @@ from utils import common_init, HALF, FLOAT,\
ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
class MyModel(torch.nn.Module):
def __init__(self, unique):
def __init__(self, unique, dtype=torch.float16):
super(MyModel, self).__init__()
self.weight0 = Parameter(unique +
torch.arange(2, device='cuda', dtype=torch.float32))
self.weight1 = Parameter(1. + unique + torch.arange(2, device='cuda', dtype=torch.float16))
self.weight1 = Parameter(1. + unique + torch.arange(2, device='cuda', dtype=dtype))
@staticmethod
def ops(input, weight0, weight1):
......@@ -51,11 +51,15 @@ class TestAddParamGroup(unittest.TestCase):
optimizer.zero_grad()
def test_add_param_group(self):
for opt_level in ("O0", "O1", "O2", "O3"):
for opt_level in ("O0", "O1", "O2", "O3", "O4", "O5"):
for zero_before_add in (True, False):
for try_accumulation in (True, False):
model0 = MyModel(1)
model1 = MyModel(2)
if opt_level in {"O4", "O5"}:
model0 = MyModel(1, torch.bfloat16)
model1 = MyModel(2, torch.bfloat16)
else:
model0 = MyModel(1)
model1 = MyModel(2)
optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
momentum=0.125)
......@@ -89,8 +93,12 @@ class TestAddParamGroup(unittest.TestCase):
[param.data.clone() for param in model1.parameters()]
for how_to_zero in "none", "model", "optimizer":
model0 = MyModel(1)
model1 = MyModel(2)
if opt_level in {"O4", "O5"}:
model0 = MyModel(1, torch.bfloat16)
model1 = MyModel(2, torch.bfloat16)
else:
model0 = MyModel(1)
model1 = MyModel(2)
optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
momentum=0.125)
......@@ -139,6 +147,9 @@ class TestAddParamGroup(unittest.TestCase):
[param.data.clone() for param in model1.parameters()]
for reference, final in zip(reference_params, final_params):
# TODO: remove the conversion once allclose supports bfloat16 type.
if final.dtype == torch.bfloat16:
final = final.float()
self.assertTrue(torch.allclose(reference.to(final.dtype), final),
"opt_level = {}, how_to_zero = {}, zero_before_add = {}".format(
opt_level, how_to_zero, zero_before_add))
......
tests/L0/run_amp/test_basic_casts.py
View file @ 6f7a8b39
......@@ -9,7 +9,9 @@ from torch import nn
import torch.nn.functional as F
from utils import common_init, HALF, FLOAT,\
ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
ALWAYS_HALF, ALWAYS_BFLOAT16, ALWAYS_FLOAT, MATCH_INPUT
from apex.testing.common_utils import skipIfRocm
def run_layer_test(test_case, fns, expected, input_shape, test_backward=True):
for fn, typ in it.product(fns, expected.keys()):
......@@ -20,124 +22,237 @@ def run_layer_test(test_case, fns, expected, input_shape, test_backward=True):
y.float().sum().backward()
test_case.assertEqual(x.grad.type(), MATCH_INPUT[typ])
class TestBasicCasts(unittest.TestCase):
def setUp(self):
self.handle = amp.init(enabled=True)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def test_linear_is_half(self):
class _TestBasicCasts(unittest.TestCase):
def _test_linear(self, expected):
m = nn.Linear(self.h, self.h)
f = ft.partial(F.linear, weight=m.weight, bias=m.bias)
run_layer_test(self, [m, f], ALWAYS_HALF, (self.b, self.h))
run_layer_test(self, [m, f], expected, (self.b, self.h))
def test_conv2d_is_half(self):
def _test_conv2d(self, expected):
m = nn.Conv2d(self.c, self.c, self.k)
f = ft.partial(F.conv2d, weight=m.weight, bias=m.bias)
run_layer_test(self, [m, f], ALWAYS_HALF, (self.b, self.c, self.h, self.h))
run_layer_test(self, [m, f], expected, (self.b, self.c, self.h, self.h))
def test_softmax_is_float(self):
def _test_softmax(self, expected):
m = nn.Softmax(dim=1)
f = ft.partial(F.softmax, dim=1)
run_layer_test(self, [m, f], ALWAYS_FLOAT, (self.b, self.h))
run_layer_test(self, [m, f], expected, (self.b, self.h))
def test_group_norm_is_float(self):
def _test_group_norm(self, expected):
m = nn.GroupNorm(num_groups=4, num_channels=self.c)
run_layer_test(self, [m], ALWAYS_FLOAT, (self.b, self.c, self.h, self.h))
run_layer_test(self, [m], expected, (self.b, self.c, self.h, self.h))
def test_mse_loss_is_float(self):
def _test_mse_loss(self, expected):
shape = (self.b, self.h)
target = torch.randn(shape)
mod = nn.MSELoss()
m = lambda x: mod(x, target)
f = ft.partial(F.mse_loss, target=target)
run_layer_test(self, [m], ALWAYS_FLOAT, shape)
run_layer_test(self, [m], expected, shape)
def test_relu_is_match(self):
run_layer_test(self, [nn.ReLU(), F.relu], MATCH_INPUT, (self.b, self.h))
def _test_relu(self, expected):
run_layer_test(self, [nn.ReLU(), F.relu], expected, (self.b, self.h))
def test_batch_norm_is_match(self):
def _test_batch_norm(self, expected):
m = nn.BatchNorm2d(num_features=self.c)
f = ft.partial(F.batch_norm, running_mean=m.running_mean, running_var=m.running_var,
weight=m.weight, bias=m.bias, training=True)
run_layer_test(self, [m], MATCH_INPUT, (self.b, self.c, self.h, self.h))
run_layer_test(self, [m], expected, (self.b, self.c, self.h, self.h))
# Test forward-only for BN inference
m.eval()
f = ft.partial(F.batch_norm, running_mean=m.running_mean, running_var=m.running_var,
weight=m.weight, bias=m.bias, training=False)
run_layer_test(self, [m, f], MATCH_INPUT, (self.b, self.c, self.h, self.h),
run_layer_test(self, [m, f], expected, (self.b, self.c, self.h, self.h),
test_backward=False)
class TestBasicCastsHalf(_TestBasicCasts):
def setUp(self):
self.handle = amp.init(enabled=True, patch_type=torch.half)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def test_linear_is_half(self):
self._test_linear(ALWAYS_HALF)
def test_conv2d_is_half(self):
self._test_conv2d(ALWAYS_HALF)
def test_softmax_is_float(self):
self._test_softmax(ALWAYS_FLOAT)
def test_group_norm_is_float(self):
self._test_group_norm(ALWAYS_FLOAT)
def test_mse_loss_is_float(self):
self._test_mse_loss(ALWAYS_FLOAT)
def test_relu_is_match(self):
self._test_relu(MATCH_INPUT)
def test_batch_norm_is_match(self):
self._test_batch_norm(MATCH_INPUT)
class TestBasicCastsBFloat16(_TestBasicCasts):
def setUp(self):
self.handle = amp.init(enabled=True, patch_type=torch.bfloat16)
common_init(self)
def tearDown(self):
self.handle._deactivate()
@skipIfRocm
def test_linear_is_bfloat16(self):
self._test_linear(ALWAYS_BFLOAT16)
@skipIfRocm
def test_conv2d_is_bfloat16(self):
self._test_conv2d(ALWAYS_BFLOAT16)
def test_softmax_is_float(self):
self._test_softmax(ALWAYS_FLOAT)
def test_group_norm_is_float(self):
self._test_group_norm(ALWAYS_FLOAT)
def test_mse_loss_is_float(self):
self._test_mse_loss(ALWAYS_FLOAT)
def test_relu_is_match(self):
self._test_relu(MATCH_INPUT)
def test_batch_norm_is_match(self):
self._test_batch_norm(MATCH_INPUT)
class TestBannedMethods(unittest.TestCase):
def setUp(self):
self.handle = amp.init(enabled=True)
self.handle = amp.init(enabled=True, patch_type=torch.half)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def bce_common(self, assertion):
def bce_common(self, assertion, dtype=torch.half):
shape = (self.b, self.h)
target = torch.rand(shape)
mod = nn.BCELoss()
m = lambda x: mod(x, target)
f = ft.partial(F.binary_cross_entropy, target=target)
for fn in [m, f]:
x = torch.rand(shape, dtype=torch.half)
x = torch.rand(shape, dtype=dtype)
assertion(fn, x)
def test_bce_raises_by_default(self):
assertion = lambda fn, x: self.assertRaises(NotImplementedError, fn, x)
self.bce_common(assertion)
self.bce_common(assertion, dtype=torch.half)
# handle with bfloat16 as patch_type
self.handle._deactivate()
self.handle = amp.init(enabled=True, patch_type=torch.bfloat16)
self.bce_common(assertion, dtype=torch.bfloat16)
def test_bce_is_float_with_allow_banned(self):
self.handle._deactivate()
self.handle = amp.init(enabled=True, allow_banned=True)
self.handle = amp.init(enabled=True, allow_banned=True, patch_type=torch.half)
assertion = lambda fn, x: self.assertEqual(fn(x).type(), FLOAT)
self.bce_common(assertion)
self.bce_common(assertion, dtype=torch.half)
class TestTensorCasts(unittest.TestCase):
def setUp(self):
self.handle = amp.init(enabled=True)
common_init(self)
def tearDown(self):
# handle with bfloat16 as patch_type
self.handle._deactivate()
self.handle = amp.init(enabled=True, allow_banned=True, patch_type=torch.bfloat16)
self.bce_common(assertion, dtype=torch.bfloat16)
def test_matmul_method_is_half(self):
class _TestTensorCasts(unittest.TestCase):
def _test_matmul_method(self, expected):
other = torch.randn(self.h, self.h)
lhs = lambda x: x.matmul(other)
rhs = lambda x: other.matmul(x)
run_layer_test(self, [lhs, rhs], ALWAYS_HALF, (self.h, self.h))
run_layer_test(self, [lhs, rhs], expected, (self.h, self.h))
def test_matmul_op_is_half(self):
def _test_matmul_op(self, expected):
other = torch.randn(self.h, self.h)
lhs = lambda x: x @ other
rhs = lambda x: other @ x
run_layer_test(self, [lhs, rhs], ALWAYS_HALF, (self.h, self.h))
run_layer_test(self, [lhs, rhs], expected, (self.h, self.h))
def test_pow_method_is_float(self):
def _test_pow_method(self, expected):
fn = lambda x: x.pow(2.)
run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
run_layer_test(self, [fn], expected, (self.b, self.h))
def test_pow_op_is_float(self):
def _test_pow_op(self, expected):
fn = lambda x: x ** 2.
run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
run_layer_test(self, [fn], expected, (self.b, self.h))
def test_cpu_is_float(self):
def _test_cpu(self, expected):
fn = lambda x: x.cpu()
run_layer_test(self, [fn], expected, (self.b, self.h))
def _test_sum(self, expected):
fn = lambda x: x.sum()
run_layer_test(self, [fn], expected, (self.b, self.h))
# TODO: maybe more tests on disabled casting?
class TestTensorCastsHalf(_TestTensorCasts):
def setUp(self):
self.handle = amp.init(enabled=True, patch_type=torch.half)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def test_matmul_method_is_half(self):
self._test_matmul_method(ALWAYS_HALF)
def test_matmul_op_is_half(self):
self._test_matmul_op(ALWAYS_HALF)
def test_pow_method_is_float(self):
self._test_pow_method(ALWAYS_FLOAT)
def test_pow_op_is_float(self):
self._test_pow_op(ALWAYS_FLOAT)
def test_cpu_is_float(self):
always_cpu_float = {torch.float: 'torch.FloatTensor',
torch.half: 'torch.FloatTensor'}
run_layer_test(self, [fn], always_cpu_float, (self.b, self.h))
self._test_cpu(always_cpu_float)
def test_sum_is_float(self):
fn = lambda x: x.sum()
run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
self._test_sum(ALWAYS_FLOAT)
class TestTensorCastsBFloat16(_TestTensorCasts):
def setUp(self):
self.handle = amp.init(enabled=True, patch_type=torch.bfloat16)
common_init(self)
def tearDown(self):
self.handle._deactivate()
@skipIfRocm
def test_matmul_method_is_bfloat16(self):
self._test_matmul_method(ALWAYS_BFLOAT16)
@skipIfRocm
def test_matmul_op_is_bfloat16(self):
self._test_matmul_op(ALWAYS_BFLOAT16)
def test_pow_method_is_float(self):
self._test_pow_method(ALWAYS_FLOAT)
def test_pow_op_is_float(self):
self._test_pow_op(ALWAYS_FLOAT)
def test_cpu_is_float(self):
always_cpu_float = {torch.float: 'torch.FloatTensor',
torch.bfloat16: 'torch.FloatTensor'}
self._test_cpu(always_cpu_float)
def test_sum_is_float(self):
self._test_sum(ALWAYS_FLOAT)
# TODO: maybe more tests on disabled casting?
if __name__ == '__main__':
unittest.main()
tests/L0/run_amp/test_cache.py
View file @ 6f7a8b39
......@@ -67,12 +67,12 @@ class TestCache(unittest.TestCase):
def tearDown(self):
pass
def train_eval_train_test(self, module, t):
def train_eval_train_test(self, module, t, opt_level):
model = module(t).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
_amp_state.allow_incoming_model_not_fp32 = True
model, optimizer = amp.initialize(model, optimizer, opt_level="O1", verbosity=0)
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level, verbosity=0)
_amp_state.allow_incoming_model_not_fp32 = False
def training_step():
......@@ -93,6 +93,8 @@ class TestCache(unittest.TestCase):
# but I'm keeping this in case we want different tolerances for fp16 and fp32 checks.
if model.weight.grad.type() == "torch.cuda.HalfTensor":
self.assertTrue(torch.allclose(model.weight.grad.float(), reference_grad))
elif model.weight.grad.type() == "torch.cuda.BFloat16Tensor":
self.assertTrue(torch.allclose(model.weight.grad.float(), reference_grad))
elif model.weight.grad.type() == "torch.cuda.FloatTensor":
self.assertTrue(torch.allclose(model.weight.grad.float(), reference_grad))
else:
......@@ -115,22 +117,41 @@ class TestCache(unittest.TestCase):
# I could easily have these as a set of for loops in a single test,
# instead of going for granularity.
def test_whitelist_module_fp16_weight(self):
self.train_eval_train_test(WhitelistModule, torch.float16)
self.train_eval_train_test(WhitelistModule, torch.float16, "O1")
def test_whitelist_module_fp32_weight(self):
self.train_eval_train_test(WhitelistModule, torch.float32)
self.train_eval_train_test(WhitelistModule, torch.float32, "O1")
def test_blacklist_module_fp16_weight(self):
self.train_eval_train_test(BlacklistModule, torch.float16)
self.train_eval_train_test(BlacklistModule, torch.float16, "O1")
def test_blacklist_module_fp32_weight(self):
self.train_eval_train_test(BlacklistModule, torch.float32)
self.train_eval_train_test(BlacklistModule, torch.float32, "O1")
def test_promote_module_fp16_weight(self):
self.train_eval_train_test(PromoteModule, torch.float16)
self.train_eval_train_test(PromoteModule, torch.float16, "O1")
def test_promote_module_fp32_weight(self):
self.train_eval_train_test(PromoteModule, torch.float32, "O1")
# opt_level = O4
def test_whitelist_module_bfp16_weight(self):
self.train_eval_train_test(WhitelistModule, torch.bfloat16, "O4")
def test_whitelist_module_fp32_weight(self):
self.train_eval_train_test(WhitelistModule, torch.float32, "O4")
def test_blacklist_module_bfp16_weight(self):
self.train_eval_train_test(BlacklistModule, torch.bfloat16, "O4")
def test_blacklist_module_fp32_weight(self):
self.train_eval_train_test(BlacklistModule, torch.float32, "O4")
def test_promote_module_bfp16_weight(self):
self.train_eval_train_test(PromoteModule, torch.bfloat16, "O4")
def test_promote_module_fp32_weight(self):
self.train_eval_train_test(PromoteModule, torch.float32)
self.train_eval_train_test(PromoteModule, torch.float32, "O4")
if __name__ == '__main__':
......
tests/L0/run_amp/test_checkpointing.py
View file @ 6f7a8b39
......@@ -6,7 +6,7 @@ import torch.nn.functional as F
import torch.optim as optim
from apex import amp
from apex.testing.common_utils import skipIfRocm
from utils import common_init, FLOAT
......@@ -28,7 +28,7 @@ class MyModel(torch.nn.Module):
class TestCheckpointing(unittest.TestCase):
def setUp(self):
self.initial_lr = 1e-3
self.test_opt_levels = ("O0", "O1", "O2", "O3")
self.test_opt_levels = ("O0", "O1", "O2", "O3", "O4", "O5")
def seed(self):
torch.manual_seed(2809)
......@@ -161,6 +161,7 @@ class TestCheckpointing(unittest.TestCase):
# skip tests for different opt_levels
continue
@skipIfRocm
def test_loss_scale_decrease(self):
num_losses = 3
nb_decrease_loss_scales = [0, 1, 2]
......@@ -236,6 +237,7 @@ class TestCheckpointing(unittest.TestCase):
state_dict = model.state_dict()
for key in state_dict:
self.assertFalse('Half' in state_dict[key].type())
self.assertFalse('BFloat16' in state_dict[key].type())
# Check, if model is still trainable
# Create dummy data
......
tests/L0/run_amp/test_fused_sgd.py
View file @ 6f7a8b39
......@@ -13,6 +13,7 @@ from torch.nn import Parameter
from utils import common_init, HALF, FLOAT,\
ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
from apex.testing.common_utils import skipIfRocm
try:
import amp_C
......
tests/L0/run_amp/test_multi_tensor_axpby.py
View file @ 6f7a8b39
......@@ -12,6 +12,8 @@ from math import floor
from utils import common_init, HALF, FLOAT,\
ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
from apex.testing.common_utils import skipIfRocm
try:
import amp_C
from amp_C import multi_tensor_axpby
......@@ -69,7 +71,10 @@ class TestMultiTensorAxpby(unittest.TestCase):
applier(multi_tensor_axpby, self.overflow_buf, [x_list, y_list, out_list], self.a, self.b, -1)
self.assertTrue(all([torch.allclose(out, self.ref.to(out_type)) for out in out_list]),
# TODO: Remove this workaround for bfloat16 after torch.allcose() support bfloat16
if out_type == torch.bfloat16:
out_list = [out.float() for out in out_list]
self.assertTrue(all([torch.allclose(out, self.ref.to(out.dtype)) for out in out_list]),
msg="{} {} {} {} {} {} {}".format(sizea, sizeb, repeat_tensors,
x_type, y_type, out_type, inplace))
self.assertTrue(self.overflow_buf.item() == 0,
......@@ -119,9 +124,9 @@ class TestMultiTensorAxpby(unittest.TestCase):
for sizea, sizeb in input_size_pairs:
for applier in appliers:
for repeat in repeat_tensors:
for x_type in (torch.float32, torch.float16):
for y_type in (torch.float32, torch.float16):
for out_type in (torch.float32, torch.float16):
for x_type in (torch.float32, torch.float16, torch.bfloat16):
for y_type in (torch.float32, torch.float16, torch.bfloat16):
for out_type in (torch.float32, torch.float16, torch.bfloat16):
for inplace in (True, False):
if inplace is True and (y_type is not out_type):
continue
......@@ -137,6 +142,7 @@ class TestMultiTensorAxpby(unittest.TestCase):
@unittest.skipIf(disabled, "amp_C is unavailable")
@unittest.skipIf(not try_nhwc, "torch version is 1.4 or earlier, may not support nhwc")
@skipIfRocm
def test_fuzz_nhwc(self):
input_size_pairs = (
((7, 77, 7, 77), (5, 55, 5, 55)),
......
tests/L0/run_amp/test_multi_tensor_l2norm.py
View file @ 6f7a8b39
......@@ -11,6 +11,8 @@ import torch.nn.functional as F
from utils import common_init, HALF, FLOAT,\
ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
from apex.testing.common_utils import skipIfRocm
try:
import amp_C
from amp_C import multi_tensor_l2norm
......@@ -56,6 +58,7 @@ class TestMultiTensorL2Norm(unittest.TestCase):
self.assertTrue(self.overflow_buf.item() == 0)
@unittest.skipIf(disabled, "amp_C is unavailable")
@skipIfRocm
def test_fuzz(self):
input_size_pairs = (
(7777*77, 555*555),
......
tests/L0/run_amp/test_multi_tensor_scale.py
View file @ 6f7a8b39
......@@ -49,7 +49,10 @@ class TestMultiTensorScale(unittest.TestCase):
applier(multi_tensor_scale, self.overflow_buf, [in_list, out_list], 1./self.scale)
self.assertTrue(all([torch.allclose(out, self.ref.to(out_type)) for out in out_list]))
# TODO: Remove this workaround for bfloat16 after torch.allcose() support bfloat16
if out_type == torch.bfloat16:
out_list = [out.float() for out in out_list]
self.assertTrue(all([torch.allclose(out, self.ref.to(out.dtype)) for out in out_list]))
self.assertTrue(self.overflow_buf.item() == 0)
def find_inf(self, sizea, sizeb, applier, repeat_tensors, in_type, out_type, t, ind, val, inplace=False):
......@@ -106,8 +109,8 @@ class TestMultiTensorScale(unittest.TestCase):
for sizea, sizeb in input_size_pairs:
for applier in appliers:
for repeat in repeat_tensors:
for in_type in (torch.float32, torch.float16):
for out_type in (torch.float32, torch.float16):
for in_type in (torch.float32, torch.float16, torch.bfloat16):
for out_type in (torch.float32, torch.float16, torch.bfloat16):
for inplace in (True, False):
if inplace is True and (out_type is not in_type):
continue
......
tests/L0/run_amp/test_multiple_models_optimizers_losses.py
View file @ 6f7a8b39
......@@ -13,6 +13,8 @@ from torch.nn import Parameter
from utils import common_init, HALF, FLOAT,\
ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
from apex.testing.common_utils import skipIfRocm
class MyModel(torch.nn.Module):
def __init__(self, unique):
super(MyModel, self).__init__()
......@@ -41,7 +43,7 @@ class TestMultipleModelsOptimizersLosses(unittest.TestCase):
def tearDown(self):
pass
def test_2models2losses1optimizer(self):
model0 = MyModel(1)
model1 = MyModel(2)
......
tests/L0/run_amp/test_promotion.py
View file @ 6f7a8b39
......@@ -7,18 +7,18 @@ import torch
from torch import nn
import torch.nn.functional as F
from utils import common_init, HALF, FLOAT, DTYPES
from utils import common_init, HALF, FLOAT, DTYPES, DTYPES2, MATCH_INPUT
class TestPromotion(unittest.TestCase):
def setUp(self):
self.handle = amp.init(enabled=True)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def run_binary_promote_test(self, fns, input_shape, x_inplace=False):
type_pairs = it.product(DTYPES, DTYPES)
class _TestPromotion(unittest.TestCase):
def run_binary_promote_test(self, fns, input_shape, lp_type, x_inplace=False):
if lp_type == torch.half:
dtypes = DTYPES
elif lp_type == torch.bfloat16:
dtypes = DTYPES2
else:
raise RuntimeError("Creating test class with invalid low_precision type. \
Supported types are torch.half and torch.bfloat16")
type_pairs = it.product(dtypes, dtypes)
for fn, (xtype, ytype) in it.product(fns, type_pairs):
x = torch.randn(input_shape, dtype=xtype).requires_grad_()
x_leaf = x
......@@ -35,41 +35,78 @@ class TestPromotion(unittest.TestCase):
if xtype == torch.float or ytype == torch.float:
self.assertEqual(out.type(), FLOAT)
else:
self.assertEqual(out.type(), HALF)
self.assertEqual(out.type(), MATCH_INPUT[lp_type])
out.float().sum().backward()
self.assertEqual(x_leaf.grad.dtype, xtype)
def _test_cat_matches_widest(self, lp_type):
shape = self.b
ys = [torch.randn(shape, dtype=lp_type) for _ in range(5)]
x_float = torch.randn(shape)
out = torch.cat(ys + [x_float])
self.assertEqual(out.type(), FLOAT)
x_lp = torch.randn(shape, dtype=lp_type)
out = torch.cat(ys + [x_lp])
self.assertEqual(out.type(), MATCH_INPUT[lp_type])
def _test_inplace_exp_is_error_for_lp(self, lp_type):
xs = torch.randn(self.b)
xs.exp_()
self.assertEqual(xs.type(), FLOAT)
xs = torch.randn(self.b, dtype=lp_type)
with self.assertRaises(NotImplementedError):
xs.exp_()
class TestPromotionHalf(_TestPromotion):
def setUp(self):
self.handle = amp.init(enabled=True, patch_type=torch.half)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def test_atan2_matches_widest(self):
fns = [lambda x, y : torch.atan2(x, y),
lambda x, y : x.atan2(y)]
self.run_binary_promote_test(fns, (self.b,))
self.run_binary_promote_test(fns, (self.b,), torch.half)
def test_mul_matches_widest(self):
fns = [lambda x, y : torch.mul(x, y),
lambda x, y: x.mul(y)]
self.run_binary_promote_test(fns, (self.b,))
self.run_binary_promote_test(fns, (self.b,), torch.half)
def test_cat_matches_widest(self):
shape = self.b
ys = [torch.randn(shape, dtype=torch.half) for _ in range(5)]
x_float = torch.randn(shape)
out = torch.cat(ys + [x_float])
self.assertEqual(out.type(), FLOAT)
x_half = torch.randn(shape, dtype=torch.half)
out = torch.cat(ys + [x_half])
self.assertEqual(out.type(), HALF)
self._test_cat_matches_widest(torch.half)
def test_inplace_exp_is_error_for_half(self):
xs = torch.randn(self.b)
xs.exp_()
self.assertEqual(xs.type(), FLOAT)
xs = torch.randn(self.b, dtype=torch.half)
with self.assertRaises(NotImplementedError):
xs.exp_()
self._test_inplace_exp_is_error_for_lp(torch.half)
def test_inplace_add_matches_self(self):
fn = lambda x, y: x.add_(y)
self.run_binary_promote_test([fn], (self.b,), torch.half, x_inplace=True)
class TestPromotionBFloat16(_TestPromotion):
def setUp(self):
self.handle = amp.init(enabled=True, patch_type=torch.bfloat16)
common_init(self)
def tearDown(self):
self.handle._deactivate()
def test_mul_matches_widest(self):
fns = [lambda x, y : torch.mul(x, y),
lambda x, y: x.mul(y)]
self.run_binary_promote_test(fns, (self.b,), torch.bfloat16)
def test_cat_matches_widest(self):
self._test_cat_matches_widest(torch.bfloat16)
def test_inplace_exp_is_error_for_bfloat16(self):
self._test_inplace_exp_is_error_for_lp(torch.bfloat16)
def test_inplace_add_matches_self(self):
fn = lambda x, y: x.add_(y)
self.run_binary_promote_test([fn], (self.b,), x_inplace=True)
self.run_binary_promote_test([fn], (self.b,), torch.bfloat16, x_inplace=True)
if __name__ == '__main__':
unittest.main()
tests/L0/run_amp/test_rnn.py
View file @ 6f7a8b39
......@@ -6,6 +6,7 @@ import torch
from torch import nn
from utils import common_init, HALF
from apex.testing.common_utils import skipIfRocm
class TestRnnCells(unittest.TestCase):
def setUp(self):
......@@ -73,6 +74,7 @@ class TestRnns(unittest.TestCase):
output[-1, :, :].float().sum().backward()
self.assertEqual(x.grad.dtype, x.dtype)
@skipIfRocm
def test_rnn_is_half(self):
configs = [(1, False), (2, False), (2, True)]
for layers, bidir in configs:
......@@ -80,6 +82,7 @@ class TestRnns(unittest.TestCase):
nonlinearity='relu', bidirectional=bidir)
self.run_rnn_test(rnn, layers, bidir)
@skipIfRocm
def test_gru_is_half(self):
configs = [(1, False), (2, False), (2, True)]
for layers, bidir in configs:
......@@ -87,6 +90,7 @@ class TestRnns(unittest.TestCase):
bidirectional=bidir)
self.run_rnn_test(rnn, layers, bidir)
@skipIfRocm
def test_lstm_is_half(self):
configs = [(1, False), (2, False), (2, True)]
for layers, bidir in configs:
......@@ -94,6 +98,7 @@ class TestRnns(unittest.TestCase):
bidirectional=bidir)
self.run_rnn_test(rnn, layers, bidir, state_tuple=True)
@skipIfRocm
def test_rnn_packed_sequence(self):
num_layers = 2
rnn = nn.RNN(input_size=self.h, hidden_size=self.h, num_layers=num_layers)
......
tests/L0/run_amp/utils.py
View file @ 6f7a8b39
......@@ -2,15 +2,21 @@ import torch
HALF = 'torch.cuda.HalfTensor'
FLOAT = 'torch.cuda.FloatTensor'
BFLOAT16 = 'torch.cuda.BFloat16Tensor'
DTYPES = [torch.half, torch.float]
DTYPES2 = [torch.bfloat16, torch.float]
ALWAYS_HALF = {torch.float: HALF,
torch.half: HALF}
ALWAYS_BFLOAT16 = {torch.bfloat16: BFLOAT16,
torch.float: BFLOAT16}
ALWAYS_FLOAT = {torch.float: FLOAT,
torch.half: FLOAT}
MATCH_INPUT = {torch.float: FLOAT,
torch.half: HALF}
torch.half: HALF,
torch.bfloat16: BFLOAT16}
def common_init(test_case):
test_case.h = 64
......
tests/L0/run_optimizers/test_adagrad.py
View file @ 6f7a8b39
......@@ -14,22 +14,28 @@ class TestFusedAdagrad(unittest.TestCase):
def tearDown(self):
pass
def gen_param_optim(self, tensors, adagrad_option):
def gen_param_optim(self, tensors, adagrad_option, apex_only=False):
ref_param = []
tst_param = []
for tensor in tensors:
ref_param.append(torch.nn.Parameter(tensor.clone()))
if apex_only:
ref_param.append(torch.nn.Parameter(tensor.clone().float()))
else:
ref_param.append(torch.nn.Parameter(tensor.clone()))
tst_param.append(torch.nn.Parameter(tensor.clone()))
ref_optim = torch.optim.Adagrad(ref_param, **adagrad_option)
if apex_only:
ref_optim = apex.optimizers.FusedAdagrad(ref_param, **adagrad_option)
else:
ref_optim = torch.optim.Adagrad(ref_param, **adagrad_option)
tst_optim = apex.optimizers.FusedAdagrad(tst_param, **adagrad_option)
return (ref_param, tst_param, ref_optim, tst_optim)
def gen_grad(self, ref_param, tst_param):
def gen_grad(self, ref_param, tst_param, apex_only=False):
for p_ref, p_tst in zip(ref_param, tst_param):
p_ref.grad = torch.rand_like(p_ref)
p_tst.grad = p_ref.grad
p_tst.grad = torch.rand_like(p_tst)
p_ref.grad = p_tst.grad.detach().float() if apex_only else p_tst.grad
def gen_mixed_grad(self, ref_param, tst_param, scale=1.0):
half_grads = []
......@@ -38,9 +44,11 @@ class TestFusedAdagrad(unittest.TestCase):
p_ref.grad = half_grads[-1].float() / scale
return half_grads
def get_max_diff(self, ref_param, tst_param):
def get_max_diff(self, ref_param, tst_param, apex_only=False):
max_abs_diff = max_rel_diff = 0
for p_ref, p_tst in zip(ref_param, tst_param):
if apex_only:
p_tst = p_tst.float()
max_abs_diff_p = (p_ref - p_tst).abs().max().item()
max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
......@@ -51,23 +59,24 @@ class TestFusedAdagrad(unittest.TestCase):
return max_abs_diff, max_rel_diff
def gen_single_type_test(self, param_type=torch.float):
def gen_single_type_test(self, param_type=torch.float, apex_only=False):
nelem = 278011
adagrad_option = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 1.0e-5}
tensor = torch.rand(nelem, dtype=param_type, device="cuda")
ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
[tensor], adagrad_option
[tensor], adagrad_option, apex_only=apex_only
)
for _ in range(self.iters):
self.gen_grad(ref_param, tst_param)
self.gen_grad(ref_param, tst_param, apex_only=apex_only)
ref_optim.step()
tst_optim.step()
max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param, apex_only=apex_only)
self.assertLessEqual(max_abs_diff, self.max_abs_diff)
self.assertLessEqual(max_rel_diff, self.max_rel_diff)
if not apex_only:
self.assertLessEqual(max_rel_diff, self.max_rel_diff)
def test_float(self):
self.gen_single_type_test(param_type=torch.float)
......@@ -76,6 +85,14 @@ class TestFusedAdagrad(unittest.TestCase):
def test_half(self):
self.gen_single_type_test(param_type=torch.float16)
# Compares bfloat16 computation against float32 as gold standard.
# Uses apex optimizers(controlled by apex_only flag) for both types.
# Doesn't use upstream optimizer like other tests as they seem to be
# numerically unstable for half types(see skip note for test above).
def test_bfloat16(self):
self.max_abs_diff = 1e-2
self.gen_single_type_test(param_type=torch.bfloat16, apex_only=True)
def test_multi_params(self):
sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
adagrad_option = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 0}
......
tests/L0/run_optimizers/test_adam.py
View file @ 6f7a8b39
......@@ -15,22 +15,28 @@ class TestFusedAdam(unittest.TestCase):
def tearDown(self):
pass
def gen_param_optim(self, tensors, adam_option):
def gen_param_optim(self, tensors, adam_option, apex_only=False):
ref_param = []
tst_param = []
for tensor in tensors:
ref_param.append(torch.nn.Parameter(tensor.clone()))
if apex_only:
ref_param.append(torch.nn.Parameter(tensor.clone().float()))
else:
ref_param.append(torch.nn.Parameter(tensor.clone()))
tst_param.append(torch.nn.Parameter(tensor.clone()))
ref_optim = torch.optim.Adam(ref_param, **adam_option)
if apex_only:
ref_optim = apex.optimizers.FusedAdam(ref_param, **adam_option)
else:
ref_optim = torch.optim.Adam(ref_param, **adam_option)
tst_optim = apex.optimizers.FusedAdam(tst_param, **adam_option)
return (ref_param, tst_param, ref_optim, tst_optim)
def gen_grad(self, ref_param, tst_param):
def gen_grad(self, ref_param, tst_param, apex_only=False):
for p_ref, p_tst in zip(ref_param, tst_param):
p_ref.grad = torch.rand_like(p_ref)
p_tst.grad = p_ref.grad
p_tst.grad = torch.rand_like(p_tst)
p_ref.grad = p_tst.grad.detach().float() if apex_only else p_tst.grad
def gen_mixed_grad(self, ref_param, tst_param, scale=1.0):
half_grads = []
......@@ -39,9 +45,11 @@ class TestFusedAdam(unittest.TestCase):
p_ref.grad = half_grads[-1].float() / scale
return half_grads
def get_max_diff(self, ref_param, tst_param):
def get_max_diff(self, ref_param, tst_param, apex_only=False):
max_abs_diff = max_rel_diff = 0
for p_ref, p_tst in zip(ref_param, tst_param):
if apex_only:
p_tst = p_tst.float()
max_abs_diff_p = (p_ref - p_tst).abs().max().item()
max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
......@@ -50,23 +58,24 @@ class TestFusedAdam(unittest.TestCase):
return max_abs_diff, max_rel_diff
def gen_single_type_test(self, param_type=torch.float):
def gen_single_type_test(self, param_type=torch.float, apex_only=False):
nelem = 278011
adam_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08,
'weight_decay':0, 'amsgrad':False}
tensor = torch.rand(nelem, dtype=param_type, device='cuda')
ref_param, tst_param, ref_optim, tst_optim = \
self.gen_param_optim([tensor], adam_option)
self.gen_param_optim([tensor], adam_option, apex_only=apex_only)
for i in range(self.iters):
self.gen_grad(ref_param, tst_param)
self.gen_grad(ref_param, tst_param, apex_only=apex_only)
ref_optim.step()
tst_optim.step()
max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param, apex_only=apex_only)
self.assertLessEqual(max_abs_diff, self.max_abs_diff)
self.assertLessEqual(max_rel_diff, self.max_rel_diff)
if not apex_only:
self.assertLessEqual(max_rel_diff, self.max_rel_diff)
def test_float(self):
self.gen_single_type_test(param_type=torch.float)
......@@ -74,6 +83,14 @@ class TestFusedAdam(unittest.TestCase):
def test_half(self):
self.gen_single_type_test(param_type=torch.float16)
# Compares bfloat16 computation against float32 as gold standard.
# Uses apex optimizers(controlled by apex_only flag) for both types.
# Doesn't use upstream optimizer like other tests as they seem to be
# numerically unstable for half types
def test_bfloat16(self):
self.max_abs_diff = 1e-2
self.gen_single_type_test(param_type=torch.bfloat16, apex_only=True)
@unittest.skip('Disable until 8/1/2019 adam/adamw upstream picked')
def test_multi_params(self):
sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
......
tests/L0/run_optimizers/test_lamb.py
View file @ 6f7a8b39
......@@ -5,6 +5,7 @@ import torch
from torch.optim import Optimizer
import apex
from apex.multi_tensor_apply import multi_tensor_applier
from apex.testing.common_utils import skipIfRocm
class RefLAMB(Optimizer):
r"""Implements Lamb algorithm.
......@@ -207,6 +208,7 @@ class TestFusedLAMB(unittest.TestCase):
self.assertLessEqual(max_abs_diff, self.max_abs_diff)
self.assertLessEqual(max_rel_diff, self.max_rel_diff)
@skipIfRocm
def test_float(self):
self.gen_single_type_test(param_type=torch.float)
......@@ -214,6 +216,7 @@ class TestFusedLAMB(unittest.TestCase):
def test_half(self):
self.gen_single_type_test(param_type=torch.float16)
@skipIfRocm
def test_multi_params(self):
sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
weight_decay = [0, 0.01]
......@@ -234,6 +237,7 @@ class TestFusedLAMB(unittest.TestCase):
self.assertLessEqual(max_abs_diff, self.max_abs_diff)
self.assertLessEqual(max_rel_diff, self.max_rel_diff)
@skipIfRocm
def test_lamb_option(self):
nelem = 1
tensor = torch.rand(nelem, dtype=torch.float, device='cuda')
......
tests/L0/run_rocm.sh 0 → 100755
View file @ 6f7a8b39
#!/bin/bash
APEX_TEST_WITH_ROCM=1 python3.6 run_test.py
tests/L0/run_test.py
View file @ 6f7a8b39
import unittest
import sys
from apex.testing.common_utils import TEST_WITH_ROCM, skipIfRocm
test_dirs = ["run_amp", "run_fp16util", "run_optimizers", "run_fused_layer_norm", "run_pyprof_nvtx", "run_pyprof_data", "run_mlp"]
ROCM_BLACKLIST = [
'run_fused_layer_norm',
'run_pyprof_nvtx',
'run_pyprof_data',
'run_mlp'
]
runner = unittest.TextTestRunner(verbosity=2)
errcode = 0
for test_dir in test_dirs:
if (test_dir in ROCM_BLACKLIST) and TEST_WITH_ROCM:
continue
suite = unittest.TestLoader().discover(test_dir)
print("\nExecuting tests from " + test_dir)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment