Efficient Multi-Head Attention (EMHA) support

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

.github/workflows/build.yml

@@ -17,6 +17,8 @@ jobs:
  steps:
  - name: 'Checkout'
  uses: actions/checkout@v3
+ with:
+ submodules: 'true'
  - name: 'Build'
  run: |
  mkdir -p wheelhouse && \

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "3rdparty/googletest"]
 	path = 3rdparty/googletest
 	url = https://github.com/google/googletest.git
+[submodule "3rdparty/cutlass"]
+	path = 3rdparty/cutlass
+	url = https://github.com/NVIDIA/cutlass.git

qa/L0_lint/pylintrc

@@ -1,6 +1,8 @@
 [MASTER]
 extension-pkg-whitelist=torch,
- transformer_engine_extensions
+ transformer_engine_extensions,
+ scaled_upper_triang_masked_softmax_dropout_cuda,
+ emha_C
 
 disable=too-many-locals,
  invalid-name,

setup.py

@@ -21,6 +21,7 @@
 with open(path + "/VERSION", "r") as f:
  te_version = f.readline()
 
+
 def get_cuda_bare_metal_version(cuda_dir):
  raw_output = subprocess.check_output(
  [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
@@ -137,6 +138,71 @@ def __init__(self, name, cmake_path, sources, **kwargs):
  )
  )
 
+ ext_modules.append(
+ CUDAExtension(
+ name="scaled_upper_triang_masked_softmax_dropout_cuda",
+ sources=[
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout.cpp",
+ ),
+ os.path.join(
+ path,
+ "transformer_engine/pytorch/csrc/emha/scaled_upper_triang_masked_softmax_dropout_cuda.cu",
+ ),
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
+ },
+ include_dirs=[
+ os.path.join(path, "transformer_engine/pytorch/csrc/emha"),
+ ],
+ )
+ )
+
+ # EMHA cannot be compiled for sm70 as it requires hardware support of bfloat16.
+ ext_modules.append(
+ CUDAExtension(
+ "emha_C",
+ sources=[
+ os.path.join(path, f)
+ for f in (
+ "transformer_engine/pytorch/csrc/emha/bmm_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/emha_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/softmax_api.cpp",
+ "transformer_engine/pytorch/csrc/emha/bmm_nn.cu",
+ "transformer_engine/pytorch/csrc/emha/bmm_nt.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_bwd_kernel.cu",
+ "transformer_engine/pytorch/csrc/emha/softmax_fwd_kernel.cu",
+ )
+ ],
+ extra_compile_args={
+ "cxx": ["-O3"],
+ "nvcc": [
+ "-O3",
+ "-U__CUDA_NO_HALF_OPERATORS__",
+ "-U__CUDA_NO_HALF_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+ "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+ "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+ "--expt-relaxed-constexpr",
+ "--expt-extended-lambda",
+ "--use_fast_math",
+ "-gencode",
+ "arch=compute_80,code=sm_80",
+ "-gencode",
+ "arch=compute_90,code=sm_90",
+ ],
+ },
+ include_dirs=[
+ os.path.join(path, "3rdparty/cutlass/include"),
+ os.path.join(path, "3rdparty/cutlass/tools/util/include"),
+ ],
+ ),
+ )
+
 
 def get_cmake_bin():
  cmake_bin = "cmake"

tests/test_fp8.py

@@ -465,3 +465,72 @@ def test_sanity_fused_qkv_params(dtype, bs, fp8_recipe, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, fp8_recipe, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("fp8_recipe", fp8_recipes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, fp8_recipe, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ with fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()

tests/test_transformerengine.py

@@ -408,3 +408,70 @@ def test_sanity_fused_qkv_params(dtype, bs, model, skip_wgrad):
  )
 
  _test_sanity_e2e(block, bs, dtype, config, skip_wgrad)
+
+
+@pytest.mark.parametrize("dtype", param_types)
+@pytest.mark.parametrize("bs", batch_sizes)
+@pytest.mark.parametrize("model", model_configs.keys())
+@pytest.mark.parametrize("skip_wgrad", skip_wgrad)
+def test_gpt_with_emha(dtype, bs, model, skip_wgrad):
+ if dtype == torch.float32:
+ return
+
+ try:
+ from transformer_engine.pytorch import emha # noqa: F401
+ except ImportError:
+ return
+ else:
+ if not (
+ torch.cuda.is_available()
+ and torch.cuda.get_device_capability("cuda") >= (8, 0)
+ ):
+ return
+ config = model_configs[model]
+
+ sigma = 0.023
+ init_method = init_method_normal(sigma)
+ output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers)
+
+ block = TransformerLayer(
+ config.hidden_size,
+ 4 * config.hidden_size,
+ config.num_attention_heads,
+ layernorm_epsilon=config.eps,
+ init_method=init_method,
+ output_layer_init_method=output_layer_init_method,
+ hidden_dropout=0.1,
+ attention_dropout=0.1,
+ kv_channels=config.embed,
+ apply_residual_connection_post_layernorm=False,
+ output_layernorm=False,
+ use_emha=True,
+ ).to(dtype=dtype, device="cuda")
+
+ te_inp_hidden_states = torch.randn(
+ config.seq_len,
+ bs,
+ config.hidden_size,
+ dtype=dtype,
+ device="cuda",
+ requires_grad=True,
+ )
+
+ if skip_wgrad:
+ _disable_wgrads(block)
+
+ with torch.no_grad():
+ te_inp_attn_mask = [0] + [config.seq_len for _ in range(bs)]
+ te_inp_attn_mask = torch.tensor(
+ te_inp_attn_mask, dtype=torch.int32, device="cuda"
+ )
+ te_inp_attn_mask = torch.cumsum(te_inp_attn_mask, dim=0)
+
+ with torch.cuda.amp.autocast(dtype=dtype):
+ te_out = block(te_inp_hidden_states, te_inp_attn_mask)
+ loss = te_out.mean()
+
+ assert te_out.dtype == dtype
+ loss.backward()
+ torch.cuda.synchronize()