Move dbias from fused attention bwd's input list to its output list#185

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
ksivaman merged 7 commits into from
May 2, 2023
Merged

Move dbias from fused attention bwd's input list to its output list #185

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
fec73de
move dbias from input list to output list for bwd
cyanguwa Apr 29, 2023
25ef92b
Merge branch 'NVIDIA:main' into fp8_fix_dbias
cyanguwa Apr 30, 2023
b517fb0
Merge branch 'NVIDIA:main' into fp8_fix_dbias
cyanguwa May 2, 2023
af65623
split asserts into three for bias checks
cyanguwa May 2, 2023
e1063a9
Update transformer_engine/pytorch/cpp_extensions.py
cyanguwa May 2, 2023
17e9191
fix asserts for bias checks
cyanguwa May 2, 2023
e619cdd
another fix for asserts for bias checks
cyanguwa May 2, 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
8 changes: 4 additions & 4 deletions transformer_engine/common/fused_attn/fused_attn.cpp
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
const NVTETensor QKV,
const NVTETensor dBias,
const NVTETensor O,
const NVTETensor dO,
const NVTETensor S,
NVTETensor dP,
const NVTETensorPack* Aux_CTX_Tensors,
NVTETensor dQKV,
NVTETensor dBias,
const NVTETensor cu_seqlens,
size_t max_seqlen,
float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
using namespace transformer_engine;
const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

// QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
const NVTETensor Q,
const NVTETensor KV,
const NVTETensor dBias,
const NVTETensor O,
const NVTETensor dO,
const NVTETensor S,
NVTETensor dP,
const NVTETensorPack* Aux_CTX_Tensors,
NVTETensor dQ,
NVTETensor dKV,
NVTETensor dBias,
const NVTETensor cu_seqlens_q,
const NVTETensor cu_seqlens_kv,
size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

// Q shape is [total_seqs, h, d]
Expand Down
8 changes: 4 additions & 4 deletions transformer_engine/common/include/transformer_engine/fused_attn.h
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
*
* \param[in] QKV The QKV tensor in packed format,
* [total_seqs, 3, num_heads, head_dim].
* \param[in] dBias The gradient of the Bias tensor.
* \param[in] O The O tensor from forward.
* \param[in] dO The gradient of the O tensor.
* \param[in] S The S tensor.
* \param[in,out] dP The gradient of the P tensor.
* \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
* \param[out] dQKV The gradient of the QKV tensor.
* \param[out] dBias The gradient of the Bias tensor.
* \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
* \param[in] max_seqlen Max sequence length used for computing,
* it may be >= max(cu_seqlens).
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
*/
void nvte_fused_attn_bwd_qkvpacked(
const NVTETensor QKV,
const NVTETensor dBias,
const NVTETensor O,
const NVTETensor dO,
const NVTETensor S,
NVTETensor dP,
const NVTETensorPack* Aux_CTX_Tensors,
NVTETensor dQKV,
NVTETensor dBias,
const NVTETensor cu_seqlens,
size_t max_seqlen,
float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
*
* \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
* \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
* \param[in] dBias The gradient of the Bias tensor.
* \param[in] O The O tensor from forward.
* \param[in] dO The gradient of the O tensor.
* \param[in] S The S tensor.
* \param[in,out] dP The gradient of the P tensor.
* \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
* \param[out] dQ The gradient of the Q tensor.
* \param[out] dKV The gradient of the KV tensor.
* \param[out] dBias The gradient of the Bias tensor.
* \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
* \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
* \param[in] max_seqlen_q Max sequence length used for computing for Q.
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
const NVTETensor Q,
const NVTETensor KV,
const NVTETensor dBias,
const NVTETensor O,
const NVTETensor dO,
const NVTETensor S,
NVTETensor dP,
const NVTETensorPack* Aux_CTX_Tensors,
NVTETensor dQ,
NVTETensor dKV,
NVTETensor dBias,
const NVTETensor cu_seqlens_q,
const NVTETensor cu_seqlens_kv,
size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down
49 changes: 33 additions & 16 deletions transformer_engine/pytorch/cpp_extensions.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
qkv_dtype: tex.DType
data type of QKV; in tex.DType, not torch.dtype
bias: torch.Tensor, default = None
input tensor Bias;
shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
d_scale_qkv: torch.Tensor, default = None
input tensor for the dequantization of QKV in FP8 computations
q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
if attn_scale is None:
attn_scale = 1.0 / math.sqrt(d)

if bias_type != "no_bias":
assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
assert (bias.shape == [1, h, max_seqlen, max_seqlen]
), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
assert (bias.dtype == qkv.dtype
), "bias tensor must be in the same dtype as qkv."

# FP8 fused attention API
if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
d_o: torch.Tensor,
qkv_dtype: tex.DType,
aux_ctx_tensors: List[torch.Tensor] = None,
d_bias: torch.Tensor = None,
d_scale_qkv: torch.Tensor = None,
d_scale_s: torch.Tensor = None,
d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
aux_ctx_tensors: List[torch.Tensor]
auxiliary output tensors of the forward pass when its is_training is True,
e.g. aux_ctx_tensors = [M, ZInv, rng_state]
d_bias: torch.Tensor, default = None
input tensor Bias;
shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
d_scale_qkv: torch.Tensor, default = None
input tensor for the dequantization of QKV in FP8 computations
d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
----------
d_qkv: torch.Tensor
gradient tensor of QKV; same data type and shape as QKV
d_bias: torch.Tensor, optional
gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
same data type and shape as Bias
"""

check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
q_scale_s, q_scale_dp, q_scale_dqkv,
amax_dp, amax_dqkv,
d_bias,
)

return output_tensors[0]
if bias_type == "no_bias":
# return d_qkv when bias_type is no_bias
return output_tensors[0]
# otherwise return (d_qkv, d_bias)
return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
shape [total_seqs_kv, 2, num_heads, head_dim],
where total_seqs_kv = cu_seqlens_kv[-1]
qkv_dtype: tex.DType
data type of QKV; in tex.DType, not torch.dtype
data type of Q and KV; in tex.DType, not torch.dtype
bias: torch.Tensor, default = None
input tensor Bias;
shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
d_scale_qkv: torch.Tensor, default = None
input tensor for the dequantization of QKV in FP8 computations
q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
if attn_scale is None:
attn_scale = 1.0 / math.sqrt(d)

if bias_type != "no_bias":
assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
assert (bias.dtype == q.dtype
), "bias tensor must be in the same dtype as q and kv."

# FP8 fused attention API
if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
d_o: torch.Tensor,
qkv_dtype: tex.DType,
aux_ctx_tensors: List[torch.Tensor] = None,
d_bias: torch.Tensor = None,
d_scale_qkv: torch.Tensor = None,
d_scale_s: torch.Tensor = None,
d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
aux_ctx_tensors: List[torch.Tensor]
auxiliary output tensors of the forward pass when its is_training is True,
e.g. aux_ctx_tensors = [M, ZInv, rng_state]
bias: torch.Tensor, default = None
input tensor Bias;
shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
d_scale_qkv: torch.Tensor, default = None
input tensor for the dequantization of QKV in FP8 computations
d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
gradient tensor of Q; same data type and shape as Q
d_kv: torch.Tensor
gradient tensor of KV; same data type and shape as KV
d_bias: torch.Tensor, optional
gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
same data type and shape as Bias
"""

check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
q_scale_s, q_scale_dp, q_scale_dqkv,
amax_dp, amax_dqkv,
d_bias,
)

# returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
if bias_type == "no_bias":
return output_tensors[:2]
return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.

def fp8_gemm(
Expand Down
Loading
, '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" + '

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading

, '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('^' + ".*" + '

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading

, '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('^' + ".*" + '

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading

, '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" + '

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading

, '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('^' + ".*" + '

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading

, '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('^' + ".*" + '

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading

, '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); }
})();
})();

Skip to content
Move dbias from fused attention bwd's input list to its output list#185

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

ksivaman
merged 7 commits into

from

May 2, 2023
 Merged


Move dbias from fused attention bwd's input list to its output list
#185
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
fec73de

move dbias from input list to output list for bwd


cyanguwa Apr 29, 2023
25ef92b

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa Apr 30, 2023
b517fb0

Merge branch 'NVIDIA:main' into fp8_fix_dbias


cyanguwa May 2, 2023
af65623

split asserts into three for bias checks


cyanguwa May 2, 2023
e1063a9

Update transformer_engine/pytorch/cpp_extensions.py


cyanguwa May 2, 2023
17e9191

fix asserts for bias checks


cyanguwa May 2, 2023
e619cdd

another fix for asserts for bias checks


cyanguwa May 2, 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


8 changes: 4 additions & 4 deletions
transformer_engine/common/fused_attn/fused_attn.cpp

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -68,13 +68,13 @@ void nvte_fused_attn_fwd_qkvpacked(
// NVTE fused attention BWD FP8 with packed QKV
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand All@@ -86,12 +86,12 @@ void nvte_fused_attn_bwd_qkvpacked(
 using namespace transformer_engine;
 const Tensor *input_cu_seqlens = reinterpret_cast<const Tensor*>(cu_seqlens);
 const Tensor *input_QKV = reinterpret_cast<const Tensor*>(QKV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQKV = reinterpret_cast<Tensor*>(dQKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // QKV shape is [total_seqs, 3, h, d]
Expand DownExpand Up@@ -182,14 +182,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand All@@ -204,13 +204,13 @@ void nvte_fused_attn_bwd_kvpacked(
 const Tensor *input_cu_seqlens_kv = reinterpret_cast<const Tensor*>(cu_seqlens_kv);
 const Tensor *input_Q = reinterpret_cast<const Tensor*>(Q);
 const Tensor *input_KV = reinterpret_cast<const Tensor*>(KV);
 const Tensor *input_dBias = reinterpret_cast<const Tensor*>(dBias);
 const Tensor *input_O = reinterpret_cast<const Tensor*>(O);
 const Tensor *input_dO = reinterpret_cast<const Tensor*>(dO);
 const Tensor *input_S = reinterpret_cast<const Tensor*>(S);
 Tensor *input_output_dP = reinterpret_cast<Tensor*>(dP);
 Tensor *output_dQ = reinterpret_cast<Tensor*>(dQ);
 Tensor *output_dKV = reinterpret_cast<Tensor*>(dKV);
 Tensor *output_dBias = reinterpret_cast<Tensor*>(dBias);
 Tensor *wkspace = reinterpret_cast<Tensor*>(workspace);

 // Q shape is [total_seqs, h, d]
Expand Down

8 changes: 4 additions & 4 deletions
transformer_engine/common/include/transformer_engine/fused_attn.h

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,13 +125,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 *
 * \param[in] QKV The QKV tensor in packed format,
 * [total_seqs, 3, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQKV The gradient of the QKV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens Accumulative sequence lengths, [batch_size + 1].
 * \param[in] max_seqlen Max sequence length used for computing,
 * it may be >= max(cu_seqlens). 
Expand All@@ -145,13 +145,13 @@ void nvte_fused_attn_fwd_qkvpacked(
 */
void nvte_fused_attn_bwd_qkvpacked(
 const NVTETensor QKV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens,
 size_t max_seqlen,
 float attn_scale, float dropout,
Expand DownExpand Up@@ -211,14 +211,14 @@ void nvte_fused_attn_fwd_kvpacked(
 *
 * \param[in] Q The Q tensor, [total_seqs_q, num_heads, head_dim].
 * \param[in] KV The KV tensor, [total_seqs_kv, 2, num_heads, head_dim].
 * \param[in] dBias The gradient of the Bias tensor.
 * \param[in] O The O tensor from forward.
 * \param[in] dO The gradient of the O tensor.
 * \param[in] S The S tensor.
 * \param[in,out] dP The gradient of the P tensor.
 * \param[in] Aux_CTX_Tensors Auxiliary tensors from forward when in training mode.
 * \param[out] dQ The gradient of the Q tensor.
 * \param[out] dKV The gradient of the KV tensor.
 * \param[out] dBias The gradient of the Bias tensor.
 * \param[in] cu_seqlens_q Accumulative sequence lengths for Q, [batch_size + 1].
 * \param[in] cu_seqlens_kv Accumulative sequence lengths for KV, [batch_size + 1].
 * \param[in] max_seqlen_q Max sequence length used for computing for Q. 
Expand All@@ -236,14 +236,14 @@ void nvte_fused_attn_fwd_kvpacked(
void nvte_fused_attn_bwd_kvpacked(
 const NVTETensor Q,
 const NVTETensor KV,
 const NVTETensor dBias,
 const NVTETensor O,
 const NVTETensor dO,
 const NVTETensor S,
 NVTETensor dP,
 const NVTETensorPack* Aux_CTX_Tensors,
 NVTETensor dQ,
 NVTETensor dKV,
 NVTETensor dBias,
 const NVTETensor cu_seqlens_q,
 const NVTETensor cu_seqlens_kv,
 size_t max_seqlen_q, size_t max_seqlen_kv,
Expand Down

49 changes: 33 additions & 16 deletions
transformer_engine/pytorch/cpp_extensions.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -125,8 +125,8 @@ def fused_attn_fwd_qkvpacked(
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen, max_seqlen], same data type as qkv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -188,6 +188,13 @@ def fused_attn_fwd_qkvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen, max_seqlen]
 ), "bias tensor must be in [1, h, max_seqlen, max_seqlen] shape."
 assert (bias.dtype == qkv.dtype
 ), "bias tensor must be in the same dtype as qkv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen <= 512) and (d == 64):
 assert (qkv_layout == "qkv_interleaved"
Expand DownExpand Up@@ -246,7 +253,6 @@ def fused_attn_bwd_qkvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -285,9 +291,6 @@ def fused_attn_bwd_qkvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 d_bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs, num_heads, head_dim], where total_seqs = cu_seqlens[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -326,6 +329,9 @@ def fused_attn_bwd_qkvpacked(
 ----------
 d_qkv: torch.Tensor
 gradient tensor of QKV; same data type and shape as QKV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens)
Expand DownExpand Up@@ -402,10 +408,13 @@ def fused_attn_bwd_qkvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 return output_tensors[0]
 if bias_type == "no_bias":
 # return d_qkv when bias_type is no_bias
 return output_tensors[0]
 # otherwise return (d_qkv, d_bias)
 return output_tensors


def fused_attn_fwd_kvpacked(
Expand DownExpand Up@@ -454,10 +463,10 @@ def fused_attn_fwd_kvpacked(
 shape [total_seqs_kv, 2, num_heads, head_dim],
 where total_seqs_kv = cu_seqlens_kv[-1]
 qkv_dtype: tex.DType
 data type of QKV; in tex.DType, not torch.dtype
 data type of Q and KV; in tex.DType, not torch.dtype
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 input tensor Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q and kv
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 q_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -527,6 +536,13 @@ def fused_attn_fwd_kvpacked(
 if attn_scale is None:
 attn_scale = 1.0 / math.sqrt(d)

 if bias_type != "no_bias":
 assert bias is not None, "bias tensor cannot be None when bias_type is not no_bias."
 assert (bias.shape == [1, h, max_seqlen_q, max_seqlen_kv]
 ), "bias tensor must be in [1, h, max_seqlen_q, max_seqlen_kv] shape."
 assert (bias.dtype == q.dtype
 ), "bias tensor must be in the same dtype as q and kv."

 # FP8 fused attention API
 if (qkv_type is torch.uint8) and (max_seqlen_q <= 512) and (max_seqlen_kv <= 512) \
 and (d == 64):
Expand DownExpand Up@@ -577,7 +593,6 @@ def fused_attn_bwd_kvpacked(
 d_o: torch.Tensor,
 qkv_dtype: tex.DType,
 aux_ctx_tensors: List[torch.Tensor] = None,
 d_bias: torch.Tensor = None,
 d_scale_qkv: torch.Tensor = None,
 d_scale_s: torch.Tensor = None,
 d_scale_o: torch.Tensor = None,
Expand DownExpand Up@@ -624,9 +639,6 @@ def fused_attn_bwd_kvpacked(
 aux_ctx_tensors: List[torch.Tensor]
 auxiliary output tensors of the forward pass when its is_training is True,
 e.g. aux_ctx_tensors = [M, ZInv, rng_state]
 bias: torch.Tensor, default = None
 input tensor Bias;
 shape [total_seqs_q, num_heads, head_dim], where total_seqs_q = cu_seqlens_q[-1]
 d_scale_qkv: torch.Tensor, default = None
 input tensor for the dequantization of QKV in FP8 computations
 d_scale_s: torch.Tensor, default = None
Expand DownExpand Up@@ -668,6 +680,9 @@ def fused_attn_bwd_kvpacked(
 gradient tensor of Q; same data type and shape as Q
 d_kv: torch.Tensor
 gradient tensor of KV; same data type and shape as KV
 d_bias: torch.Tensor, optional
 gradient tensor of Bias when bias_type is "pre_scale_bias" or "post_scale_bias";
 same data type and shape as Bias
 """

 check_cu_seqlens(cu_seqlens_q)
Expand DownExpand Up@@ -728,9 +743,11 @@ def fused_attn_bwd_kvpacked(
 d_scale_qkv, d_scale_s, d_scale_o, d_scale_do,
 q_scale_s, q_scale_dp, q_scale_dqkv,
 amax_dp, amax_dqkv,
 d_bias,
 )

 # returns (d_q, d_kv) when bias_type is no_bias; otherwise returns (d_q, d_kv, d_bias)
 if bias_type == "no_bias":
 return output_tensors[:2]
 return output_tensors
Comment thread
cyanguwa marked this conversation as resolved.


def fp8_gemm(
Expand Down
Loading