diff --git a/bench_folding_grad.py b/bench_folding_grad.py new file mode 100644 index 0000000..d710a36 --- /dev/null +++ b/bench_folding_grad.py @@ -0,0 +1,120 @@ +"""Forward+backward benchmark: vanilla vs cueq-fast protenij. + +Uses an MSE-on-distogram-logits loss that exercises the full trunk +(input embedder, MSA module, template embedder, Pairformer stack, distogram +head). Diffusion + confidence head are excluded — they host no cueq-swapped +kernels and the sampler adds stochasticity that complicates grad timing. + +This is the relevant regime for training: cueq's cuda-fwd/cuda-bwd kernels for +triangle_attention (and triton VJP for triangle_multiplicative_update) get +exercised end-to-end. +""" +import os +os.environ["PROTENIX_DATA_ROOT_DIR"] = os.path.expanduser("~/.protenix") + +import argparse +import statistics +import time + +import equinox as eqx +import jax +import jax.numpy as jnp + +from protenix.backend import load_model + +from bench_folding import SEQUENCES, enable_cueq, build_features + + +def trunk_loss(model, features, key, n_cycle=4): + """MSE of distogram logits against a zero target. Scalar — grad OK. + + recycle() stop_gradients between cycles, so only the last cycle contributes + to grad. We still run n_cycle forwards to match production wall time. + """ + emb = model.embed_inputs(input_feature_dict=features) + trunk = model.recycle( + initial_embedding=emb, input_feature_dict=features, + recycling_steps=n_cycle, key=key, + ) + logits = model.distogram_head(trunk.z) + return jnp.mean(jnp.square(logits)) + + +def _block(tree): + jax.tree.map( + lambda v: v.block_until_ready() if hasattr(v, "block_until_ready") else None, + tree, + ) + + +def _time(fn, *, n_iter): + # First call includes compile; time it separately. + t0 = time.perf_counter() + out = fn() + _block(out) + compile_plus_run = time.perf_counter() - t0 + + times = [] + for _ in range(n_iter): + t0 = time.perf_counter() + out = fn() + _block(out) + times.append(time.perf_counter() - t0) + return compile_plus_run, times + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default="protenix_base_default_v1.0.0") + parser.add_argument("--n_cycle", type=int, default=4) + parser.add_argument("--n_iter", type=int, default=3) + parser.add_argument("--msa_dir", default="output_test_predict/msa", + help="mmseqs2 MSA cache dir; set to '' for dummy MSA") + args = parser.parse_args() + + print(f"device: {jax.devices()[0]}") + print(f"loading {args.model}...") + baseline = load_model(args.model) + cueq = enable_cueq(baseline) + + msa_dir = args.msa_dir if args.msa_dir else None + n_cycle = args.n_cycle + + def _fwd(model, features, key): + return trunk_loss(model, features, key, n_cycle=n_cycle) + + fwd = eqx.filter_jit(_fwd) + fwd_bwd = eqx.filter_jit(eqx.filter_value_and_grad(_fwd)) + + print(f"\nn_cycle={n_cycle} n_iter={args.n_iter} " + f"msa={'real (mmseqs2)' if msa_dir else 'dummy'}") + print("loss = mean(distogram_logits ** 2)") + + for name, seq in SEQUENCES.items(): + features = build_features(seq, name, msa_dir=msa_dir) + n_tok = int(features["token_index"].shape[0]) + print(f"\n=== {name} ({len(seq)} residues, {n_tok} tokens, " + f"msa={tuple(features['msa'].shape)}) ===") + key = jax.random.PRNGKey(0) + + results = {} + for label, model in [("baseline", baseline), ("cueq-fast", cueq)]: + c_fwd, t_fwd = _time(lambda: fwd(model, features, key), n_iter=args.n_iter) + c_bb, t_bb = _time(lambda: fwd_bwd(model, features, key), n_iter=args.n_iter) + mf = statistics.mean(t_fwd) * 1e3 + mb = statistics.mean(t_bb) * 1e3 + results[label] = (mf, mb, mb - mf) + print(f" {label:10s} " + f"fwd mean={mf:6.1f}ms (compile+1st={c_fwd:6.2f}s) " + f"fwd+bwd mean={mb:6.1f}ms (compile+1st={c_bb:6.2f}s) " + f"bwd≈{mb - mf:6.1f}ms") + + mf_b, mfb_b, bwd_b = results["baseline"] + mf_c, mfb_c, bwd_c = results["cueq-fast"] + print(f" speedup (cueq vs baseline): " + f"fwd={mf_b/mf_c:.2f}x bwd≈{bwd_b/bwd_c:.2f}x " + f"fwd+bwd={mfb_b/mfb_c:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/protenix/protenij.py b/protenix/protenij.py index f43d0b8..dd5b3f4 100644 --- a/protenix/protenij.py +++ b/protenix/protenij.py @@ -33,6 +33,25 @@ import protenix.model.protenix import protenix.openfold_local.model.primitives + +# Monkey-patch cuequivariance_jax 0.8.x: its triangle_attention custom_vjp fwd +# rule returns (a, lse, amax) as a tuple, but the primitive binding returns a +# list — JAX's custom_vjp matcher rejects the mismatch during backward tracing. +# Harmless if cueq isn't installed or is never used for a backward pass. +try: + from cuequivariance_jax.triangle._triangle_attention import ( + triangle_attention_custom_vjp as _cueq_ta_cvjp, + ) + _cueq_orig_fwd = _cueq_ta_cvjp.fwd + + def _cueq_patched_fwd(*args, **kwargs): + primal, residuals = _cueq_orig_fwd(*args, **kwargs) + return list(primal), residuals + + _cueq_ta_cvjp.fwd = _cueq_patched_fwd +except ImportError: + pass + def move_final_dim_to_dim(x, dim: int): # permute_final_dims n_dim = len(x.shape) @@ -283,21 +302,40 @@ class TriangleMultiplication(AbstractFromTorch): linear_z: backend.Linear sigmoid: any _outgoing: bool + # Static field (part of treedef) so stacking/partitioning never treats the + # bool as a dynamic leaf — older skeletons default it to False on load. + use_cueq: bool = eqx.field(default=False, static=True) + + @staticmethod + def from_torch(m): + return TriangleMultiplication( + linear_a_p=from_torch(m.linear_a_p), + linear_a_g=from_torch(m.linear_a_g), + linear_b_p=from_torch(m.linear_b_p), + linear_b_g=from_torch(m.linear_b_g), + layer_norm_in=from_torch(m.layer_norm_in), + layer_norm_out=from_torch(m.layer_norm_out), + linear_g=from_torch(m.linear_g), + linear_z=from_torch(m.linear_z), + sigmoid=from_torch(m.sigmoid), + _outgoing=m._outgoing, + ) def __call__( self, z_in: Float[Array, "... N N C_z"], mask: Bool[Array, "... N N"] | None, ) -> Float[Array, "... N N C_z"]: - #with jax.default_matmul_precision("float32"): if mask is None: mask = jnp.ones_like(z_in, shape=z_in.shape[:-1]) - mask = mask[..., None] + if getattr(self, "use_cueq", False): + return self._cueq_forward(z_in, mask) + mask_e = mask[..., None] z = self.layer_norm_in(z_in) - a = self.linear_a_p(z) * mask * self.sigmoid(self.linear_a_g(z)) - b = self.linear_b_p(z) * mask * self.sigmoid(self.linear_b_g(z)) + a = self.linear_a_p(z) * mask_e * self.sigmoid(self.linear_a_g(z)) + b = self.linear_b_p(z) * mask_e * self.sigmoid(self.linear_b_g(z)) ### permute_final_dims (2, 0, 1) if self._outgoing: @@ -307,7 +345,6 @@ def __call__( a = einops.rearrange(a, "... a b c -> ... c b a") b = einops.rearrange(b, "... a b c -> ... c a b") p = a @ b - # return p x = einops.rearrange(p, "... a b c -> ... b c a") x = self.layer_norm_out(x) @@ -317,6 +354,47 @@ def __call__( return x * g + def _cueq_forward(self, z_in, mask): + """Triangle multiplicative update via NVIDIA cuequivariance JAX kernel. + + Re-uses existing Linear/LayerNorm weights — cueq wants a/b projections + concatenated along the output axis (2*D_in, D_in). No new parameters. + """ + import cuequivariance_jax as cuex + + def stack_bias(ba, bb): + if ba is None and bb is None: + return None + if ba is None: + ba = jnp.zeros_like(bb) + if bb is None: + bb = jnp.zeros_like(ba) + return jnp.concatenate([ba, bb], axis=0) + + p_in_w = jnp.concatenate([self.linear_a_p.weight, self.linear_b_p.weight], axis=0) + g_in_w = jnp.concatenate([self.linear_a_g.weight, self.linear_b_g.weight], axis=0) + p_in_b = stack_bias(self.linear_a_p.bias, self.linear_b_p.bias) + g_in_b = stack_bias(self.linear_a_g.bias, self.linear_b_g.bias) + + return cuex.triangle_multiplicative_update( + z_in, + direction="outgoing" if self._outgoing else "incoming", + mask=mask, + norm_in_weight=self.layer_norm_in.weight, + norm_in_bias=self.layer_norm_in.bias, + p_in_weight=p_in_w, + p_in_bias=p_in_b, + g_in_weight=g_in_w, + g_in_bias=g_in_b, + norm_out_weight=self.layer_norm_out.weight, + norm_out_bias=self.layer_norm_out.bias, + p_out_weight=self.linear_z.weight, + p_out_bias=self.linear_z.bias, + g_out_weight=self.linear_g.weight, + g_out_bias=self.linear_g.bias, + eps=self.layer_norm_in.eps, + ) + class Attention(AbstractFromTorch): c_q: int # input dimension of query @@ -388,6 +466,18 @@ class TriangleAttention(AbstractFromTorch): linear: Linear mha: Attention inf: float + # Static: see TriangleMultiplication.use_cueq for rationale. + use_cueq: bool = eqx.field(default=False, static=True) + + @staticmethod + def from_torch(m): + return TriangleAttention( + starting=m.starting, + layer_norm=from_torch(m.layer_norm), + linear=from_torch(m.linear), + mha=from_torch(m.mha), + inf=m.inf, + ) def __call__( self, @@ -403,21 +493,70 @@ def __call__( x = self.layer_norm(x) - # [*, I, 1, 1, J] - mask_bias = (self.inf * (mask - 1))[..., :, None, None, :] - # [*, H, I, J] triangle_bias = einops.rearrange(self.linear(x), "... A B C -> ... C A B") # [*, 1, H, I, J] triangle_bias = einops.rearrange(triangle_bias, "... H I J -> ... 1 H I J") - x = self.mha(q_x=x, kv_x=x, biases=[mask_bias, triangle_bias]) + if getattr(self, "use_cueq", False): + x = self._cueq_forward(x, mask, triangle_bias) + else: + # [*, I, 1, 1, J] + mask_bias = (self.inf * (mask - 1))[..., :, None, None, :] + x = self.mha(q_x=x, kv_x=x, biases=[mask_bias, triangle_bias]) if not self.starting: x = einops.rearrange(x, "... I J C -> ... J I C") return x + def _cueq_forward(self, x, mask, triangle_bias): + """Triangle attention via NVIDIA cuequivariance JAX kernel. + + Reuses mha's linear projections, gate, and output projection; only the + attention kernel itself is swapped. Checkpoint weights are unchanged. + """ + import cuequivariance_jax as cuex + + mha = self.mha + H = mha.no_heads + D = mha.c_hidden + + q = einops.rearrange(mha.linear_q(x), "... N S (H D) -> ... N H S D", H=H) + k = einops.rearrange(mha.linear_k(x), "... N S (H D) -> ... N H S D", H=H) + v = einops.rearrange(mha.linear_v(x), "... N S (H D) -> ... N H S D", H=H) + + # cueq expects rank-5 [B, N, H, S, D]; add a B=1 if absent. + added_batch = q.ndim == 4 + if added_batch: + q, k, v = q[None], k[None], v[None] + + bias = triangle_bias if triangle_bias.ndim == 5 else triangle_bias[None] + mask_cueq = mask.astype(bool)[..., :, None, None, :] + if mask_cueq.ndim == 4: + mask_cueq = mask_cueq[None] + + scale = 1.0 / math.sqrt(D) + # TF32 / tensor-core path. HIGHEST (IEEE fp32) lacks a backward kernel + # in cueq 0.8.x and only matches baseline precision, which is itself + # TF32 on Ampere+, so DEFAULT is the only mode worth supporting. + out = cuex.triangle_attention( + q, k, v, bias, mask_cueq, scale, precision=jax.lax.Precision.DEFAULT, + ) + # cueq returns [output, logsumexp, max_val]; keep just the output. + o = out[0] if isinstance(out, (tuple, list)) else out + + if added_batch: + o = o[0] + + o = einops.rearrange(o, "... N H S D -> ... N S H D") + if mha.linear_g is not None: + g = jax.nn.sigmoid(mha.linear_g(x)) + g = einops.rearrange(g, "... N S (H D) -> ... N S H D", H=H) + o = o * g + o = einops.rearrange(o, "... N S H D -> ... N S (H D)") + return mha.linear_o(o) + def _attention( q, diff --git a/pyproject.toml b/pyproject.toml index 77a91a9..526dbf6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "einops>=0.8.0", "equinox>=0.13.0", "gemmi>=0.6.5", - "jax>=0.6.2", + "jax>=0.8.1", "ml-collections>=1.1.0", "modelcif>=0.7", "numpy>=2.1.0", @@ -26,6 +26,9 @@ dependencies = [ "setuptools>=80.9.0", "tqdm>=4.67.1", "huggingface_hub>=0.20.0", + "cuequivariance-jax>=0.8.1,<0.9", + "cuequivariance-ops-jax-cu12>=0.8.1,<0.9", + "jax-triton>=0.3.1", ] [project.optional-dependencies] diff --git a/test_predict.py b/test_predict.py index 53f28fe..c40e9cd 100644 --- a/test_predict.py +++ b/test_predict.py @@ -48,7 +48,7 @@ eqx_path = os.path.join(CACHE_DIR, f"{MODEL_NAME}") print(f"Loading JAX model from {eqx_path}...") -jax_model = load_model(eqx_path) +jax_model = load_model(MODEL_NAME) # Override diffusion parameters for vanilla ODE sampling jax_model = eqx.tree_at(lambda m: m.gamma0, jax_model, 0.0)