Skip to content

Increase number of FP8 tensors per GEMM#22

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
ptrendx merged 12 commits into from
Feb 3, 2023
Merged

Increase number of FP8 tensors per GEMM #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
95910d8
Increase number of FP8 tensors per GEMM
vasunvidia Oct 28, 2022
fcbce36
Enable FP8 output tensor for fp8_gemm
vasunvidia Nov 7, 2022
8a336ab
[BERT FP8] Initial TE review comments
vasunvidia Dec 14, 2022
320182a
Temporary fix for cuda graph non convergence
vasunvidia Dec 15, 2022
ffe4852
Address review comments-2
vasunvidia Jan 8, 2023
075e287
Review comments-3
vasunvidia Jan 28, 2023
784fb47
Cleanup
vasunvidia Jan 28, 2023
0ccb2e5
Change for New API
vasunvidia Jan 29, 2023
397a1e4
Remove unnecessary clone for D_scale, D_amax
vasunvidia Jan 31, 2023
a575b5e
Avoid Roll for AMAX history size = 1
vasunvidia Feb 1, 2023
2dad360
Update onnx_te_gemm API
vasunvidia Feb 2, 2023
66b98e3
Fix Lint errors
vasunvidia Feb 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 30 additions & 4 deletions transformer_engine/common/gemm/cublaslt_gemm.cu
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
void *A_scale_inverse = inputA->scale_inv.dptr;
void *B = inputB->data.dptr;
void *B_scale_inverse = inputB->scale_inv.dptr;
void *C = outputD->data.dptr;
void *D = outputD->data.dptr;
void *D_scale = outputD->scale.dptr;
void *D_amax = outputD->amax.dptr;
void *bias_ptr = inputBias->data.dptr;
const bool bias = bias_ptr != nullptr;
void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
if (use_fp8) {
NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
}
if (is_fp8_dtype(outputD->data.dtype)) {
NVTE_CHECK(!accumulate,
"Accumulation mode not supported with FP8 GEMM output!");
}

float one = 1.0;
float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

cublasLtMatmulDesc_t operationDesc = nullptr;
cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
cublasLtMatmulPreference_t preference = nullptr;
int returnedResults = 0;
cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
&B_scale_inverse,
sizeof(B_scale_inverse)));
if (is_fp8_dtype(outputD->data.dtype)) {
// Accumulation mode not supported for FP8 output
C = nullptr;
NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
&D_scale,
sizeof(D_scale)));
NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
&D_amax,
sizeof(D_amax)));
// For FP8 output, cuBLAS requires C_type to be same as bias_type
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.
Outdated
} else {
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
}
if (bias) {
NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
&bias_type, sizeof(bias_type)));
}
} else {
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
}

if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&workspaceSize, sizeof(workspaceSize)));

NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
Ddesc, preference, 1, &heuristicResult,
&returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
B, /* B */
Bdesc,
static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
D, /* D */
Ddesc,
&heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down
15 changes: 15 additions & 0 deletions transformer_engine/pytorch/cpp_extensions.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
workspace: torch.Tensor,
accumulate: bool = False,
out: Optional[torch.Tensor] = None,
out_index = None,
fp8_meta_tensor: tex.FP8TensorMeta = None,
bias: Optional[torch.Tensor] = None,
use_bias: bool = False,
fp32_output: bool = False,
use_split_accumulator: bool = False,
D_dtype: tex.DType = None,
) -> torch.Tensor:
"""TN layout GEMM with fp8 inputs."""

empty_tensor = torch.Tensor()
if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
assert fp8_meta_tensor is not None and out_index is not None

return_output = False
if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
return_output = True

out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.
# Use bfloat16 as default bias_dtype
bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
out_dtype = D_dtype if D_dtype is not None else out_dtype

_ = torch.ops.tex_ts.te_gemm_ts(
A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
B_dtype,
False, # transb
out,
empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
out_dtype,
empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
bias if use_bias else empty_tensor,
bias_dtype,
empty_tensor, # this is pre_gelu_out
False, # grad
workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

input_dtype = TE_DType[dtype]
output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

return_output = False
if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
input_dtype,
transb,
out,
empty_tensor, # out_scale
output_dtype,
empty_tensor, # out_amax
grad_bias if grad else bias,
bias_dtype,
gelu_input,
grad,
workspace,
Expand Down
10 changes: 7 additions & 3 deletions transformer_engine/pytorch/csrc/common.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
GEMM1_INPUT = 0,
GEMM1_WEIGHT = 1,
GEMM2_INPUT = 2,
GEMM2_WEIGHT = 3
GEMM1_OUTPUT = 2,
GEMM2_INPUT = 3,
GEMM2_WEIGHT = 4,
GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
GRAD_OUTPUT1 = 0,
GRAD_OUTPUT2 = 1
GRAD_INPUT1 = 1,
GRAD_OUTPUT2 = 2,
GRAD_INPUT2 = 3
};


Expand Down
16 changes: 12 additions & 4 deletions transformer_engine/pytorch/csrc/extensions.cu
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
transformer_engine::DType B_type,
bool transb,
at::Tensor D,
at::Tensor D_scale,
transformer_engine::DType D_type,
at::Tensor D_amax,
at::Tensor bias,
transformer_engine::DType bias_type,
at::Tensor pre_gelu_out,
bool grad,
at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
auto te_D = makeTransformerEngineTensor(D.data_ptr(),
{static_cast<size_t>(D.size(0)),
static_cast<size_t>(D.size(1))},
D_type);
D_type, D_amax.data_ptr(),
D_scale.data_ptr(), nullptr);
auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
.value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
.value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
.value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
.value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
.value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
.value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
.value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
.value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
.value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
.value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
.value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
.value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}
3 changes: 3 additions & 0 deletions transformer_engine/pytorch/csrc/extensions.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
transformer_engine::DType B_type,
bool transb,
at::Tensor D,
at::Tensor D_scale,
transformer_engine::DType D_type,
at::Tensor D_amax,
at::Tensor bias,
transformer_engine::DType bias_type,
at::Tensor pre_gelu_out,
bool grad,
at::Tensor workspace,
Expand Down
7 changes: 7 additions & 0 deletions transformer_engine/pytorch/csrc/ts_fp8_op.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
int64_t B_type,
int64_t transb,
at::Tensor D,
at::Tensor D_scale,
int64_t D_type,
at::Tensor D_amax,
at::Tensor bias,
int64_t bias_type,
at::Tensor pre_gelu_out,
int64_t grad,
at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
bool transb_arg = static_cast<bool>(transb);
transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
bool grad_arg = static_cast<bool>(grad);
size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
B_type_arg,
transb_arg,
D,
D_scale,
D_type_arg,
D_amax,
bias,
bias_type_arg,
pre_gelu_out,
grad_arg,
workspace,
Expand Down
3 changes: 2 additions & 1 deletion transformer_engine/pytorch/fp8.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
"""Update amax history and set next amax to zero."""
amax_history = torch.roll(amax_history, -1, 0)
if amax_history.shape[0] > 1:
amax_history = torch.roll(amax_history, -1, 0)
amax_history[0].fill_(0.0)
return amax_history

Expand Down
4 changes: 3 additions & 1 deletion transformer_engine/pytorch/module.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
def set_meta_tensor(self, fwd: bool) -> None:
"""Init scales and amaxes for fwd | bwd."""
fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
# Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
# 2 (grad_output and grad_input) for bwd
num_fp8_tensors = (
self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.
Outdated
)

self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down
5 changes: 4 additions & 1 deletion transformer_engine/pytorch/te_onnx_extensions.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
"v", "fs", "i", "i", "i",
"v", "i", "v", "v", "i",
"v", "fs", "i", "fs", "v", "i", "v", "i",
"v", "i", "i", "i")
def onnx_te_gemm(
g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
input_type,
trans_input,
out,
out_scale,
out_type,
out_amax,
bias,
bias_type,
pre_gelu_out,
grad,
workspace,
Expand Down
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
Increase number of FP8 tensors per GEMM by vasunvidia · Pull Request #22 · NVIDIA/TransformerEngine · GitHub
Skip to content
Increase number of FP8 tensors per GEMM#22

New issue

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

 Sign up for GitHub

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

Already on GitHub?
Sign in
to your account

 Merged

ptrendx
merged 12 commits into

from

Feb 3, 2023
 Merged


Increase number of FP8 tensors per GEMM
#22
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


12 commits


Select commit
Hold shift + click to select a range
95910d8

Increase number of FP8 tensors per GEMM


vasunvidia Oct 28, 2022
fcbce36

Enable FP8 output tensor for fp8_gemm


vasunvidia Nov 7, 2022
8a336ab

[BERT FP8] Initial TE review comments


vasunvidia Dec 14, 2022
320182a

Temporary fix for cuda graph non convergence


vasunvidia Dec 15, 2022
ffe4852

Address review comments-2


vasunvidia Jan 8, 2023
075e287

Review comments-3


vasunvidia Jan 28, 2023
784fb47

Cleanup


vasunvidia Jan 28, 2023
0ccb2e5

Change for New API


vasunvidia Jan 29, 2023
397a1e4

Remove unnecessary clone for D_scale, D_amax


vasunvidia Jan 31, 2023
a575b5e

Avoid Roll for AMAX history size = 1


vasunvidia Feb 1, 2023
2dad360

Update onnx_te_gemm API


vasunvidia Feb 2, 2023
66b98e3

Fix Lint errors


vasunvidia Feb 3, 2023

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


34 changes: 30 additions & 4 deletions
transformer_engine/common/gemm/cublaslt_gemm.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -55,7 +55,10 @@ void cublas_gemm(const Tensor *inputA,
 void *A_scale_inverse = inputA->scale_inv.dptr;
 void *B = inputB->data.dptr;
 void *B_scale_inverse = inputB->scale_inv.dptr;
 void *C = outputD->data.dptr;
 void *D = outputD->data.dptr;
 void *D_scale = outputD->scale.dptr;
 void *D_amax = outputD->amax.dptr;
 void *bias_ptr = inputBias->data.dptr;
 const bool bias = bias_ptr != nullptr;
 void *pre_gelu_out = outputPreGelu->data.dptr;
Expand All@@ -78,6 +81,10 @@ void cublas_gemm(const Tensor *inputA,
 if (use_fp8) {
 NVTE_CHECK(!gelu, "fp8 gemm + gelu fusion is unavailable right now!");
 }
 if (is_fp8_dtype(outputD->data.dtype)) {
 NVTE_CHECK(!accumulate,
 "Accumulation mode not supported with FP8 GEMM output!");
 }

 float one = 1.0;
 float zero = 0.0;
Expand All@@ -87,7 +94,7 @@ void cublas_gemm(const Tensor *inputA,
 NVTE_CHECK_CUBLAS(cublasLtCreate(&handle));

 cublasLtMatmulDesc_t operationDesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Ddesc = nullptr;
 cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr, Ddesc = nullptr;
 cublasLtMatmulPreference_t preference = nullptr;
 int returnedResults = 0;
 cublasLtMatmulHeuristicResult_t heuristicResult = {};
Expand DownExpand Up@@ -135,11 +142,29 @@ void cublas_gemm(const Tensor *inputA,
 CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
 &B_scale_inverse,
 sizeof(B_scale_inverse)));
 if (is_fp8_dtype(outputD->data.dtype)) {
 // Accumulation mode not supported for FP8 output
 C = nullptr;
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_D_SCALE_POINTER,
 &D_scale,
 sizeof(D_scale)));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_AMAX_D_POINTER,
 &D_amax,
 sizeof(D_amax)));
 // For FP8 output, cuBLAS requires C_type to be same as bias_type
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, bias_type, m, n, ldd));
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }
 if (bias) {
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc,
 CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE,
 &bias_type, sizeof(bias_type)));
 }
 } else {
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutCreate(&Cdesc, D_type, m, n, ldd));
 }

 if (bias && gelu) {
Expand DownExpand Up@@ -190,7 +215,7 @@ void cublas_gemm(const Tensor *inputA,
 preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
 &workspaceSize, sizeof(workspaceSize)));

 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Ddesc,
 NVTE_CHECK_CUBLAS(cublasLtMatmulAlgoGetHeuristic(handle, operationDesc, Adesc, Bdesc, Cdesc,
 Ddesc, preference, 1, &heuristicResult,
 &returnedResults));

Expand All@@ -205,8 +230,8 @@ void cublas_gemm(const Tensor *inputA,
 B, /* B */
 Bdesc,
 static_cast<const void*>(&beta), /* beta */
D, /* C */
Ddesc,
C, /* C */
Cdesc,
 D, /* D */
 Ddesc,
 &heuristicResult.algo, /* algo */
Expand All@@ -217,6 +242,7 @@ void cublas_gemm(const Tensor *inputA,

 NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceDestroy(preference));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Ddesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Cdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Bdesc));
 NVTE_CHECK_CUBLAS(cublasLtMatrixLayoutDestroy(Adesc));
 NVTE_CHECK_CUBLAS(cublasLtMatmulDescDestroy(operationDesc));
Expand Down

15 changes: 15 additions & 0 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -22,14 +22,19 @@ def fp8_gemm(
 workspace: torch.Tensor,
 accumulate: bool = False,
 out: Optional[torch.Tensor] = None,
 out_index = None,
 fp8_meta_tensor: tex.FP8TensorMeta = None,
 bias: Optional[torch.Tensor] = None,
 use_bias: bool = False,
 fp32_output: bool = False,
 use_split_accumulator: bool = False,
 D_dtype: tex.DType = None,
) -> torch.Tensor:
 """TN layout GEMM with fp8 inputs."""

 empty_tensor = torch.Tensor()
 if D_dtype is not None and D_dtype in [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]:
 assert fp8_meta_tensor is not None and out_index is not None

 return_output = False
 if out is None:
Expand All@@ -42,6 +47,9 @@ def fp8_gemm(
 return_output = True

 out_dtype = tex.DType.kFloat32 if fp32_output else TE_DType[out_dtype]
Comment thread
vasunvidia marked this conversation as resolved.

 # Use bfloat16 as default bias_dtype
 bias_dtype = tex.DType.kBFloat16 if bias is None else TE_DType[bias.dtype]
 out_dtype = D_dtype if D_dtype is not None else out_dtype

 _ = torch.ops.tex_ts.te_gemm_ts(
 A,
Expand All@@ -55,8 +63,11 @@ def fp8_gemm(
 B_dtype,
 False, # transb
 out,
 empty_tensor if out_index is None else fp8_meta_tensor.scale[out_index],
 out_dtype,
 empty_tensor if out_index is None else fp8_meta_tensor.amax_history[0][out_index],
 bias if use_bias else empty_tensor,
 bias_dtype,
 empty_tensor, # this is pre_gelu_out
 False, # grad
 workspace,
Expand DownExpand Up@@ -95,6 +106,7 @@ def gemm(

 input_dtype = TE_DType[dtype]
 output_dtype = tex.DType.kFloat32 if fp32_output else input_dtype
 bias_dtype = output_dtype if bias is None else TE_DType[bias.dtype]

 return_output = False
 if out is None:
Expand DownExpand Up@@ -132,8 +144,11 @@ def gemm(
 input_dtype,
 transb,
 out,
 empty_tensor, # out_scale
 output_dtype,
 empty_tensor, # out_amax
 grad_bias if grad else bias,
 bias_dtype,
 gelu_input,
 grad,
 workspace,
Expand Down

10 changes: 7 additions & 3 deletions
transformer_engine/pytorch/csrc/common.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -48,15 +48,19 @@ class FP8TensorMeta {
enum FP8FwdTensors {
 GEMM1_INPUT = 0,
 GEMM1_WEIGHT = 1,
 GEMM2_INPUT = 2,
 GEMM2_WEIGHT = 3
 GEMM1_OUTPUT = 2,
 GEMM2_INPUT = 3,
 GEMM2_WEIGHT = 4,
 GEMM2_OUTPUT = 5
};

// Used as named indices on the `scale`, `scale_inv`,
// and `amax` tensors in the `FP8TensorMeta` class.
enum FP8BwdTensors {
 GRAD_OUTPUT1 = 0,
 GRAD_OUTPUT2 = 1
 GRAD_INPUT1 = 1,
 GRAD_OUTPUT2 = 2,
 GRAD_INPUT2 = 3
};


Expand Down

16 changes: 12 additions & 4 deletions
transformer_engine/pytorch/csrc/extensions.cu

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand All@@ -39,9 +42,10 @@ void te_gemm(at::Tensor A,
 auto te_D = makeTransformerEngineTensor(D.data_ptr(),
 {static_cast<size_t>(D.size(0)),
 static_cast<size_t>(D.size(1))},
 D_type);
 D_type, D_amax.data_ptr(),
 D_scale.data_ptr(), nullptr);
 auto te_bias = makeTransformerEngineTensor(bias.data_ptr(), {static_cast<size_t>(bias.size(0))},
GetTransformerEngineDType(bias.scalar_type()));
bias_type);

 const auto gelu_shape = pre_gelu_out.data_ptr() == nullptr
 ? std::vector<size_t>{static_cast<size_t>(pre_gelu_out.size(0))}
Expand DownExpand Up@@ -869,10 +873,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 py::enum_<transformer_engine::FP8FwdTensors>(m, "FP8FwdTensors")
 .value("GEMM1_INPUT", transformer_engine::FP8FwdTensors::GEMM1_INPUT)
 .value("GEMM1_WEIGHT", transformer_engine::FP8FwdTensors::GEMM1_WEIGHT)
 .value("GEMM1_OUTPUT", transformer_engine::FP8FwdTensors::GEMM1_OUTPUT)
 .value("GEMM2_INPUT", transformer_engine::FP8FwdTensors::GEMM2_INPUT)
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT);
 .value("GEMM2_WEIGHT", transformer_engine::FP8FwdTensors::GEMM2_WEIGHT)
 .value("GEMM2_OUTPUT", transformer_engine::FP8FwdTensors::GEMM2_OUTPUT);

 py::enum_<transformer_engine::FP8BwdTensors>(m, "FP8BwdTensors")
 .value("GRAD_OUTPUT1", transformer_engine::FP8BwdTensors::GRAD_OUTPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2);
 .value("GRAD_INPUT1", transformer_engine::FP8BwdTensors::GRAD_INPUT1)
 .value("GRAD_OUTPUT2", transformer_engine::FP8BwdTensors::GRAD_OUTPUT2)
 .value("GRAD_INPUT2", transformer_engine::FP8BwdTensors::GRAD_INPUT2);
}

3 changes: 3 additions & 0 deletions
transformer_engine/pytorch/csrc/extensions.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -16,8 +16,11 @@ void te_gemm(at::Tensor A,
 transformer_engine::DType B_type,
 bool transb,
 at::Tensor D,
 at::Tensor D_scale,
 transformer_engine::DType D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 transformer_engine::DType bias_type,
 at::Tensor pre_gelu_out,
 bool grad,
 at::Tensor workspace,
Expand Down

7 changes: 7 additions & 0 deletions
transformer_engine/pytorch/csrc/ts_fp8_op.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -73,8 +73,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 int64_t B_type,
 int64_t transb,
 at::Tensor D,
 at::Tensor D_scale,
 int64_t D_type,
 at::Tensor D_amax,
 at::Tensor bias,
 int64_t bias_type,
 at::Tensor pre_gelu_out,
 int64_t grad,
 at::Tensor workspace,
Expand All@@ -87,6 +90,7 @@ at::Tensor te_gemm_ts(at::Tensor A,
 transformer_engine::DType B_type_arg = reverse_map_dtype(B_type);
 bool transb_arg = static_cast<bool>(transb);
 transformer_engine::DType D_type_arg = reverse_map_dtype(D_type);
 transformer_engine::DType bias_type_arg = reverse_map_dtype(bias_type);
 bool grad_arg = static_cast<bool>(grad);
 size_t workspaceSize_arg = static_cast<size_t>(workspaceSize);
 bool accumulate_arg = static_cast<bool>(accumulate);
Expand All@@ -109,8 +113,11 @@ at::Tensor te_gemm_ts(at::Tensor A,
 B_type_arg,
 transb_arg,
 D,
 D_scale,
 D_type_arg,
 D_amax,
 bias,
 bias_type_arg,
 pre_gelu_out,
 grad_arg,
 workspace,
Expand Down

3 changes: 2 additions & 1 deletion
transformer_engine/pytorch/fp8.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -299,7 +299,8 @@ def get_fp8_group() -> Union[dist_group_type, None]:

def update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
 """Update amax history and set next amax to zero."""
 amax_history = torch.roll(amax_history, -1, 0)
 if amax_history.shape[0] > 1:
 amax_history = torch.roll(amax_history, -1, 0)
 amax_history[0].fill_(0.0)
 return amax_history

Expand Down

4 changes: 3 additions & 1 deletion
transformer_engine/pytorch/module.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -158,8 +158,10 @@ def __init__(self) -> None:
 def set_meta_tensor(self, fwd: bool) -> None:
 """Init scales and amaxes for fwd | bwd."""
 fp8_meta_tensor_key = "scaling_fwd" if fwd else "scaling_bwd"
 # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and
 # 2 (grad_output and grad_input) for bwd
 num_fp8_tensors = (
 self.fp8_meta["num_gemms"] * 2 if fwd else self.fp8_meta["num_gemms"]
 self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2
Comment thread
vasunvidia marked this conversation as resolved.


Outdated

 )

 self.fp8_meta[fp8_meta_tensor_key] = tex.FP8TensorMeta()
Expand Down

5 changes: 4 additions & 1 deletion
transformer_engine/pytorch/te_onnx_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -99,7 +99,7 @@ def onnx_fp8_gelu(g, inputs, scale, amax, scale_inv, fp8_tensor, otype):

@symbolic_helper.parse_args("v", "fs", "i", "i", "i",
 "v", "fs", "i", "i", "i",
 "v", "i", "v", "v", "i",
 "v", "fs", "i", "fs", "v", "i", "v", "i",
 "v", "i", "i", "i")
def onnx_te_gemm(
 g,
Expand All@@ -114,8 +114,11 @@ def onnx_te_gemm(
 input_type,
 trans_input,
 out,
 out_scale,
 out_type,
 out_amax,
 bias,
 bias_type,
 pre_gelu_out,
 grad,
 workspace,
Expand Down