From 85315a637c0b71e3ee340589b8c275c8ad8613f3 Mon Sep 17 00:00:00 2001 From: Clay Moore Date: Thu, 23 Jul 2026 15:25:30 -0500 Subject: [PATCH] perf(inference): drop training-only targets the forward never reads The forward pass never reads disto_target or token_to_rep_atom, two tensors inherited from the Boltz training featurizer: disto_target is the ground-truth distogram label (an O(N^2 * num_bins) one-hot built with cdist + one_hot), and token_to_rep_atom is the token-to-representative-atom gather (O(N_tok * M_atoms) one-hot) for the structure/confidence heads. The codebase is inference-only, so they are no longer built at all (the _compute_disto_target helper goes with them). They were previously constructed, collated, and copied to the device every batch, then ignored. Removing them is output-identical. Measured on an RTX 5080 at N=740 tokens: - input batch GPU memory: 155.1 -> 16.8 MiB (about 9x smaller) - host to device copy: about 9.8 -> 1.3 ms per batch (about 7x faster) disto_target alone is 64 MiB at N=512 and 144 MiB at N=768. Adds tests/test_inference_features.py: the featurizer omits both keys, and the forward (trunk and affinity head) is bit-identical whether or not they are injected into the batch (CI-runnable, ligand-only, no CCD). --- nesso/data/featurizer.py | 42 +-------- tests/test_inference_features.py | 143 +++++++++++++++++++++++++++++++ 2 files changed, 145 insertions(+), 40 deletions(-) create mode 100644 tests/test_inference_features.py diff --git a/nesso/data/featurizer.py b/nesso/data/featurizer.py index 742304a..c3948b4 100644 --- a/nesso/data/featurizer.py +++ b/nesso/data/featurizer.py @@ -70,21 +70,6 @@ def extract_esm_features( return s_chain -def _compute_disto_target( - disto_coords: Tensor, - min_dist: float = 2.0, - max_dist: float = 22.0, - num_bins: int = 64, -) -> Tensor: - """Compute one-hot distogram from token disto coordinates.""" - t_dists = torch.cdist(disto_coords.float(), disto_coords.float()) - boundaries = torch.linspace( - min_dist, max_dist, num_bins - 1, device=disto_coords.device - ) - distogram = (t_dists.unsqueeze(-1) > boundaries).sum(dim=-1).long().contiguous() - return one_hot(distogram, num_classes=num_bins).float() - - def select_subset_from_mask(mask: np.ndarray, p: float) -> np.ndarray: """Subsample True entries in mask using a geometric draw.""" num_true = int(np.sum(mask)) @@ -164,12 +149,6 @@ def process_token_features( # noqa: C901, PLR0915, PLR0912 pad_mask = torch.ones(len(token_data), dtype=torch.float) disto_mask = from_numpy(token_data["disto_mask"].copy()).float() - # Distogram target from token disto_coords - disto_coords = from_numpy(token_data["disto_coords"].copy()) - disto_target = _compute_disto_target( - disto_coords, min_dist=min_dist, max_dist=max_dist, num_bins=num_dist_bins - ) - # Token bond features if max_tokens is not None: pad_len = max_tokens - len(token_data) @@ -275,7 +254,6 @@ def process_token_features( # noqa: C901, PLR0915, PLR0912 res_type = pad_dim(res_type, 0, pad_len) pad_mask = pad_dim(pad_mask, 0, pad_len) disto_mask = pad_dim(disto_mask, 0, pad_len) - disto_target = pad_dim(pad_dim(disto_target, 0, pad_len), 1, pad_len) unspecified_oh = torch.zeros( pad_len, len(const.pocket_contact_info), dtype=torch.float32 ) @@ -294,7 +272,6 @@ def process_token_features( # noqa: C901, PLR0915, PLR0912 "type_bonds": bonds_type, "token_pad_mask": pad_mask, "token_disto_mask": disto_mask, - "disto_target": disto_target, "pocket_feature": pocket_feature, } @@ -353,7 +330,6 @@ def process_atom_features( ref_space_uid_list = [] coord_data_list = [] atom_to_token_list = [] - token_to_rep_atom_list = [] resolved_mask_list = [] chain_res_ids = {} @@ -428,11 +404,6 @@ def process_atom_features( conformer_pos = random.randn(n_atoms, 3).astype(np.float32) atom_conformer_list.append(conformer_pos) - disto_in_valid = next( - i for i, v in enumerate(valid_indices) if v == token["disto_idx"] - ) - token_to_rep_atom_list.append(atom_idx + disto_in_valid) - token_coords = structure.coords[offset + np.array(valid_indices)]["coords"] coord_data_list.append(token_coords[np.newaxis, ...]) resolved_mask_list.append(token_atoms["is_present"].copy()) @@ -482,14 +453,6 @@ def process_atom_features( torch.tensor(atom_to_token_list, dtype=torch.long), num_classes=num_token_classes, ) - token_to_rep_atom = one_hot( - torch.tensor(token_to_rep_atom_list, dtype=torch.long).clamp( - 0, max(0, num_atoms - 1) - ), - num_classes=num_atoms, - ) - if max_tokens is not None and L < max_tokens: - token_to_rep_atom = pad_dim(token_to_rep_atom, 0, max_tokens - L) pad_len = ( (num_atoms - 1) // atoms_per_window_queries + 1 @@ -507,9 +470,8 @@ def process_atom_features( ref_space_uid = pad_dim(ref_space_uid, 0, pad_len) coords = pad_dim(coords, 1, pad_len) atom_to_token = pad_dim(atom_to_token, 0, pad_len) - token_to_rep_atom = pad_dim(token_to_rep_atom, 1, pad_len) - return { + atom_features = { "ref_pos": ref_pos, "atom_resolved_mask": resolved_mask, "ref_atom_name_chars": ref_atom_name_chars, @@ -521,9 +483,9 @@ def process_atom_features( "coords": coords, "atom_pad_mask": pad_mask, "atom_to_token": atom_to_token, - "token_to_rep_atom": token_to_rep_atom, "res_index_to_conf_id": res_index_to_conf_id, } + return atom_features def process_esm_features( diff --git a/tests/test_inference_features.py b/tests/test_inference_features.py new file mode 100644 index 0000000..401c624 --- /dev/null +++ b/tests/test_inference_features.py @@ -0,0 +1,143 @@ +"""Inference featurization does not build training-only targets. + +The forward pass never reads ``disto_target`` or ``token_to_rep_atom`` (the +ground-truth distogram label and the token-to-representative-atom gather +inherited from the Boltz training featurizer). Since the codebase is +inference-only, those tensors are not built at all. + +These tests pin that: the featurizer omits both keys, and the model forward is +bit-identical whether or not those tensors are present in the batch (proving it +never consumed them). Ligand-only input keeps this CCD-free so it runs on plain +CI. +""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from nesso.data.featurizer import NessoFeaturizer +from nesso.data.inference import InferenceDataset, inference_collate +from nesso.data.tokenize import tokenize_structure +from nesso.data.types import Manifest, Structure, Tokenized +from nesso.data.yaml_input import parse_yaml +from nesso.model.models.nesso1 import Nesso1 +from numpy.random import RandomState + +REMOVED_KEYS = {"disto_target", "token_to_rep_atom"} + +_LIGAND_SMILES = "Cc1ccc(NC(=O)c2ccc(CN3CCN(C)CC3)cc2)cc1Nc1nccc(-c2cccnc2)n1" +_YAML = ( + "version: 1\n" + "sequences:\n" + " - ligand:\n" + " id: B\n" + f' smiles: "{_LIGAND_SMILES}"\n' +) + + +def _featurize(tmp_path: Path) -> dict: + mol_dir = tmp_path / "rdkit_conformers" + structures_dir = tmp_path / "structures" + esm_dir = tmp_path / "esm" + for d in (mol_dir, structures_dir, esm_dir): + d.mkdir(parents=True, exist_ok=True) + + yaml_path = tmp_path / "lig.yaml" + yaml_path.write_text(_YAML) + struct, record, _, _ = parse_yaml( + yaml_path, mol_dir, ccd_dict=None, record_id="lig" + ) + struct.dump(structures_dir / f"{record.id}.npz") + + struct = Structure.load(structures_dir / f"{record.id}.npz") + struct = struct.remove_invalid_chains(struct.mask.copy()) + tokens, bonds = tokenize_structure(struct) + tokenized = Tokenized(tokens=tokens, bonds=bonds, structure=struct, record=record) + + ds = InferenceDataset( + manifest=Manifest([record]), + target_dir=tmp_path, + featurizer=NessoFeaturizer( + esm_emb_dir=esm_dir, esm_emb_dim=1280, esm_num_layers=33 + ), + ligand_dir=mol_dir, + ccd_pkl=None, + ) + molecules = ds._setup_molecules(struct, str(record.id)) + torch.manual_seed(0) + feats = ds.featurizer.process( + tokenized, + record=record, + molecules=molecules, + random=RandomState(0), + atoms_per_window_queries=32, + binder_pocket_conditioned_prop=0.0, + max_tokens=None, + ) + feats["affinity_token_mask"] = (feats["mol_type"] == 3).float() + return inference_collate([feats]) + + +def _tiny_model(*, affinity: bool = False) -> Nesso1: + return Nesso1( + atom_s=16, + atom_z=16, + token_s=32, + token_z=32, + atom_feature_dim=387, + embedder_args={"atom_encoder_depth": 1, "atom_encoder_heads": 2}, + pairformer_model_args={"num_blocks": 1}, + esm_module_args={"esm_embed_dim": 1280}, + affinity_prediction=affinity, + affinity_model_args={ + "pairformer_args": {"num_blocks": 1}, + "transformer_args": {}, + } + if affinity + else None, + use_kernels=False, + predict_args={ + "refine_protein_inference": False, + "affinity_protein_cutoff": 15.0, + }, + ).eval() + + +def _assert_forward_ignores_injection(batch: dict, model: Nesso1) -> None: + n = batch["token_pad_mask"].shape[1] + m = batch["atom_pad_mask"].shape[1] + with torch.no_grad(): + out_lean = model(dict(batch), recycling_steps=2) + # Re-add the removed tensors; a forward that truly ignores them is unchanged. + injected = dict(batch) + injected["disto_target"] = torch.randn(1, n, n, 64) + injected["token_to_rep_atom"] = torch.zeros(1, n, m) + with torch.no_grad(): + out_injected = model(injected, recycling_steps=2) + + assert set(out_lean) == set(out_injected) + for key, value in out_lean.items(): + if isinstance(value, torch.Tensor): + assert torch.equal(value, out_injected[key]), key + + +def test_featurizer_omits_training_targets(tmp_path: Path) -> None: + batch = _featurize(tmp_path) + keys = {k for k, v in batch.items() if isinstance(v, torch.Tensor)} + assert REMOVED_KEYS.isdisjoint(keys), keys & REMOVED_KEYS + # Sanity: the live features the model does read are still present. + for live in ("res_type", "s_esm", "atom_to_token", "token_pad_mask"): + assert live in keys, live + + +def test_trunk_forward_ignores_injected_targets(tmp_path: Path) -> None: + torch.manual_seed(0) + _assert_forward_ignores_injection(_featurize(tmp_path), _tiny_model()) + + +def test_affinity_forward_ignores_injected_targets(tmp_path: Path) -> None: + torch.manual_seed(0) + model = _tiny_model(affinity=True) + _assert_forward_ignores_injection(_featurize(tmp_path), model)