Open ae-foster opened 1 month ago
Hi,
So this is relatively complicated issue. The reason is that the current code would be incompatible with anything that uses custom_vjp
which the pallas attention does. This is due to the fact that internally kfac-jax
differentiates value_and_grad_func(*args)[0]
-- e.g. extracting the forward pass from the user provided value_and_grad_func
. However, as I have now found out - the jax.value_and_grad
strips away the custom_vjp
primitive. Hence, at the point where kfac-jax
differentiates the function there is no custom_vjp
and instead there is just mha_forward
. For this we might need to think about, how to address it. A few workarounds, while we decide what is best:
Extend the optimizer, such that it accepts the forward_func, and internally creates a value_and_grad_func
with jax.value_and_grad
. This way you can pass the real forward_func to the curvature estimator.
Instead of using jax.value_and_grad
pass a function that returns f(x), jax.grad(f)(x)
. This way the forward pass that kfac-jax
will extract will contain the custom_vjp
.
PS: I tried 1) and it works, however there is one more complication that arises later - KFAC uses forward mode autodiff, while you can not apply forward mode AD to a custom_vjp
function.
To address this my understanding is, that because Jax currently does not support custom_vjp and custom_jvp together, and a pallas_call would not work with the automatic derivation of vjp for custom_jvp, what needs to happen is to create a new jax.core.Primitive for the pallas_call which to wrap correctly the vjp and jvp. Also someone to implement the jvp rule for the attention Pallas call. Atm I don't think I have the capacity to do that though.
Thanks @botev , this has given me a good direction to look in.
I have hacked my kfac-jax
as per your suggestion 1), and I can now get the new error
TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.
presumably the same as you got.
I can take a look at implementing a pallas kernel for custom_jvp
for attention. Will keep you posted
Adding a custom_jvp might be unfeasible as the current way that Jax works it allows for only one custom decorator. Most likely the whole attention call needs to be wrapped as jax primitive. But if you get anything to work let me know.
Could I (temporarily) drop the custom vjp
and have only a custom jvp
? We could at least see if it works
you can try? I think jax might be able to infer the vjp from the jvp?
Quick update on this. I tried implementing a custom_jvp
on a simple pallas
kernel for addition instead of full blown attention. Unfortunately, we hit the issue that transposition rules are not implemented for pallas_call
and hence we cannot run vjp
, see https://github.com/google/jax/issues/19146
Fortunately, it seems the jax
folks are working on related issues although things are not yet fully documented / released. I found this snippet based on https://github.com/google/jax/issues/9129 . Importantly, custom_transpose
is available but not fully documented in the latest jax.
The simplest example of this is
def add_vectors_kernel(x_ref, y_ref, o_ref):
x, y = x_ref[...], y_ref[...]
o_ref[...] = x + y
@jax.jit
def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
return pl.pallas_call(add_vectors_kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype))(
x, y
)
@jax.jit
def add_vectors_fwd(x: jax.Array, y: jax.Array) -> jax.Array:
s = pl.pallas_call(add_vectors_kernel, out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype))(
x, y
)
return s, s
@jax.jit
def add_vectors_tangent(residuals, tangents, out_avals=None) -> jax.Array:
x_dot, y_dot = tangents
grad_out = pl.pallas_call(
add_vectors_kernel, out_shape=jax.ShapeDtypeStruct(x_dot.shape, x_dot.dtype)
)(x_dot, y_dot)
return grad_out
@jax.jit
def add_vectors_bwd(residuals, cotangents) -> jax.Array:
return (cotangents, cotangents)
def custom_vjp_by_custom_transpose(fun, tangent, fwd, bwd):
fun = custom_jvp(fun)
@fun.defjvp
def jvp(primals, tangents):
outs, residuals = fwd(*primals)
tan_out_types = tree_map(lambda o: core.get_aval(o).at_least_vspace(), outs)
tan_fn = custom_transpose(tangent)
tan_fn.def_transpose(bwd)
return outs, tan_fn(tan_out_types, residuals, tangents)
return fun
pallas_add = custom_vjp_by_custom_transpose(
add_vectors, add_vectors_tangent, add_vectors_fwd, add_vectors_bwd
)
This results in a function that can be differentiated with both jvp
and vjp
. If in my earlier script, I replace
attended = attention.mha(k, k, k, None)
with
attended = pallas_add(k, k)
then I can successfully run a step of KFAC :)
We may be on the right path.
I can confirm that the above approach also works with attention. I guess the only change needed from the kfac-jax
end then is the ability to directly input the value function so that the custom_jvp
instruction is preserved.
Can you try replacing with:
def custom_vjp_by_custom_transpose(fun, tangent, fwd, bwd):
fun = custom_jvp(fun)
@fun.defjvp
def jvp(primals, tangents):
_, residuals = fwd(*primals)
outs = fun(*primals) # Make the primal output a `custom_jvp` call.
tan_out_types = tree_map(lambda o: core.get_aval(o).at_least_vspace(), outs)
tan_fn = custom_transpose(tangent)
tan_fn.def_transpose(bwd)
return outs, tan_fn(tan_out_types, residuals, tangents)
return fun
Hi @botev Thanks for your tip above about calling fun
itself inside the custom_jvp
rule. We tried this out and it works with the old convert_value_and_grad_to_value
pattern.
Some further updates on migrating attention from our testing script into our real code
custom_tranpose
pattern doesn't work in more complex settings because it is not compatible with vmap
. (This is a known issue with that feature, which is probably why it's not documented. Seems like there is a new suggestion for how to combine jvp
and vjp
https://github.com/google/jax/pull/22457 )kfac-jax
by default does not correctly register dense layers with batching dimensions. In our production code, we have extra patterns to fix thisImplicitCurvature
estimation that needs jvp
and vjp
vjp
on the model. That is, we do not actually need to patch the attention layer to support jvp
if we only have dense layers in the model.
First, thank you for creating and releasing this invaluable resource.
What I am trying to do
I would like to combine
kfax-jax
with fused attention frompallas
.As far as I understand, this should theoretically be trivial: the attention operation itself contains no tagged parameters and no tagged losses, so KFAC should simply propagate the forward and backward passes.
Reproducing failure with the current repo
Environment and set-up
jax==0.4.28
jaxlib==0.4.28+cuda12.cudnn89
dm-haiku==0.0.12
kfac-jax
cloned frommain
Reproduction
This fails with the following error message
My own attempts to investigate
We can see that the issue arises when running
vjp
on the manipulated version of the model. To try and diagnose the issue, I tried a minimal re-implementation of this part of the KFAC algorithmThis runs without error. This reimplementation was based on my reading and understanding of the
kfac-jax
codebase and might possibly miss something important. Of course, it is missing the loss and layer tagging part; I had hoped that wasn't relevant.Request for assistance
I would really appreciate your advice on this task. Specifically
kfac-jax
to work withpallas
attention? I would be happy to help work on a fix with guidance on where to lookkfac-jax
to work specifically with the attention operation, assuming that this operation contains no layer tags and no loss tags?