Issue with scan + custom_vjp where backward pass is in higher precision than forward pass

[liamclarkza]

Description

I am trying to use a function that requires a gradient to be in higher precision on the backward pass than on a forward pass. A minimal reproducible example of this is shown below:

import jax
import jax.numpy as jnp

# Dummy function that requires float32 gradient
@jax.custom_vjp
def identity(x):
    return x

def identity_fwd(x):
    return x, None

def identity_bwd(res, g):
    return (g.astype(jnp.float32),)

identity.defvjp(identity_fwd, identity_bwd)


def model(x):

    def scannable_identity_fn(carry, _):
        return identity(carry), None
    
    x = identity(x) # without these identity calls, the scan in/out types are different
    x = jax.lax.scan(scannable_identity_fn, x, None, length=2)[0]
    x = identity(x)
    return x


# Running with x in bfloat16 fails due to scan, as the gradient is fp32
x = jnp.asarray(8, dtype=jnp.bfloat16)
grad = jax.grad(model)(x)
print(grad)

When using this with a scan I get the following error.

---------------------------------------------------------------------------
MLIRError                                 Traceback (most recent call last)
File ~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1212, in lower_jaxpr_to_module(***failed resolving arguments***)
   [1211](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1211) try:
-> [1212](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1212)   if not ctx.module.operation.verify():
   [1213](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1213)     raise ValueError(
   [1214](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1214)         "Cannot lower jaxpr with verifier errors. " +
   [1215](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/jax/_src/interpreters/mlir.py:1215)         dump_module_message(ctx.module, "verification"))

MLIRError: Verification failed:
error: "jit(scan)/jit(main)/while/body/closed_call"(callsite("model"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":26:8) at callsite("<module>"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":33:7) at callsite("InteractiveShell.run_code"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3577:20) at callsite("InteractiveShell.run_ast_nodes"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3517:19) at callsite("InteractiveShell.run_cell_async"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3334:29) at callsite("_pseudo_sync_runner"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/async_helpers.py":128:8) at callsite("InteractiveShell._run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3130:21) at callsite("InteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3075:21) at callsite("ZMQInteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/zmqshell.py":549:15) at "IPythonKernel.do_execute"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/ipkernel.py":449:26))))))))))): 'func.call' op operand type mismatch: expected operand type 'tensor<bf16>', but provided 'tensor<f32>' for operand number 0
 note: "jit(scan)/jit(main)/while/body/closed_call"(callsite("model"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":26:8) at callsite("<module>"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":33:7) at callsite("InteractiveShell.run_code"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3577:20) at callsite("InteractiveShell.run_ast_nodes"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3517:19) at callsite("InteractiveShell.run_cell_async"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3334:29) at callsite("_pseudo_sync_runner"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/async_helpers.py":128:8) at callsite("InteractiveShell._run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3130:21) at callsite("InteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3075:21) at callsite("ZMQInteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/zmqshell.py":549:15) at "IPythonKernel.do_execute"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/ipkernel.py":449:26))))))))))): see current operation: %6 = "func.call"(%arg2) <{callee = @None}> : (tensor<f32>) -> tensor<f32>

The above exception was the direct cause of the following exception:

JaxStackTraceBeforeTransformation         Traceback (most recent call last)
File <frozen runpy>:198, in _run_module_as_main()

File <frozen runpy>:88, in _run_code()

File ~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel_launcher.py:18
     [16](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel_launcher.py:16) from ipykernel import kernelapp as app
---> [18](https://file+.vscode-resource.vscode-cdn.net/Users/liam/PycharmProjects/waffle/~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel_launcher.py:18) app.launch_new_instance()

File ~/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/traitlets/config/application.py:1075, in launch_instance()
...
	'func.call' op operand type mismatch: expected operand type 'tensor<bf16>', but provided 'tensor<f32>' for operand number 0
		at loc("jit(scan)/jit(main)/while/body/closed_call"(callsite("model"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":26:8) at callsite("<module>"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":33:7) at callsite("InteractiveShell.run_code"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3577:20) at callsite("InteractiveShell.run_ast_nodes"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3517:19) at callsite("InteractiveShell.run_cell_async"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3334:29) at callsite("_pseudo_sync_runner"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/async_helpers.py":128:8) at callsite("InteractiveShell._run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3130:21) at callsite("InteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3075:21) at callsite("ZMQInteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/zmqshell.py":549:15) at "IPythonKernel.do_execute"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/ipkernel.py":449:26))))))))))))
	see current operation: %6 = "func.call"(%arg2) <{callee = @None}> : (tensor<f32>) -> tensor<f32>
		at loc("jit(scan)/jit(main)/while/body/closed_call"(callsite("model"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":26:8) at callsite("<module>"("/var/folders/84/ltzfvw692yl9h64qv9g44c5c0000gn/T/ipykernel_75635/2738124691.py":33:7) at callsite("InteractiveShell.run_code"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3577:20) at callsite("InteractiveShell.run_ast_nodes"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3517:19) at callsite("InteractiveShell.run_cell_async"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3334:29) at callsite("_pseudo_sync_runner"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/async_helpers.py":128:8) at callsite("InteractiveShell._run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3130:21) at callsite("InteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/IPython/core/interactiveshell.py":3075:21) at callsite("ZMQInteractiveShell.run_cell"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/zmqshell.py":549:15) at "IPythonKernel.do_execute"("/Users/liam/PycharmProjects/waffle/.venv/lib/python3.11/site-packages/ipykernel/ipkernel.py":449:26))))))))))))
The module was dumped to jax_ir0034_jit_scan_verification.mlir.

While I can resolve this by using the same dtype on the forward and backward pass, I am using a similar setup to perform quantization of activations of a forward pass, but trying to keep gradients in a higher precision. Is this a limitation of using scan or is there a way to achieve different types for fwd and bwd in a scan?

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.5.0
jaxlib: 0.5.0
numpy:  2.0.2
python: 3.11.6 (v3.11.6:8b6ee5ba3b, Oct  2 2023, 11:18:21) [Clang 13.0.0 (clang-1300.0.29.30)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='Liams-Work-MacBook-Pro.local', release='24.3.0', version='Darwin Kernel Version 24.3.0: Thu Jan  2 20:24:24 PST 2025; root:xnu-11215.81.4~3/RELEASE_ARM64_T6030', machine='arm64')

Note: same issue occurs on other platforms

[liamclarkza]

Also, if you print the jaxpr for the grad from the above using print(jax.make_jaxpr(jax.grad(model))(x)), you can see there is an issue with the types in the reverse scan where the carry value types don't match. The bwd type conversion is happening in the scan now, which then causes it to break, instead of the conversion happening before the reverse scan.

{ lambda ; a:bf16[]. let
    _:bf16[] b:bf16[2] = scan[
      _split_transpose=False
      jaxpr={ lambda ; c:bf16[]. let  in (c, c) }
      length=2
      linear=(False,)
      num_carry=1
      num_consts=0
      reverse=False
      unroll=1
    ] a
    d:f32[] = scan[
      _split_transpose=False
      jaxpr={ lambda ; e:bf16[] f:bf16[]. let
          g:f32[] = convert_element_type[new_dtype=float32 weak_type=False] e
        in (g,) }
      length=2
      linear=(True, False)
      num_carry=1
      num_consts=0
      reverse=True
      unroll=1
    ] 1.0 b

[dfm]

Great questions - thanks for bringing this up! It's actually not totally surprising to me that this fails because we frequently assume that primal and tangent types match throughout JAX core. I'm actually more surprised that the error isn't raised sooner... I would have expected custom_vjp to throw a type mismatch error!

In the long run, this should be handled by defining a custom type that uses a different tangent dtype, but we don't yet have any APIs for implementing this. @dougalm is working on adding support for this, but I think it's still mostly in the brainstorming phase.

In the shorter term, you might be able to get the performance that you want by moving the custom_vjp up the stack, putting it above the scan. Then you could isolate the inconsistent dtypes within the backward function, explicitly casting at the boundaries. I'm not sure exactly what the implementation would look like, but that might be a direction to explore?

[dfm]

Following up: @mattjj tells me that there actually is an undocumented experimental API that can be used to get what you want today! It's not the most ergonomic API, and it has some limitations, but he shared the example that I've included below which I've annotated to add some context that I found useful when trying to understand it:

import jax
import jax.numpy as jnp
from jax.experimental import primal_tangent_dtype

# "primal_tangent_dtype" allows us to define a custom dtype that permits a
# different tangent type.
primal_dtype = jnp.dtype('bfloat16')
tangent_dtype = jnp.dtype('float32')
dt = primal_tangent_dtype(primal_dtype, tangent_dtype)

# A version of jax.value_and_grad that handles the custom tangent dtype
# by passing it explicitly to vjp. If you use 
def value_and_grad2(f):
  # Note the jit here. Dtypes defined using primal_tangent_dtype cannot be used
  # outside of jit, and they can't be passed as input or output.
  @jax.jit
  def f_grad(x):
    value, f_vjp = jax.vjp(f, x.astype(dt))
    x_bar, = f_vjp(jnp.ones((), dtype=tangent_dtype))
    return value.astype(primal_dtype), x_bar
  return f_grad

# An example calculation that operates in "primal_dtype".
@jax.custom_vjp
def add(x, y):
  assert x.dtype == dt
  assert y.dtype == dt
  # Here we explicitly cast to "primal_dtype" to perform the computatation,
  # then cast back to our custom type at the output.
  return (x.astype(primal_dtype) + y.astype(primal_dtype)).astype(dt)

def add_fwd(x, y):
  return add(x, y), None

# The backwards pass operates in "tangent_dtype".
def add_bwd(_, g):
  assert g.dtype == tangent_dtype
  return g, g

add.defvjp(add_fwd, add_bwd)

# This composes with scan properly.
def csum(xs):
  def body(c, x): return (add(c, x), x)
  zero = jnp.zeros((), dtype=primal_dtype).astype(dt)
  _, result = jax.lax.scan(body, zero, xs)
  return result

xs = jnp.arange(5., dtype=primal_dtype)
value_and_grad2(lambda x: csum(x)[-1])(xs)

Take a look to see if you can map your problem onto this example, and we're happy to help, with the one caveat that this is an experimental approach that @dougalm will hopefully be replacing in the (near?) future!

[liamclarkza]

Thanks, @dfm and @mattjj; it's really nice to see that this is in the pipeline! I was also initially thinking that a higher-level custom vjp might be necessary, but this would be nice to avoid if possible. By the way, is the experimental approach that @dougalm is working on related to this PR?

From a bit of experimenting with the experimental primal_tangent_dtype approach, I don't think this will work without a lot of re-writing/refactoring, as it appears that one has to ensure that all computations in a chain relying on the primal_tangent_dtype need to explicitly handle the "unwrapping" and "wrapping" the values to that it can be used in computations like standard Jax types.

For more context, I have provided an example of one of the ways in which I am trying to utilise higher precision gradients. Below is an example of a Flax model where I am quantizing checkpointed activations between "blocks" of a model that are being scanned over - something that would potentially be useful to apply to a transformer block, substantially reducing the memory footprint during training.

The example below works in its current form; however, we are forced into quantizing the gradient, which is something I would like to avoid for training stability. Changing the tangent_dtype when calling the dequantize_fn will result in the same failure shown in the minimalist example I provided above.

Thanks again for the feedback, and if you have any other thoughts on how to approach this, I'd be keen to hear them. For now I will look into the custom vjp approach.

import flax.linen as nn
from flax.training.train_state import TrainState
import jax
import jax.numpy as jnp
from jax import custom_vjp
import optax


def quantize_fn(primal_dtype, quantized_primal_dtype, tangent_dtype):
    def quantize_fwd(x, scale, amax_history):
        new_scale, new_history = nn.fp8_ops.update_fp8_meta(
            x, quantized_primal_dtype, scale, amax_history
        )
        qx = nn.fp8_ops.quantize(
            x, quantized_primal_dtype, nn.fp8_ops._fm32_to_float32(new_scale), primal_dtype
        )
        return (qx, new_scale), (new_scale, new_history)

    def quantize_bwd(res, g):
        dqx, _ = g
        new_scale, new_history = res
        return dqx.astype(tangent_dtype), new_scale, new_history

    @custom_vjp
    def quantize(x, scale, amax_history):
        (qx, new_scale), _ = quantize_fwd(x, scale, amax_history)
        return qx, new_scale

    quantize.defvjp(quantize_fwd, quantize_bwd)

    return quantize


def dequantize_fn(primal_dtype, tangent_dtype):
    def dequantize_fwd(x, scale):
        dqx = nn.fp8_ops.dequantize(x, primal_dtype, nn.fp8_ops._fm32_to_float32(scale))
        return dqx, None

    def dequantize_bwd(_, g):
        return g.astype(tangent_dtype), None # issue is that tangent_dtype has to be the same as the quantized_primal_dtype

    @custom_vjp
    def dequantize(x, scale):
        dqx, _ = dequantize_fwd(x, scale)
        return dqx

    dequantize.defvjp(dequantize_fwd, dequantize_bwd)

    return dequantize


def create_variables(module: nn.Module, amax_hist_length: int = 1024):
    """Create scale and amax history variables for quantization."""
    scale = module.variable(
        nn.fp8_ops.OVERWRITE_WITH_GRADIENT,
        "scale",
        nn.initializers.ones_init(),
        jax.random.key(0),
        (1,),
    )
    amax_history = module.variable(
        nn.fp8_ops.OVERWRITE_WITH_GRADIENT,
        "amax_history",
        nn.initializers.zeros_init(),
        jax.random.key(1),
        (amax_hist_length,),
    )
    return scale, amax_history


# EXAMPLE: Quantized activation checkpointing in Flax Model
EMB_DIM = 128
N_BLOCKS = 10
BATCH_SIZE = 64
N_CLASSES = 10
AMAX_HIST_LENGTH = 128

# FP8 Quantization Functions
quantize = quantize_fn(jnp.float32, jnp.float8_e4m3fn, jnp.float32)
dequantize = dequantize_fn(jnp.float32, jnp.float8_e4m3fn)

class Block(nn.Module):

    @nn.compact
    def __call__(self, carry):
        x, scale = carry
        x = dequantize(x=x, scale=scale)
        x = nn.Dense(EMB_DIM)(x)
        x = nn.relu(x)
        out_scale, out_amax_history = create_variables(
            self, amax_hist_length=AMAX_HIST_LENGTH
        )
        x, new_scale = quantize(x=x, scale=out_scale.value, amax_history=out_amax_history.value)
        return x, new_scale


class Model(nn.Module):

    @nn.compact
    def __call__(self, x):
        scanned_blocks = nn.remat_scan(
            Block,
            lengths=(N_BLOCKS,),
            policy=jax.checkpoint_policies.nothing_saveable,
            variable_axes={"params": 0, nn.fp8_ops.OVERWRITE_WITH_GRADIENT: 0},
            split_rngs={"params": True},
        )()
        scale, amax_history = create_variables(
            self, amax_hist_length=AMAX_HIST_LENGTH
        )
        x, new_scale = quantize(x=x, scale=scale.value, amax_history=amax_history.value)
        x, out_scale = scanned_blocks((x, new_scale))
        x = dequantize(x=x, scale=out_scale)
        return x


@jax.jit
def step(state, x, y):
    def loss_fn(params):
        logits = state.apply_fn(params, x)
        loss = optax.softmax_cross_entropy_with_integer_labels(logits, y)
        return jnp.mean(loss)

    grad_fn = jax.value_and_grad(loss_fn)
    loss, grads = grad_fn(state.params)
    new_state = state.apply_gradients(grads=grads)
    return new_state, loss


# Basic Training Loop:
key = jax.random.key(0)
x = jax.random.normal(key, (BATCH_SIZE, EMB_DIM))
y = jax.random.randint(key, (BATCH_SIZE,), 0, N_CLASSES)
model = Model()
train_state = TrainState.create(
    apply_fn=model.apply,
    params=model.init(key, x),
    tx=optax.adam(1e-4),
)
for i in range(100):
    train_state, loss = step(train_state, x, y)
    print(f"Step {i} loss: {loss}")