Fixed the ring attention backward to not do 2x work. Although this means it doesn't quite overlap the communications.

Author: nshepperd

setup.py

@@ -106,7 +106,7 @@ def append_nvcc_threads(nvcc_extra_args):
     #         cc_flag.append("arch=compute_90,code=sm_90")
     ext_modules.append(
         CUDAExtension(
-            name="flash_attn_jax.flash_api",
+            name="flash_attn_jax_lib.flash_api",
             sources=[
                 "csrc/flash_attn/flash_api.cpp",
                 "csrc/flash_attn/flash_common.cpp",

src/flash_attn_jax/flash.py

@@ -39,7 +39,7 @@
 _flash_mha_bwd_p.def_impl(partial(xla.apply_primitive, _flash_mha_bwd_p))
 
 try:
-    # JAX 0.4.24 and above requires this.
+    # JAX 0.4.24 and above requires this because of custom partitioning.
     import jax._src.dispatch
     jax._src.dispatch.prim_requires_devices_during_lowering.add(_flash_mha_bwd_p)
     jax._src.dispatch.prim_requires_devices_during_lowering.add(_flash_mha_fwd_p)

src/flash_attn_jax/flash_hlo.py

@@ -21,7 +21,7 @@
 from einops import rearrange
 import math
 
-import flash_attn_jax.flash_api as flash_api
+import flash_attn_jax_lib.flash_api as flash_api
 
 # ==== Register primitives ====
 

src/flash_attn_jax/flash_sharding.py

@@ -22,6 +22,7 @@
 import math
 
 from .flash_hlo import _flash_mha_fwd_hlo, _flash_mha_bwd_hlo
+from .ring_attention import ring_fwd, ring_bwd
 
 # ==== Sharding ====
 
@@ -30,12 +31,20 @@
 
 from jax._src.ad_checkpoint import _optimization_barrier
 
+def is_replicated(sharding):
+    return (isinstance(sharding, PositionalSharding) and sharding.shape == (1,)) or (isinstance(sharding, NamedSharding) and len(sharding.spec) == 0)
+
 def partition_fwd(softmax_scale, is_causal, window_size, mesh, arg_shapes, result_shape):
     result_shardings = jax.tree_map(lambda x: x.sharding, result_shape)
     arg_shardings = jax.tree_map(lambda x: x.sharding, arg_shapes)
 
     q_sharding = arg_shardings[0]
-    if isinstance(q_sharding, PositionalSharding):
+    k_sharding = arg_shardings[1]
+    v_sharding = arg_shardings[2]
+    assert q_sharding == k_sharding and q_sharding == v_sharding, "Only support q, k, v sharing the same sharding."
+    if is_replicated(q_sharding):
+        result_sharding = (q_sharding, q_sharding)
+    elif isinstance(q_sharding, PositionalSharding):
         (n,l,h,d) = q_sharding.shape
         assert d == 1, "Sharding across `d` won't be efficient, so it's not supported."
         assert l == 1, "For ring attention, use `with Mesh(...) as mesh` and NamedSharding."
@@ -53,7 +62,7 @@ def partition_fwd(softmax_scale, is_causal, window_size, mesh, arg_shapes, resul
             axis_name = l
             axis_size = mesh.shape[axis_name]
             # ring attention
-            return mesh, partial(ring_fwd, softmax_scale, is_causal, axis_name, axis_size), result_shardings, arg_shardings
+            return mesh, partial(ring_fwd, softmax_scale=softmax_scale, is_causal=is_causal, axis_name=axis_name, axis_size=axis_size, mha_fwd=_flash_mha_fwd_hlo), result_shardings, arg_shardings
         else:
             result_shardings = q_sharding, NamedSharding(mesh, P(n,h,l))
             arg_shardings = q_sharding, q_sharding, q_sharding
@@ -64,7 +73,12 @@ def fwd(q,k,v):
 def infer_sharding_fwd(softmax_scale, is_causal, window_size, mesh, arg_shapes, result_shape):
     arg_shardings = jax.tree_map(lambda x: x.sharding, arg_shapes)
     q_sharding = arg_shardings[0]
-    if isinstance(q_sharding, PositionalSharding):
+    k_sharding = arg_shardings[1]
+    v_sharding = arg_shardings[2]
+    assert q_sharding == k_sharding and q_sharding == v_sharding, "Only support q, k, v sharing the same sharding."
+    if is_replicated(q_sharding):
+        result_sharding = (q_sharding, q_sharding)
+    elif isinstance(q_sharding, PositionalSharding):
         [n,l,h,d] = q_sharding.shape
         result_sharding = (q_sharding, # [n,l,h,d]
                            q_sharding.replicate(3).reshape(n,l,h).transpose((0,2,1)) # [n,h,l]
@@ -73,6 +87,8 @@ def infer_sharding_fwd(softmax_scale, is_causal, window_size, mesh, arg_shapes,
         [n,l,h,d] = q_sharding.spec
         result_sharding = (q_sharding,
                            NamedSharding(q_sharding.mesh, P(n,h,l)))
+    else:
+        raise ValueError("Unsupported sharding type.", type(q_sharding))
     return result_sharding
 
 _flash_mha_fwd_hlo_sharded.def_partition(
@@ -99,7 +115,10 @@ def partition_bwd(softmax_scale, is_causal, window_size, mesh, arg_shapes, resul
     v_sharding = arg_shardings[3]
     o_sharding = arg_shardings[4]
     lse_sharding = arg_shardings[5]
-    if isinstance(q_sharding, PositionalSharding):
+    assert q_sharding == k_sharding and q_sharding == v_sharding, "Only support q, k, v sharing the same sharding."
+    if is_replicated(q_sharding):
+        result_shardings = (q_sharding,)*3
+    elif isinstance(q_sharding, PositionalSharding):
         assert q_sharding == k_sharding, "Expect q and k sharding to match"
         assert q_sharding == v_sharding, "Expect q and v sharding to match"
         [n, l, h, d] = q_sharding.shape
@@ -121,7 +140,7 @@ def partition_bwd(softmax_scale, is_causal, window_size, mesh, arg_shapes, resul
             axis_name = l
             axis_size = mesh.shape[axis_name]
             # ring attention
-            return mesh, partial(ring_bwd, softmax_scale, is_causal, axis_name, axis_size), result_shardings, arg_shardings
+            return mesh, partial(ring_bwd, softmax_scale=softmax_scale, is_causal=is_causal, axis_name=axis_name, axis_size=axis_size, mha_bwd=_flash_mha_bwd_hlo), result_shardings, arg_shardings
         else:
             result_shardings = q_sharding, q_sharding, q_sharding
             lse_sharding = NamedSharding(mesh, P(n,h,l))
@@ -133,103 +152,3 @@ def fwd(*args):
 _flash_mha_bwd_hlo_sharded.def_partition(
     infer_sharding_from_operands=infer_sharding_bwd,
     partition=partition_bwd)
-
-# ==== Ring Forward ====
-
-def ring_fwd(softmax_scale, is_causal, axis_name, axis_size, q,k,v):
-    [n,l,h,d] = q.shape
-
-    q_ix = jax.lax.axis_index(axis_name)
-    k_ix = jax.lax.axis_index(axis_name)
-
-    o = jnp.zeros([n,l,h,d], jnp.float32)
-    lse = jnp.full([n,h,l], float('-inf'), jnp.float32)
-
-    # scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])
-    def f(c, a):
-        (k, v, o, lse, k_ix) = c
-
-        o1, lse1 = o, lse
-        if is_causal:
-            o2, lse2 = jax.lax.switch((k_ix < q_ix).astype(jnp.int32) + (k_ix <= q_ix).astype(jnp.int32),
-                                    [
-                                        lambda q,k,v: (jnp.zeros([n,l,h,d], q.dtype), jnp.full([n,h,l], float('-inf'), jnp.float32)),
-                                        lambda q,k,v: _flash_mha_fwd_hlo(q,k,v, softmax_scale=softmax_scale, is_causal=True, window_size=(-1,-1)),
-                                        lambda q,k,v: _flash_mha_fwd_hlo(q,k,v, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1)),
-                                    ], q, k, v)
-        else:
-            o2, lse2 = _flash_mha_fwd_hlo(q,k,v, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))
-        o2 = o2.astype(jnp.float32)
-
-        mx = jnp.maximum(lse1,lse2)
-        mn = jnp.minimum(lse1,lse2)
-        lse = jnp.log1p(jnp.exp(mn-mx)) + mx
-
-        o = (o1 * rearrange(jnp.exp(lse1 - lse), 'n h l -> n l h 1') +
-             o2 * rearrange(jnp.exp(lse2 - lse), 'n h l -> n l h 1'))
-        
-        k2 = jax.lax.ppermute(k, axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
-        v2 = jax.lax.ppermute(v, axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
-        k_ix = jax.lax.ppermute(k_ix, axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
-
-        return ((k2, v2, o, lse, k_ix), None)
-    acc = (k,v,o,lse,k_ix)
-    # We sadly have to manually unroll this because scan breaks the axis context preventing us from using ppermute (unroll=axis_size doesn't help either).
-    # Optimization barrier prevents instruction reordering so that ppermute and flash_mha execute concurrently.
-    for _ in range(axis_size):
-        acc, _ = f(acc, None)
-        acc = _optimization_barrier(acc)
-    (_,_,o,lse,_) = acc
-    # (_,_,o,lse), _ = jax.lax.scan(f,init,None,axis_size)
-    return o.astype(q.dtype), lse
-
-# ==== Ring Backward ===
-
-# This doesn't seem like the most efficient way to do this, kind of wasting compute by calculating every dq,dk,dv twice.
-# Should we send the accumulator for dk,dv cross-device instead? Relying on the fact that after a full cycle, they return to the starting device.
-def ring_bwd(softmax_scale, is_causal, axis_name, axis_size, do,q,k,v,o,lse):
-    [n,l,h,d] = q.shape
-
-    ix = jax.lax.axis_index(axis_name)
-
-    dq = jnp.zeros([n,l,h,d], jnp.float32)
-    dk = jnp.zeros([n,l,h,d], jnp.float32)
-    dv = jnp.zeros([n,l,h,d], jnp.float32)
-
-    # scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])
-    def f(acc, a):
-        (do2,q2,k2,v2,o2,lse2,ix2, dq,dk,dv) = acc
-
-        cmp = (ix2 < ix).astype(jnp.int32) + (ix2 <= ix).astype(jnp.int32) 
-        # 0: ix < ix2
-        # 1: ix = ix2
-        # 2: ix > ix2
-        if is_causal:
-            dqa = jax.lax.switch(cmp, [
-                                lambda q,k,v: jnp.zeros([n,l,h,d], q.dtype),
-                                lambda q,k,v: _flash_mha_bwd_hlo(do,q,k2,v2,o,lse, softmax_scale=softmax_scale, is_causal=True, window_size=(-1,-1))[0],
-                                lambda q,k,v: _flash_mha_bwd_hlo(do,q,k2,v2,o,lse, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))[0],
-                            ], q, k, v)
-            dka,dva = jax.lax.switch(cmp, [
-                                lambda q,k,v: _flash_mha_bwd_hlo(do2,q2,k,v,o2,lse2, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))[1:],
-                                lambda q,k,v: _flash_mha_bwd_hlo(do2,q2,k,v,o2,lse2, softmax_scale=softmax_scale, is_causal=True, window_size=(-1,-1))[1:],
-                                lambda q,k,v: (jnp.zeros([n,l,h,d], q.dtype),jnp.zeros([n,l,h,d], q.dtype)),
-                            ], q, k, v)
-        else:
-            dqa,_,_ = _flash_mha_bwd_hlo(do,q,k2,v2,o,lse, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))
-            _,dka,dva = _flash_mha_bwd_hlo(do2,q2,k,v,o2,lse2, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))
-
-        dq += dqa
-        dk += dka
-        dv += dva
-
-        (do2,q2,k2,v2,o2,lse2,ix2) = jax.lax.ppermute((do2,q2,k2,v2,o2,lse2,ix2), axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
-
-        return ((do2,q2,k2,v2,o2,lse2,ix2, dq,dk,dv), None)
-    acc = (do,q,k,v,o,lse,ix,dq,dk,dv)
-    # Unrolled as above.
-    for _ in range(axis_size):
-        acc, _ = f(acc, None)
-        acc = _optimization_barrier(acc)
-    (do2,q2,k2,v2,o2,lse2,ix2, dq,dk,dv) = acc
-    return dq.astype(q.dtype),dk.astype(q.dtype),dv.astype(q.dtype)

src/flash_attn_jax/ring_attention.py

@@ -0,0 +1,121 @@
+from functools import partial, wraps
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from jax import core, dtypes
+from jax.core import ShapedArray
+from jax.interpreters import batching
+from jax.interpreters import mlir
+from jax.interpreters import xla
+from jax.interpreters.mlir import ir
+from jax.lib import xla_client
+from jaxlib.hlo_helpers import custom_call
+from jax.experimental.custom_partitioning import custom_partitioning
+
+from jax.sharding import PartitionSpec as P
+from jax.sharding import Mesh
+from jax.sharding import NamedSharding
+from jax.sharding import PositionalSharding
+from jax._src.ad_checkpoint import _optimization_barrier
+
+from einops import rearrange
+import math
+
+# ==== Ring Forward ====
+
+def ring_fwd(q,k,v, axis_name, axis_size, mha_fwd, softmax_scale=None, is_causal=False):
+    [n,l,h,d] = q.shape
+    if softmax_scale is None:
+        softmax_scale = 1/math.sqrt(d)
+
+    q_ix = jax.lax.axis_index(axis_name)
+    k_ix = jax.lax.axis_index(axis_name)
+
+    o = jnp.zeros([n,l,h,d], jnp.float32)
+    lse = jnp.full([n,h,l], float('-inf'), jnp.float32)
+
+    # scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])
+    def f(c, a):
+        (k, v, o, lse, k_ix) = c
+
+        o1, lse1 = o, lse
+        if is_causal:
+            cmp = (k_ix < q_ix).astype(jnp.int32) + (k_ix <= q_ix).astype(jnp.int32) 
+            o2, lse2 = jax.lax.switch(cmp,
+                                    [
+                                        lambda: (jnp.zeros([n,l,h,d], q.dtype), jnp.full([n,h,l], float('-inf'), jnp.float32)),
+                                        lambda: mha_fwd(q,k,v, softmax_scale=softmax_scale, is_causal=True, window_size=(-1,-1)),
+                                        lambda: mha_fwd(q,k,v, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1)),
+                                    ])
+        else:
+            o2, lse2 = mha_fwd(q,k,v, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))
+        o2 = o2.astype(jnp.float32)
+
+        mx = jnp.maximum(lse1,lse2)
+        mn = jnp.minimum(lse1,lse2)
+        lse = jnp.log1p(jnp.exp(mn-mx)) + mx
+
+        o = (o1 * rearrange(jnp.exp(lse1 - lse), 'n h l -> n l h 1') +
+             o2 * rearrange(jnp.exp(lse2 - lse), 'n h l -> n l h 1'))
+        
+        k2 = jax.lax.ppermute(k, axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
+        v2 = jax.lax.ppermute(v, axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
+        k_ix = jax.lax.ppermute(k_ix, axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
+
+        return ((k2, v2, o, lse, k_ix), None)
+    acc = (k,v,o,lse,k_ix)
+    # Manually unroll this until https://github.com/google/jax/pull/20884 is merged.
+    # Optimization barrier prevents instruction reordering across loop iters, so that
+    # ppermute and flash_mha execute concurrently (though this is unreliable).
+    for _ in range(axis_size):
+        acc, _ = f(acc, None)
+        acc = _optimization_barrier(acc)
+    (_,_,o,lse,_) = acc
+    # (_,_,o,lse,_), _ = jax.lax.scan(f,acc,None,axis_size)
+    return o.astype(q.dtype), lse
+
+# ==== Ring Backward ===
+
+def ring_bwd(do,q,k,v,o,lse, axis_name, axis_size, mha_bwd, softmax_scale=None, is_causal=False):
+    [n,l,h,d] = q.shape
+    if softmax_scale is None:
+        softmax_scale = 1/math.sqrt(d)
+
+    ix = jax.lax.axis_index(axis_name)
+
+    dq = jnp.zeros([n,l,h,d], jnp.float32)
+    dk = jnp.zeros([n,l,h,d], jnp.float32)
+    dv = jnp.zeros([n,l,h,d], jnp.float32)
+
+    # scan :: (c -> a -> (c, b)) -> c -> [a] -> (c, [b])
+    def f(acc, _):
+        (k2,v2,dk2,dv2,ix2, dq) = acc
+
+        cmp = (ix2 < ix).astype(jnp.int32) + (ix2 <= ix).astype(jnp.int32) 
+        # 0: ix < ix2
+        # 1: ix = ix2
+        # 2: ix > ix2
+        if is_causal:
+            dqa, dka, dva = jax.lax.switch(cmp, (
+                                lambda: (jnp.zeros(q.shape, q.dtype), jnp.zeros(k.shape, k.dtype), jnp.zeros(v.shape, v.dtype)),
+                                lambda: mha_bwd(do,q,k2,v2,o,lse, softmax_scale=softmax_scale, is_causal=True, window_size=(-1,-1)),
+                                lambda: mha_bwd(do,q,k2,v2,o,lse, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))
+            ))
+        else:
+            dqa, dka, dva = mha_bwd(do,q,k2,v2,o,lse, softmax_scale=softmax_scale, is_causal=False, window_size=(-1,-1))
+
+        dq += dqa
+        dk2 += dka
+        dv2 += dva
+
+        (k2,v2,dk2,dv2,ix2) = jax.lax.ppermute((k2,v2,dk2,dv2,ix2), axis_name, [(i, (i+1)%axis_size) for i in range(axis_size)])
+
+        return ((k2,v2,dk2,dv2,ix2, dq), None)
+    acc = (k,v,dk,dv,ix, dq)
+    # See above (#20884).
+    for _ in range(axis_size):
+        acc, _ = f(acc, None)
+        acc = _optimization_barrier(acc)
+    (k,v,dk,dv,ix2, dq) = acc
+    return dq.astype(q.dtype),dk.astype(q.dtype),dv.astype(q.dtype)