import jax
from jax import numpy as jnp
from jax import jit, vmap
from flax import nnx
from typing import Callable, Optional
from jaxtyping import Array, PyTree
from jax.typing import DTypeLike
# layer = nnx.MultiHeadAttention(
# num_heads=8, in_features=5, qkv_features=16, decode=False, rngs=nnx.Rngs(0)
# )
[docs]
class AttentionBlock(nnx.Module):
"""Self-attention block with a fixed fp32 compute island.
``flax.nnx.MultiHeadAttention`` (flax 0.12.7) computes its softmax in
whatever ``dtype`` it is given -- there is no internal fp32 upcast for
the attention logits/softmax the way some other implementations provide.
Rather than bolt on a custom ``attention_fn`` just to force fp32 softmax,
this block keeps *all* of its internal math (LayerNorm + MultiHeadAttention)
at ``dtype=jnp.float32``, ignoring the ``dtype`` compute-precision knob for
those internals. This is a deliberately larger fp32 island than a single
softmax op, but the bulk of the model's FLOPs live in the ``DenseBlock``
MLP stack (``widening_factor`` x wider), which does honor the bf16
``dtype`` knob, so this island has a small cost in practice.
The island is for the *math* only: the block's output (after the
optional skip connection) is downcast to the requested ``dtype`` before
being returned, mirroring the codebase's established fp32-island idiom
(e.g. ``QKNorm.__call__``'s ``.astype(v.dtype)`` in ``flux1/layers.py``).
Without this downcast, the fp32 residual (``x_in``, captured post
fp32-LayerNorm) would silently re-promote every downstream block's
output back to fp32 via JAX's bf16+fp32 promotion rule on the skip-add,
defeating the bf16 knob's memory/bandwidth benefit for the whole
inter-block residual stream.
``param_dtype`` (master-weight storage) is unaffected and still threads
through normally.
"""
def __init__(
self,
din: int,
num_heads: int,
features: int,
skip_connection: bool,
rngs: nnx.Rngs,
dtype: DTypeLike = jnp.float32,
param_dtype: DTypeLike = jnp.float32,
):
[docs]
self.skip_connection = skip_connection
# Only used to downcast the block's output back to compute dtype
# after the fp32-island math below (see class docstring).
# fp32 island: dtype is intentionally fixed to float32 regardless of
# the dtype knob above (see class docstring).
[docs]
self.layer_norm = nnx.LayerNorm(
din, rngs=rngs, dtype=jnp.float32, param_dtype=param_dtype
)
[docs]
self.attn = nnx.MultiHeadAttention(
in_features=din,
num_heads=num_heads,
qkv_features=features,
decode=False,
rngs=rngs,
dtype=jnp.float32,
param_dtype=param_dtype,
)
[docs]
def __call__(self, x: jnp.ndarray, mask: jnp.ndarray | None) -> jnp.ndarray:
x = self.layer_norm(x)
x_in = x
x = self.attn(x, mask=mask)
if self.skip_connection:
x = x + x_in
return jnp.asarray(x, dtype=self.dtype)
[docs]
class DenseBlock(nnx.Module):
"""MLP block (the FLOPs-dominant part of the transformer) with an fp32
LayerNorm island.
The LayerNorm math runs in fp32 (mirrors ``AttentionBlock``'s
normalization treatment); the wide hidden-layer matmuls run in the
requested compute ``dtype``. The block's output (after the context-merge
and optional skip connection) is downcast to ``dtype`` before being
returned -- otherwise the fp32 residual (``x_in``, captured post
fp32-LayerNorm) would silently re-promote the output back to fp32 via
JAX's bf16+fp32 promotion rule on the skip-add, defeating the bf16
knob's memory/bandwidth benefit for the whole inter-block residual
stream (see ``AttentionBlock`` docstring for the same pattern).
"""
def __init__(
self,
din,
dcontext,
num_hidden_layers,
widening_factor: int,
act: Callable,
skip_connection: bool,
rngs: nnx.Rngs,
dtype: DTypeLike = jnp.float32,
param_dtype: DTypeLike = jnp.float32,
):
[docs]
self.skip_connection = skip_connection
# Only used to downcast the block's output back to compute dtype
# after the fp32-island LayerNorm math (see class docstring).
n_features = din
# fp32 island: LayerNorm stays fp32 regardless of the compute dtype
# knob (mirrors AttentionBlock's normalization treatment).
[docs]
self.layer_norm = nnx.LayerNorm(
din, rngs=rngs, dtype=jnp.float32, param_dtype=param_dtype
)
hidden_blocks = []
hidden_blocks.append(
nnx.Linear(
n_features,
widening_factor * n_features,
rngs=rngs,
dtype=dtype,
param_dtype=param_dtype,
)
)
n_features *= widening_factor
for i in range(1, num_hidden_layers):
hidden_blocks.append(
nnx.Linear(
n_features,
n_features,
rngs=rngs,
dtype=dtype,
param_dtype=param_dtype,
)
)
hidden_blocks.append(
nnx.Linear(n_features, din, rngs=rngs, dtype=dtype, param_dtype=param_dtype)
)
[docs]
self.hidden_blocks = nnx.List(hidden_blocks)
[docs]
self.context_block = nnx.Linear(
dcontext, din, rngs=rngs, dtype=dtype, param_dtype=param_dtype
)
return
[docs]
def __call__(self, x, context):
x = self.layer_norm(x)
x_in = x
for i in range(len(self.hidden_blocks) - 1):
x = self.hidden_blocks[i](x)
x = self.act(x)
x = self.hidden_blocks[-1](x)
if context is not None:
context_emb = self.context_block(context)
context_emb = self.act(context_emb)
while context_emb.ndim < x.ndim:
context_emb = context_emb[..., None, :]
x = x + context_emb
if self.skip_connection:
x = x + x_in
return jnp.asarray(x, dtype=self.dtype)
