Skip to content

GitLab

Unverified Commit 6e848924 authored by Michael Goldfarb's avatar Michael Goldfarb Committed by GitHub
Browse files

[JAX] Consolidate the distributed fused attention test code (#1405) 



Consolidate the distributed fused attention tests to shared input generation and execition logic.
Signed-off-by: default avatarMichael Goldfarb <mgoldfarb@nvidia.com>
parent c2937c5a
Hide whitespace changes
Inline Side-by-side
Showing with 323 additions and 392 deletions
tests/jax/distributed_test_base.py
View file @ 6e848924
...@@ -18,14 +18,22 @@ from utils import assert_allclose, is_devices_enough ...@@ -18,14 +18,22 @@ from utils import assert_allclose, is_devices_enough
def generate_configs(): def generate_configs():
configs = [] configs = []
if is_devices_enough(2): if is_devices_enough(2):
configs.append([2, (2,), "dp", MeshResource(dp_resource="dp")]) configs.append(
configs.append([2, (2,), "tp", MeshResource(tp_resource="tp")]) pytest.param(2, (2,), ("dp",), MeshResource(dp_resource="dp"), id="n2_dp2_tp1")
)
configs.append(
pytest.param(2, (2,), ("tp",), MeshResource(tp_resource="tp"), id="n2_dp1_tp2")
)
if is_devices_enough(4): if is_devices_enough(4):
TP_size = 2
DP_size = 2
configs.append( configs.append(
[4, (DP_size, TP_size), ("dp", "tp"), MeshResource(dp_resource="dp", tp_resource="tp")] pytest.param(
4,
(2, 2),
("dp", "tp"),
MeshResource(dp_resource="dp", tp_resource="tp"),
id=f"n4_dp2_tp2",
)
) )
return configs return configs
...@@ -33,7 +41,8 @@ def generate_configs(): ...@@ -33,7 +41,8 @@ def generate_configs():
def generate_context_parallel_configs(): def generate_context_parallel_configs():
configs = [] configs = []
mr = MeshResource(dp_resource="dp", cp_resource="cp", tp_resource="tp")
axes = ("dp", "cp", "tp")
DP_sizes = (1, 2) DP_sizes = (1, 2)
CP_sizes = (1, 2, 4, 8) CP_sizes = (1, 2, 4, 8)
TP_sizes = (1, 2) TP_sizes = (1, 2)
...@@ -41,13 +50,7 @@ def generate_context_parallel_configs(): ...@@ -41,13 +50,7 @@ def generate_context_parallel_configs():
ndev = cp * tp * dp ndev = cp * tp * dp
if is_devices_enough(ndev): if is_devices_enough(ndev):
configs.append( configs.append(
pytest.param( pytest.param(ndev, (dp, cp, tp), axes, mr, id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}")
ndev,
(dp, cp, tp),
("dp", "cp", "tp"),
MeshResource(dp_resource="dp", cp_resource="cp", tp_resource="tp"),
id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}",
)
) )
return configs return configs
......
tests/jax/test_distributed_fused_attn.py
View file @ 6e848924
...@@ -37,8 +37,7 @@ from transformer_engine.jax.attention import ( ...@@ -37,8 +37,7 @@ from transformer_engine.jax.attention import (
) )
from transformer_engine.jax.sharding import MeshResource from transformer_engine.jax.sharding import MeshResource
# We will use the golden reference model from our non distributed attention test fixture. from test_fused_attn import FusedAttnRunner, BiasShape, general_dot_product_attention, make_mask
from test_fused_attn import general_dot_product_attention, make_mask
DTYPES = [jnp.float16, jnp.bfloat16] DTYPES = [jnp.float16, jnp.bfloat16]
...@@ -49,7 +48,7 @@ class TestDistributedSelfAttn: ...@@ -49,7 +48,7 @@ class TestDistributedSelfAttn:
self, mesh_shape, mesh_axes, mesh_resource, with_bias, shape, dtype self, mesh_shape, mesh_axes, mesh_resource, with_bias, shape, dtype
): ):
jax_dtype = jax.dtypes.canonicalize_dtype(dtype) jax_dtype = jax.dtypes.canonicalize_dtype(dtype)
_, seqlen, _, heads, _ = shape _, seqlen, heads, _ = shape
is_dp_enabled = mesh_resource.dp_resource is not None is_dp_enabled = mesh_resource.dp_resource is not None
tp_size = 1 tp_size = 1
if mesh_resource.tp_resource is not None: if mesh_resource.tp_resource is not None:
...@@ -62,45 +61,28 @@ class TestDistributedSelfAttn: ...@@ -62,45 +61,28 @@ class TestDistributedSelfAttn:
# for loss and dbias # for loss and dbias
return generate_collectives_count(allreduce=allreduce_total_bytes, allgather=0, other=0) return generate_collectives_count(allreduce=allreduce_total_bytes, allgather=0, other=0)
def generate_inputs(self, shape, mesh_resource, with_bias, attn_mask_type, dtype):
batch, seqlen, _, heads, _ = shape
qkv = random.normal(random.PRNGKey(1124), shape, dtype=dtype)
bias = (
random.normal(random.PRNGKey(1125), (1, heads, seqlen, seqlen), dtype)
if with_bias
else None
)
mask = None
if attn_mask_type == AttnMaskType.PADDING_MASK:
mask = make_causal_mask(batch, seqlen)
elif attn_mask_type == AttnMaskType.CAUSAL_MASK:
mask = make_self_mask(batch, seqlen)
qkv_pspec = PartitionSpec(
mesh_resource.dp_resource, None, None, mesh_resource.tp_resource, None
)
bias_pspec = (
PartitionSpec(None, mesh_resource.tp_resource, None, None) if with_bias else None
)
mask_pspec = (
PartitionSpec(mesh_resource.dp_resource, None, None, None)
if attn_mask_type != AttnMaskType.NO_MASK
else None
)
return (qkv, bias, mask), (qkv_pspec, bias_pspec, mask_pspec)
@pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
@pytest.mark.parametrize("data_shape", [[32, 512, 3, 12, 64], [32, 1024, 3, 16, 128]])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"attn_bias_type", "data_shape",
[AttnBiasType.NO_BIAS, AttnBiasType.PRE_SCALE_BIAS, AttnBiasType.POST_SCALE_BIAS], [
pytest.param((32, 512, 12, 64), id="32-512-12-64"),
pytest.param((32, 1024, 16, 128), id="32-1024-16-128"),
],
) )
@pytest.mark.parametrize( @pytest.mark.parametrize(
"attn_mask_type", [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK] "attn_bias_type, bias_shape",
[
pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"),
pytest.param(AttnBiasType.PRE_SCALE_BIAS, BiasShape._1HSS, id="PRE_SCALE_BIAS-1HSS"),
pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape._1HSS, id="POST_SCALE_BIAS-1HSS"),
],
)
@pytest.mark.parametrize(
"attn_mask_type",
[
pytest.param(AttnMaskType.PADDING_MASK, id="PADDING_MASK"),
pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL_MASK"),
],
) )
@pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("dtype", DTYPES)
def test_self_attn( def test_self_attn(
...@@ -111,14 +93,14 @@ class TestDistributedSelfAttn: ...@@ -111,14 +93,14 @@ class TestDistributedSelfAttn:
mesh_resource, mesh_resource,
data_shape, data_shape,
attn_bias_type, attn_bias_type,
bias_shape,
attn_mask_type, attn_mask_type,
dtype, dtype,
): ):
dropout_prob = 0.0 dropout_prob = 0.0
is_training = True is_training = True
scaling_factor = 1.0
_, seqlen, _, num_head, hidden = data_shape batch, seqlen, num_head, hidden = data_shape
if not is_fused_attn_kernel_available( if not is_fused_attn_kernel_available(
dtype, dtype,
...@@ -136,74 +118,36 @@ class TestDistributedSelfAttn: ...@@ -136,74 +118,36 @@ class TestDistributedSelfAttn:
): ):
pytest.skip(f"No FusedAttn backend found") pytest.skip(f"No FusedAttn backend found")
def target_func(qkv, bias, mask): col_ref = self.generate_collectives_count_ref(
return jnp.mean( mesh_shape,
fused_attn( mesh_axes,
(qkv,), mesh_resource,
bias, attn_bias_type != AttnBiasType.NO_BIAS,
mask, data_shape,
None, dtype,
attn_bias_type=attn_bias_type,
attn_mask_type=attn_mask_type,
qkv_layout=QKVLayout.BS3HD,
scaling_factor=scaling_factor,
dropout_probability=dropout_prob,
is_training=is_training,
)
)
def ref_func(qkv, bias, mask):
query, key, value = jnp.split(qkv, [1, 2], axis=-3)
query = jnp.squeeze(query)
key = jnp.squeeze(key)
value = jnp.squeeze(value)
output = dot_product_attention(
query,
key,
value,
bias=bias,
mask=mask,
deterministic=is_training,
dropout_rate=dropout_prob,
dropout_rng=None,
dtype=jnp.float32,
)
return jnp.mean(output).astype(dtype)
with_bias = attn_bias_type != AttnBiasType.NO_BIAS
(qkv, bias, mask), (qkv_pspec, bias_pspec, mask_pspec) = self.generate_inputs(
data_shape, mesh_resource, with_bias, attn_mask_type, dtype
) )
collective_count_ref = self.generate_collectives_count_ref( runner = FusedAttnRunner(
mesh_shape, mesh_axes, mesh_resource, with_bias, data_shape, dtype batch,
seqlen,
seqlen,
num_head,
num_head,
hidden,
attn_bias_type,
attn_mask_type,
dropout_prob,
dtype,
is_training,
QKVLayout.BS3HD,
bias_shape,
None,
number_of_devices=device_count,
mesh_shape=mesh_shape,
mesh_axes=mesh_axes,
mesh_resource=mesh_resource,
coll_count_ref=col_ref,
) )
devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) runner.test_backward()
mesh = Mesh(devices, mesh_axes)
with mesh, fp8_autocast(mesh_resource=mesh_resource):
qkv_ = jax.device_put(qkv, NamedSharding(mesh, qkv_pspec))
bias_ = (
jax.device_put(bias, NamedSharding(mesh, bias_pspec)) if bias is not None else bias
)
mask_ = (
jax.device_put(mask, NamedSharding(mesh, mask_pspec)) if mask is not None else mask
)
grad_args = (0, 1) if with_bias else (0,)
out_grad_shardings = (qkv_pspec, bias_pspec) if with_bias else (qkv_pspec,)
compare_ops(
target_func,
ref_func,
[qkv_, bias_, mask_],
collective_count_ref,
grad_args=grad_args,
metric_fwd_dtype=dtype,
metric_bwd_dtype=dtype,
in_shardings=(qkv_pspec, bias_pspec, mask_pspec),
out_shardings=(None, out_grad_shardings),
)
class TestDistributedCrossAttn: class TestDistributedCrossAttn:
...@@ -213,31 +157,6 @@ class TestDistributedCrossAttn: ...@@ -213,31 +157,6 @@ class TestDistributedCrossAttn:
all_reduce_loss_bytes = 4 # 1 * FP32 all_reduce_loss_bytes = 4 # 1 * FP32
return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0) return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0)
def generate_inputs(self, shape, mesh_resource, attn_mask_type, dtype):
batch, seqlen, heads, hidden = shape
q = random.normal(random.PRNGKey(1124), shape, dtype=dtype)
kv = random.normal(random.PRNGKey(1125), (batch, seqlen, 2, heads, hidden), dtype=dtype)
mask = None
if attn_mask_type == AttnMaskType.PADDING_MASK:
mask = make_causal_mask(batch, seqlen)
elif attn_mask_type == AttnMaskType.CAUSAL_MASK:
mask = make_self_mask(batch, seqlen)
q_pspec = PartitionSpec(mesh_resource.dp_resource, None, mesh_resource.tp_resource, None)
kv_pspec = PartitionSpec(
mesh_resource.dp_resource, None, None, mesh_resource.tp_resource, None
)
mask_pspec = (
PartitionSpec(mesh_resource.dp_resource, None, None, None)
if attn_mask_type != AttnMaskType.NO_MASK
else None
)
return (q, kv, mask), (q_pspec, kv_pspec, mask_pspec)
@pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
@pytest.mark.parametrize("data_shape", [[32, 128, 12, 64], [32, 512, 16, 64]]) @pytest.mark.parametrize("data_shape", [[32, 128, 12, 64], [32, 512, 16, 64]])
@pytest.mark.parametrize( @pytest.mark.parametrize(
...@@ -248,11 +167,11 @@ class TestDistributedCrossAttn: ...@@ -248,11 +167,11 @@ class TestDistributedCrossAttn:
self, device_count, mesh_shape, mesh_axes, mesh_resource, data_shape, attn_mask_type, dtype self, device_count, mesh_shape, mesh_axes, mesh_resource, data_shape, attn_mask_type, dtype
): ):
attn_bias_type = AttnBiasType.NO_BIAS attn_bias_type = AttnBiasType.NO_BIAS
bias_shape = None
dropout_prob = 0.0 dropout_prob = 0.0
is_training = True is_training = True
scaling_factor = 1.0
_, seqlen, num_head, hidden = data_shape batch, seqlen, num_head, hidden = data_shape
if not is_fused_attn_kernel_available( if not is_fused_attn_kernel_available(
dtype, dtype,
...@@ -270,67 +189,29 @@ class TestDistributedCrossAttn: ...@@ -270,67 +189,29 @@ class TestDistributedCrossAttn:
): ):
pytest.skip(f"No FusedAttn backend found") pytest.skip(f"No FusedAttn backend found")
def target_func(q, kv, mask): col_ref = self.generate_collectives_count_ref()
return jnp.mean( runner = FusedAttnRunner(
fused_attn( batch,
(q, kv), seqlen,
None, seqlen,
mask, num_head,
None, num_head,
attn_bias_type=attn_bias_type, hidden,
attn_mask_type=attn_mask_type, attn_bias_type,
qkv_layout=QKVLayout.BSHD_BS2HD, attn_mask_type,
scaling_factor=scaling_factor, dropout_prob,
dropout_probability=dropout_prob, dtype,
is_training=is_training, is_training,
), QKVLayout.BSHD_BS2HD,
dtype=jnp.float32, bias_shape,
) None,
number_of_devices=device_count,
def ref_func(query, kv, mask): mesh_shape=mesh_shape,
key, value = jnp.split(kv, [1], axis=-3) mesh_axes=mesh_axes,
query = jnp.squeeze(query) mesh_resource=mesh_resource,
key = jnp.squeeze(key) coll_count_ref=col_ref,
value = jnp.squeeze(value)
output = dot_product_attention(
query,
key,
value,
bias=None,
mask=mask,
deterministic=is_training,
dropout_rate=dropout_prob,
dropout_rng=None,
dtype=jnp.float32,
)
return jnp.mean(output, dtype=jnp.float32)
(q, kv, mask), (q_pspec, kv_pspec, mask_pspec) = self.generate_inputs(
data_shape, mesh_resource, attn_mask_type, dtype
) )
collective_count_ref = self.generate_collectives_count_ref() runner.test_backward()
devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape)
mesh = Mesh(devices, mesh_axes)
with mesh, fp8_autocast(mesh_resource=mesh_resource):
q_ = jax.device_put(q, NamedSharding(mesh, q_pspec))
kv_ = jax.device_put(kv, NamedSharding(mesh, kv_pspec))
mask_ = (
jax.device_put(mask, NamedSharding(mesh, mask_pspec)) if mask is not None else mask
)
compare_ops(
target_func,
ref_func,
[q_, kv_, mask_],
collective_count_ref,
grad_args=(0, 1),
metric_fwd_dtype=dtype,
metric_bwd_dtype=dtype,
in_shardings=(q_pspec, kv_pspec, mask_pspec),
out_shardings=(None, (q_pspec, kv_pspec)),
)
@pytest.mark.parametrize( @pytest.mark.parametrize(
...@@ -366,41 +247,6 @@ class TestDistributedCrossAttn: ...@@ -366,41 +247,6 @@ class TestDistributedCrossAttn:
) )
class TestDistributedContextParallelSelfAttn: class TestDistributedContextParallelSelfAttn:
def generate_inputs(self, shape, kv_groups: int, attn_mask_type: AttnMaskType, dtype):
batch, seqlen, heads, hidden = shape
kv_shape = (batch, seqlen, heads // kv_groups, hidden)
qkey, kkey, vkey = random.split(random.PRNGKey(1124), 3)
q = random.normal(qkey, shape, dtype=dtype)
k = random.normal(kkey, (batch, seqlen, heads // kv_groups, hidden), dtype=dtype)
v = random.normal(vkey, (batch, seqlen, heads // kv_groups, hidden), dtype=dtype)
def gen_valid(bs, max_seqlen, pad_ratio):
pad_len = int(max_seqlen * pad_ratio)
valid_len = max_seqlen - pad_len
tokens = jnp.concatenate([jnp.ones((bs, valid_len)), jnp.zeros((bs, pad_len))], axis=-1)
return tokens, jnp.logical_not(tokens)
from test_fused_attn import make_mask
q_idx, _ = gen_valid(batch, seqlen, 0.0)
kv_idx, _ = gen_valid(batch, seqlen, 0.0)
mask = make_mask(q_idx, kv_idx, None, None, attn_mask_type)
return q, k, v, mask
def qkv_to_layout(self, q, k, v, qkv_layout):
qkv_args = ()
match qkv_layout:
case QKVLayout.BSHD_BS2HD:
k, v = map(partial(jnp.expand_dims, axis=-3), [k, v])
kv = jnp.concatenate((k, v), axis=-3)
qkv_args = (q, kv)
case QKVLayout.BSHD_BSHD_BSHD:
qkv_args = (q, k, v)
case _:
raise ValueError(f"Unsupported {qkv_layout=}")
return qkv_args
def impl_test_context_parallel_attn( def impl_test_context_parallel_attn(
self, self,
device_count, device_count,
...@@ -416,6 +262,7 @@ class TestDistributedContextParallelSelfAttn: ...@@ -416,6 +262,7 @@ class TestDistributedContextParallelSelfAttn:
cp_strategy, cp_strategy,
): ):
attn_bias_type = AttnBiasType.NO_BIAS attn_bias_type = AttnBiasType.NO_BIAS
bias_shape = None
dropout_prob = 0.0 dropout_prob = 0.0
is_training = True is_training = True
dp_size, cp_size, tp_size = mesh_shape dp_size, cp_size, tp_size = mesh_shape
...@@ -431,6 +278,29 @@ class TestDistributedContextParallelSelfAttn: ...@@ -431,6 +278,29 @@ class TestDistributedContextParallelSelfAttn:
num_kv_heads = num_head // kv_groups num_kv_heads = num_head // kv_groups
scaling_factor = 1.0 / np.sqrt(num_head) scaling_factor = 1.0 / np.sqrt(num_head)
runner = FusedAttnRunner(
batch,
seqlen,
seqlen,
num_head,
num_kv_heads,
hidden,
attn_bias_type,
attn_mask_type,
dropout_prob,
dtype,
is_training,
qkv_layout,
bias_shape,
None,
number_of_devices=device_count,
mesh_shape=mesh_shape,
mesh_axes=mesh_axes,
mesh_resource=mesh_resource,
cp_strategy=cp_strategy,
cp_load_balanced=load_balanced,
)
def check_has_backend_for_mask(mask_type): def check_has_backend_for_mask(mask_type):
return is_fused_attn_kernel_available( return is_fused_attn_kernel_available(
dtype, dtype,
...@@ -465,123 +335,7 @@ class TestDistributedContextParallelSelfAttn: ...@@ -465,123 +335,7 @@ class TestDistributedContextParallelSelfAttn:
if num_head % kv_groups != 0 or (num_head // kv_groups) % tp_size != 0: if num_head % kv_groups != 0 or (num_head // kv_groups) % tp_size != 0:
pytest.skip(f"Skipping {kv_groups=} not multiple of {data_shape=} or {tp_size=}") pytest.skip(f"Skipping {kv_groups=} not multiple of {data_shape=} or {tp_size=}")
def target_func(q, k, v, mask): runner.test_backward()
return fused_attn(
self.qkv_to_layout(q, k, v, qkv_layout),
None, # bias
mask,
None, # seed
attn_bias_type=attn_bias_type,
attn_mask_type=attn_mask_type,
qkv_layout=qkv_layout,
scaling_factor=scaling_factor,
dropout_probability=dropout_prob,
is_training=is_training,
context_parallel_strategy=cp_strategy,
context_parallel_causal_load_balanced=load_balanced,
context_parallel_axis="cp",
).astype(dtype)
def ref_func(q, k, v, mask):
output = general_dot_product_attention(
q,
k,
v,
bias=None,
mask=mask,
deterministic=not is_training,
scale_factor=scaling_factor,
dropout_rate=dropout_prob,
dropout_rng=None,
dtype=jnp.float32,
)
return output.astype(dtype)
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.is_causal():
gradient_multiplier /= 10
ret_valid = func(*args, **kwargs)
return (jnp.mean(ret_valid, dtype=jnp.float32) * gradient_multiplier).astype(dtype)
q, k, v, mask = self.generate_inputs(data_shape, kv_groups, attn_mask_type, dtype)
diff_argnums = (0, 1, 2)
# Single GPU (reference)
ref_func_jit = jax.jit(
jax.value_and_grad(
lambda q, k, v, mask: grad_func(ref_func, q, k, v, mask), argnums=diff_argnums
)
)
ref_fwd, ref_grads = ref_func_jit(q, k, v, mask)
# Multi GPU (function under test)
devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape)
mesh = Mesh(devices, mesh_axes)
with mesh, fp8_autocast(mesh_resource=mesh_resource, enabled=False):
qkv_ps = PartitionSpec(
mesh_resource.dp_resource,
mesh_resource.cp_resource,
mesh_resource.tp_resource,
None,
)
qkv_sharding = NamedSharding(mesh, qkv_ps)
mask_ps = PartitionSpec(
mesh_resource.dp_resource, None, mesh_resource.cp_resource, None
)
mask_sharding = NamedSharding(mesh, mask_ps)
reorder = partial(
reorder_causal_load_balancing, cp_size=cp_size, tensor_format=qkv_format
)
inverse_reorder = partial(
inverse_reorder_causal_load_balancing, cp_size=cp_size, tensor_format=qkv_format
)
if load_balanced:
q, k, v = jax.tree.map(reorder, (q, k, v))
q_, k_, v_ = map(partial(jax.device_put, device=qkv_sharding), [q, k, v])
mask_ = jax.device_put(mask, device=mask_sharding)
target_func_jit = jax.jit(
jax.value_and_grad(
lambda q, k, v, mask: grad_func(target_func, q, k, v, mask),
argnums=diff_argnums,
),
in_shardings=[qkv_sharding, qkv_sharding, qkv_sharding, mask_sharding],
out_shardings=(None, (qkv_sharding, qkv_sharding, qkv_sharding)),
)
target_fwd, target_grads = target_func_jit(q_, k_, v_, mask_)
if load_balanced:
target_dq, target_dk, target_dv = jax.tree.map(inverse_reorder, target_grads[0:3])
target_grads = (target_dq, target_dk, target_dv, *target_grads[3:])
has_diffs = False
print_debug_tensor_stats("target", target_fwd)
print_debug_tensor_stats("ref", ref_fwd)
print_debug_tensor_stats("diff", jnp.abs(target_fwd - ref_fwd))
assert_allclose(target_fwd, ref_fwd, dtype=dtype)
for i in range(len(target_grads)):
if ref_grads[i] is None or target_grads[i] is None:
# expect both none if one is
assert target_grads[i] is None and ref_grads[i] is None
else:
print_debug_tensor_stats(f"target_grad[{i}]", target_grads[i])
print_debug_tensor_stats(f"ref_grad[{i}]", ref_grads[i])
print_debug_tensor_stats(
f"diff_grad[{i}]", jnp.abs(target_grads[i] - ref_grads[i])
)
assert_allclose(target_grads[i], ref_grads[i], dtype=dtype)
def test_context_parallel_allgather_attn( def test_context_parallel_allgather_attn(
self, self,
......
tests/jax/test_fused_attn.py
View file @ 6e848924
...@@ -3,10 +3,10 @@ ...@@ -3,10 +3,10 @@
# See LICENSE for license information. # See LICENSE for license information.
"""Tests for fused attention""" """Tests for fused attention"""
from enum import Enum from enum import Enum
from dataclasses import dataclass from dataclasses import dataclass, field
from functools import partial from functools import partial
from math import sqrt from math import sqrt
from typing import Tuple, Optional from typing import Tuple, Optional, Dict
import random import random
import jax import jax
...@@ -19,16 +19,22 @@ from flax.linen import make_attention_mask ...@@ -19,16 +19,22 @@ from flax.linen import make_attention_mask
from flax.linen.dtypes import promote_dtype from flax.linen.dtypes import promote_dtype
from jax import Array from jax import Array
from jax import value_and_grad, jit from jax import value_and_grad, jit
from jax.sharding import Mesh, NamedSharding, PartitionSpec
from jax.typing import ArrayLike, DTypeLike from jax.typing import ArrayLike, DTypeLike
from transformer_engine.jax import fp8_autocast
from transformer_engine.jax.sharding import MeshResource
from transformer_engine.jax.attention import ( from transformer_engine.jax.attention import (
AttnBiasType, AttnBiasType,
AttnMaskType, AttnMaskType,
QKVLayout, QKVLayout,
QKVFormat, QKVFormat,
reorder_causal_load_balancing,
inverse_reorder_causal_load_balancing,
fused_attn, fused_attn,
fused_attn_thd, fused_attn_thd,
make_swa_mask, make_swa_mask,
CPStrategy,
) )
from transformer_engine.jax.cpp_extensions import FusedAttnHelper from transformer_engine.jax.cpp_extensions import FusedAttnHelper
from transformer_engine.transformer_engine_jax import ( from transformer_engine.transformer_engine_jax import (
...@@ -36,7 +42,8 @@ from transformer_engine.transformer_engine_jax import ( ...@@ -36,7 +42,8 @@ from transformer_engine.transformer_engine_jax import (
get_cudnn_version, get_cudnn_version,
) )
from utils import assert_allclose from distributed_test_base import assert_equal_collectives
from utils import assert_allclose, print_debug_tensor_stats
@pytest.fixture(autouse=True, scope="module") @pytest.fixture(autouse=True, scope="module")
...@@ -304,6 +311,19 @@ class FusedAttnRunner: ...@@ -304,6 +311,19 @@ class FusedAttnRunner:
bias_shape: BiasShape bias_shape: BiasShape
window_size: Optional[Tuple[int, int]] = None window_size: Optional[Tuple[int, int]] = None
# Specifies sharding resources for distributed tests
number_of_devices: int = 1
mesh_shape: tuple[int, ...] = (1, 1, 1)
mesh_axes: tuple[str, ...] = ("dp", "cp", "tp")
mesh_resource: MeshResource = field(default_factory=partial(MeshResource, "dp", "cp", "tp"))
# Context parallel aux arguments
cp_strategy: CPStrategy = CPStrategy.DEFAULT
cp_load_balanced: bool = True
# dictionary of expected collective comm bytes
coll_count_ref: Optional[Dict[str, int]] = None
# See https://docs.nvidia.com/deeplearning/cudnn/latest/release-notes.html#cudnn-9-4-0 for known issue # See https://docs.nvidia.com/deeplearning/cudnn/latest/release-notes.html#cudnn-9-4-0 for known issue
# generating zero-length ragged tensors. This setting adjusts the test to avoid the zero-length cases. # generating zero-length ragged tensors. This setting adjusts the test to avoid the zero-length cases.
def _get_max_segments_per_sequence(self): def _get_max_segments_per_sequence(self):
...@@ -362,6 +382,14 @@ class FusedAttnRunner: ...@@ -362,6 +382,14 @@ class FusedAttnRunner:
def _setup_inputs(self): def _setup_inputs(self):
self._check_configs() self._check_configs()
# Create a mesh for distributed tests
self.devices = np.asarray(jax.devices()[: self.number_of_devices]).reshape(*self.mesh_shape)
self.mesh = Mesh(self.devices, self.mesh_axes)
self.dp_size = self.mesh.shape.get(self.mesh_resource.dp_resource, 1)
self.cp_size = self.mesh.shape.get(self.mesh_resource.cp_resource, 1)
self.tp_size = self.mesh.shape.get(self.mesh_resource.tp_resource, 1)
key = jax.random.PRNGKey(0) key = jax.random.PRNGKey(0)
q_key, k_key, v_key, bias_key, dropout_key = jax.random.split(key, 5) q_key, k_key, v_key, bias_key, dropout_key = jax.random.split(key, 5)
...@@ -527,6 +555,66 @@ class FusedAttnRunner: ...@@ -527,6 +555,66 @@ class FusedAttnRunner:
self.dropout_rng = dropout_key if self.dropout_prob > 0 else None self.dropout_rng = dropout_key if self.dropout_prob > 0 else None
self.scaling_factor = 1.0 / sqrt(self.head_dim) self.scaling_factor = 1.0 / sqrt(self.head_dim)
# Setup distributed sharding specs
# Setup shardings for distributed tests
self.qkvo_psec = PartitionSpec(
self.mesh_resource.dp_resource,
self.mesh_resource.cp_resource,
self.mesh_resource.tp_resource,
None,
)
self.qkvo_sharding = NamedSharding(self.mesh, self.qkvo_psec)
self.mask_pspec = PartitionSpec(
self.mesh_resource.dp_resource, None, self.mesh_resource.cp_resource, None
)
self.mask_sharding = NamedSharding(self.mesh, self.mask_pspec)
if self.bias_shape == BiasShape._1HSS:
self.bias_pspec = PartitionSpec(
None, self.mesh_resource.tp_resource, self.mesh_resource.cp_resource, None
)
elif self.bias_shape == BiasShape._B1SS:
self.bias_pspec = PartitionSpec(
self.mesh_resource.dp_resource, None, self.mesh_resource.cp_resource, None
)
elif self.bias_shape == BiasShape._11SS:
self.bias_pspec = PartitionSpec(None, None, self.mesh_resource.cp_resource, None)
else:
self.bias_pspec = PartitionSpec()
self.bias_sharding = NamedSharding(self.mesh, self.bias_pspec)
self.dropout_rng_pspec = PartitionSpec(
None,
)
self.dropout_rng_sharding = NamedSharding(self.mesh, self.dropout_rng_pspec)
self.logit_scale_pspec = PartitionSpec(None, None, self.mesh_resource.cp_resource, None)
self.logit_scale_sharding = NamedSharding(self.mesh, self.logit_scale_pspec)
# [batch][max_segments_per_batch]
# TODO(mgoldfarb-nvidia): Will need to handle CP cases of replicated or distributed length/offset.
self.seq_length_offset_pspec = PartitionSpec(self.mesh_resource.dp_resource, None)
self.seq_length_offset_sharding = NamedSharding(self.mesh, self.seq_length_offset_pspec)
# Softmax aux sharding
if self.cp_size > 1 and self.cp_load_balanced:
self.cp_reorder_fn = partial(
reorder_causal_load_balancing,
cp_size=self.cp_size,
tensor_format=self.qkv_layout.get_qkv_format(),
)
self.cp_inverse_reorder_fn = partial(
inverse_reorder_causal_load_balancing,
cp_size=self.cp_size,
tensor_format=self.qkv_layout.get_qkv_format(),
)
else:
# no-ops for non cp or non load balanced
self.cp_reorder_fn = lambda x: x
self.cp_inverse_reorder_fn = lambda x: x
def test_forward(self): def test_forward(self):
""" """
Test forward without JIT Test forward without JIT
...@@ -534,17 +622,21 @@ class FusedAttnRunner: ...@@ -534,17 +622,21 @@ class FusedAttnRunner:
self._setup_inputs() self._setup_inputs()
args = [self.q, self.k, self.v, self.bias, self.mask, self.dropout_rng] args = [self.q, self.k, self.v, self.bias, self.mask, self.dropout_rng]
customcall_args = [ customcall_args = [
self.q, # Put test data onto each GPU for distributed.
self.k, # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and
self.v, # THD params once we support those features on CP.
self.bias, jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding),
self.mask_for_customcall, jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
self.seqlens_q, jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding),
self.seqlens_kv, jax.device_put(self.bias, self.bias_sharding),
self.offsets_q, jax.device_put(self.mask_for_customcall, self.mask_sharding),
self.offsets_kv, jax.device_put(self.seqlens_q, self.seq_length_offset_sharding),
self.dropout_rng, jax.device_put(self.seqlens_kv, self.seq_length_offset_sharding),
jax.device_put(self.offsets_q, self.seq_length_offset_sharding),
jax.device_put(self.offsets_kv, self.seq_length_offset_sharding),
jax.device_put(self.dropout_rng, self.dropout_rng_sharding),
] ]
kwargs = { kwargs = {
"attn_bias_type": self.attn_bias_type, "attn_bias_type": self.attn_bias_type,
...@@ -555,10 +647,31 @@ class FusedAttnRunner: ...@@ -555,10 +647,31 @@ class FusedAttnRunner:
"qkv_layout": self.qkv_layout, "qkv_layout": self.qkv_layout,
"max_segments_per_seq": self._get_max_segments_per_sequence(), "max_segments_per_seq": self._get_max_segments_per_sequence(),
"window_size": self.window_size, "window_size": self.window_size,
"context_parallel_strategy": self.cp_strategy,
"context_parallel_causal_load_balanced": self.cp_load_balanced,
} }
# Convert the outputs to float32 for the elementwise comparison customcall_fused_dpa_jit = jit(
primitive_out = customcall_fused_dpa(*customcall_args, **kwargs) partial(customcall_fused_dpa, **kwargs),
static_argnames=kwargs.keys(),
in_shardings=[
self.qkvo_sharding,
self.qkvo_sharding,
self.qkvo_sharding,
self.bias_sharding,
self.mask_sharding,
self.seq_length_offset_sharding,
self.seq_length_offset_sharding,
self.seq_length_offset_sharding,
self.seq_length_offset_sharding,
self.dropout_rng_sharding,
],
)
with self.mesh, fp8_autocast(mesh_resource=self.mesh_resource):
primitive_out = customcall_fused_dpa_jit(*customcall_args)
primitive_out = self.cp_inverse_reorder_fn(primitive_out)
reference_out = jax_dpa(*args, **kwargs) reference_out = jax_dpa(*args, **kwargs)
if self.is_training and self.dropout_prob > 0.0: if self.is_training and self.dropout_prob > 0.0:
...@@ -571,9 +684,19 @@ class FusedAttnRunner: ...@@ -571,9 +684,19 @@ class FusedAttnRunner:
assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype) assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype)
assert_allclose(primitive_valid, reference_valid, dtype=self.dtype) assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
if self.coll_count_ref is not None:
with self.mesh, fp8_autocast(mesh_resource=self.mesh_resource):
target_hlo = (
customcall_fused_dpa_jit.lower(*customcall_args, **kwargs).compile().as_text()
)
assert_equal_collectives(target_hlo, self.coll_count_ref)
def test_backward(self): def test_backward(self):
""" """
Test value_and_grad with JIT, which includes both forward and backward Test value_and_grad with JIT, which includes both forward and backward.
If coll_count_ref is not None then the HLO of the backwrds function
HLO will be examined for the expected comms.
""" """
self._setup_inputs() self._setup_inputs()
...@@ -587,20 +710,24 @@ class FusedAttnRunner: ...@@ -587,20 +710,24 @@ class FusedAttnRunner:
ret_valid = jnp.where( ret_valid = jnp.where(
self.pad_q[..., jnp.newaxis, jnp.newaxis], 0, func(*args, **kwargs) self.pad_q[..., jnp.newaxis, jnp.newaxis], 0, func(*args, **kwargs)
) )
return (jnp.mean(ret_valid, dtype=jnp.float32) * gradient_multiplier).astype(self.dtype) return (
jnp.mean(ret_valid.astype(jnp.float32), dtype=jnp.float32) * gradient_multiplier
).astype(self.dtype)
args = [self.q, self.k, self.v, self.bias, self.mask, self.dropout_rng] args = [self.q, self.k, self.v, self.bias, self.mask, self.dropout_rng]
customcall_args = [ customcall_args = [
self.q, # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and
self.k, # THD params once we support those features on CP.
self.v, jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding),
self.bias, jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
self.mask_for_customcall, jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding),
self.seqlens_q, jax.device_put(self.bias, self.bias_sharding),
self.seqlens_kv, jax.device_put(self.mask_for_customcall, self.mask_sharding),
self.offsets_q, jax.device_put(self.seqlens_q, self.seq_length_offset_sharding),
self.offsets_kv, jax.device_put(self.seqlens_kv, self.seq_length_offset_sharding),
self.dropout_rng, jax.device_put(self.offsets_q, self.seq_length_offset_sharding),
jax.device_put(self.offsets_kv, self.seq_length_offset_sharding),
jax.device_put(self.dropout_rng, self.dropout_rng_sharding),
] ]
kwargs = { kwargs = {
"attn_bias_type": self.attn_bias_type, "attn_bias_type": self.attn_bias_type,
...@@ -611,10 +738,22 @@ class FusedAttnRunner: ...@@ -611,10 +738,22 @@ class FusedAttnRunner:
"qkv_layout": self.qkv_layout, "qkv_layout": self.qkv_layout,
"max_segments_per_seq": self._get_max_segments_per_sequence(), "max_segments_per_seq": self._get_max_segments_per_sequence(),
"window_size": self.window_size, "window_size": self.window_size,
"context_parallel_strategy": self.cp_strategy,
"context_parallel_causal_load_balanced": self.cp_load_balanced,
} }
# We can compute dBias only for the [1, h, s, s] layout # We can compute dBias only for the [1, h, s, s] layout
arg_nums = (0, 1, 2, 3) if self.bias_shape == BiasShape._1HSS else (0, 1, 2) if self.bias_shape == BiasShape._1HSS:
arg_nums = (0, 1, 2, 3)
grad_shardings = (
self.qkvo_sharding,
self.qkvo_sharding,
self.qkvo_sharding,
self.bias_sharding,
)
else:
arg_nums = (0, 1, 2)
grad_shardings = (self.qkvo_sharding, self.qkvo_sharding, self.qkvo_sharding)
# Use FP16/BF16 to sum the results may cause overflow, use FP32 for the summation # Use FP16/BF16 to sum the results may cause overflow, use FP32 for the summation
jitted_primitive = jit( jitted_primitive = jit(
...@@ -623,7 +762,20 @@ class FusedAttnRunner: ...@@ -623,7 +762,20 @@ class FusedAttnRunner:
customcall_fused_dpa, q, k, v, bias, *args, **kwargs customcall_fused_dpa, q, k, v, bias, *args, **kwargs
), ),
arg_nums, arg_nums,
) ),
in_shardings=(
self.qkvo_sharding,
self.qkvo_sharding,
self.qkvo_sharding,
self.bias_sharding,
self.mask_sharding,
self.seq_length_offset_sharding,
self.seq_length_offset_sharding,
self.seq_length_offset_sharding,
self.seq_length_offset_sharding,
self.dropout_rng_sharding,
),
out_shardings=(None, grad_shardings),
) )
jitted_reference = jit( jitted_reference = jit(
value_and_grad( value_and_grad(
...@@ -632,20 +784,31 @@ class FusedAttnRunner: ...@@ -632,20 +784,31 @@ class FusedAttnRunner:
) )
) )
primitive_out, primitive_dgrad = jitted_primitive(*customcall_args) with self.mesh, fp8_autocast(mesh_resource=self.mesh_resource):
primitive_out, primitive_dgrad = jitted_primitive(*customcall_args)
reference_out, reference_dgrad = jitted_reference(*args) reference_out, reference_dgrad = jitted_reference(*args)
# Skip elementwise comparison when dropout enabled # Skip elementwise comparison when dropout enabled
if self.dropout_prob > 0.0: if self.dropout_prob > 0.0:
return return
print_debug_tensor_stats(f"primitive_out", primitive_out)
print_debug_tensor_stats(f"reference_grad_valid", reference_out)
print_debug_tensor_stats(f"diff_grad", jnp.abs(primitive_out - reference_out))
assert_allclose(primitive_out, reference_out, dtype=self.dtype) assert_allclose(primitive_out, reference_out, dtype=self.dtype)
def check_dqkv(primitive, reference, pad): def check_dqkv(primitive, reference, pad, idx):
primitive_valid, primitive_invalid, reference_valid, reference_invalid = ( primitive_valid, primitive_invalid, reference_valid, reference_invalid = (
_split_valid_and_invalid(primitive, reference, pad) _split_valid_and_invalid(primitive, reference, pad)
) )
print_debug_tensor_stats(f"primitive_grad_valid[{idx}]", primitive_valid[idx])
print_debug_tensor_stats(f"reference_grad_valid[{idx}]", reference_valid[idx])
print_debug_tensor_stats(
f"diff_grad[{idx}]", jnp.abs(primitive_valid[idx] - reference_valid[idx])
)
assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype) assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype)
assert_allclose(primitive_invalid, reference_invalid, dtype=self.dtype) assert_allclose(primitive_invalid, reference_invalid, dtype=self.dtype)
assert_allclose(primitive_valid, reference_valid, dtype=self.dtype) assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
...@@ -653,11 +816,17 @@ class FusedAttnRunner: ...@@ -653,11 +816,17 @@ class FusedAttnRunner:
primitive_dq, primitive_dk, primitive_dv = primitive_dgrad[:3] primitive_dq, primitive_dk, primitive_dv = primitive_dgrad[:3]
reference_dq, reference_dk, reference_dv = reference_dgrad[:3] reference_dq, reference_dk, reference_dv = reference_dgrad[:3]
check_dqkv(primitive_dq, reference_dq, self.pad_q) primitive_dq = self.cp_inverse_reorder_fn(primitive_dq)
check_dqkv(primitive_dk, reference_dk, self.pad_kv) primitive_dk = self.cp_inverse_reorder_fn(primitive_dk)
check_dqkv(primitive_dv, reference_dv, self.pad_kv) primitive_dv = self.cp_inverse_reorder_fn(primitive_dv)
check_dqkv(primitive_dq, reference_dq, self.pad_q, 0)
check_dqkv(primitive_dk, reference_dk, self.pad_kv, 1)
check_dqkv(primitive_dv, reference_dv, self.pad_kv, 2)
if self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape == BiasShape._1HSS: if self.attn_bias_type != AttnBiasType.NO_BIAS and self.bias_shape == BiasShape._1HSS:
# TODO(mgoldfarb-nvidia): Inverse reorder bias once supported by a CP implementation.
primitive_dbias = primitive_dgrad[3] primitive_dbias = primitive_dgrad[3]
reference_dbias = reference_dgrad[3] reference_dbias = reference_dgrad[3]
...@@ -685,6 +854,11 @@ class FusedAttnRunner: ...@@ -685,6 +854,11 @@ class FusedAttnRunner:
dtype=self.dtype, dtype=self.dtype,
) )
if self.coll_count_ref is not None:
with self.mesh, fp8_autocast(mesh_resource=self.mesh_resource):
target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text()
assert_equal_collectives(target_hlo, self.coll_count_ref)
@pytest.mark.parametrize( @pytest.mark.parametrize(
"attn_mask_type", "attn_mask_type",
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment