Skip to content
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

Jump to bottom

Fix QKV dtype in the bwd of FP8+CP #1134

Merged
merged 9 commits into from
Aug 30, 2024
Merged

Fix QKV dtype in the bwd of FP8+CP #1134

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
7c926f5
fix qkv_dtype of FP8+CP
xrennvidia Aug 23, 2024
80c4af5
Merge branch 'main' into xren/cp_fp8_fix
xrennvidia Aug 26, 2024
184ad60
config cp correction dtype of FP8+CP
xrennvidia Aug 27, 2024
e383bbe
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 27, 2024
62e620a
code style change
xrennvidia Aug 27, 2024
bda0506
Merge branch 'main' into xren/cp_fp8_fix
xrennvidia Aug 27, 2024
24c626d
always do FP8 CP correction in FP32
xrennvidia Aug 29, 2024
5b5d8d7
Merge branch 'main' into xren/cp_fp8_fix
xrennvidia Aug 29, 2024
8e34de1
Merge branch 'main' into xren/cp_fp8_fix
cyanguwa Aug 30, 2024
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
5 changes: 3 additions & 2 deletions transformer_engine/pytorch/attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2261,8 +2261,9 @@ def backward(ctx, dout):

if ctx.fp8:
if ctx.use_fused_attention:
fp8_dtype_forward = get_fp8_te_dtype(ctx.fp8_meta["recipe"], fprop_tensor=True)
fp8_dtype_backward = get_fp8_te_dtype(ctx.fp8_meta["recipe"], fprop_tensor=False)
fused_attn_qkv_dtype = fp8_dtype_backward
fused_attn_qkv_dtype = fp8_dtype_forward
fused_attn_dqkv_dtype = fp8_dtype_backward
fused_attn_backend = FusedAttnBackend["FP8"]
dq_fp8 = torch.empty((cp_size, *q.shape), dtype=q.dtype, device=q.device)
Expand Down Expand Up @@ -2304,7 +2305,7 @@ def backward(ctx, dout):
if ctx.use_fused_attention:
fp8_meta_kwargs = {}
fused_attn_qkv_dtype = TE_DType[q.dtype]
fused_attn_dqkv_dtype = TE_DType[q.dtype]
fused_attn_dqkv_dtype = TE_DType[dout.dtype]
fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"]

out = out.view(*q.shape)
Expand Down
Loading