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

[JAX] Fused attention unit tests fixes and refinements #1352

Merged
merged 10 commits into from
Dec 17, 2024
Merged

[JAX] Fused attention unit tests fixes and refinements #1352

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
bd8b8e9
Add util functions to attn_mask_type
zlsh80826 Nov 11, 2024
2978fc7
Add util functions to qkv_layout
zlsh80826 Nov 19, 2024
80efad1
Fix THD cross reference code
zlsh80826 Dec 2, 2024
50b841e
Remove explicit segment_pad, encoding it to segment_ids
zlsh80826 Dec 2, 2024
c16f572
Add jax.jit, replace _token with segment_ids, rename bias shape enum
zlsh80826 Dec 2, 2024
93f21fc
Add comment for make_mask
zlsh80826 Dec 2, 2024
f6456e7
Clean code
zlsh80826 Dec 2, 2024
011517e
Add doc strings for the added functions
zlsh80826 Dec 2, 2024
921fbb9
Remove cache for fa deterministic which causes UT failed
zlsh80826 Dec 5, 2024
74f0603
Rename fixture to avoid conflict
zlsh80826 Dec 12, 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
2 changes: 1 addition & 1 deletion tests/jax/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ def clear_live_arrays():


@pytest.fixture(autouse=True, scope="module")
def enable_fused_attn():
def enable_fused_attn_after_hopper():
"""
Enable fused attn for hopper+ arch.
Fused attn kernels on pre-hopper arch are not deterministic.
Expand Down
6 changes: 2 additions & 4 deletions tests/jax/test_distributed_fused_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from utils import (
make_causal_mask,
make_self_mask,
assert_tree_like_allclose,
assert_allclose,
print_debug_tensor_stats,
)
Expand All @@ -32,7 +31,6 @@
AttnMaskType,
QKVLayout,
QKVFormat,
get_qkv_format,
reorder_causal_load_balancing,
inverse_reorder_causal_load_balancing,
CPStrategy,
Expand Down Expand Up @@ -421,7 +419,7 @@ def impl_test_contex_parallel_attn(
dropout_prob = 0.0
is_training = True
dp_size, cp_size, tp_size = mesh_shape
qkv_format = get_qkv_format(qkv_layout)
qkv_format = qkv_layout.get_qkv_format()

batch, seqlen, num_head, hidden = data_shape

Expand Down Expand Up @@ -503,7 +501,7 @@ def grad_func(func, *args, **kwargs):
# Gradient is small, use a gradient multiplier to amplify the gradient
_, max_seq_len, num_heads, _ = data_shape
gradient_multiplier = max_seq_len * num_heads
if attn_mask_type in [AttnMaskType.CAUSAL_MASK, AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK]:
if attn_mask_type.is_causal():
gradient_multiplier /= 10
ret_valid = func(*args, **kwargs)
return (jnp.mean(ret_valid, dtype=jnp.float32) * gradient_multiplier).astype(dtype)
Expand Down
Loading
Loading