diff --git a/experiments/it_opt_1gpu.toml b/experiments/it_opt_1gpu.toml new file mode 100644 index 00000000..3722d33e --- /dev/null +++ b/experiments/it_opt_1gpu.toml @@ -0,0 +1,41 @@ +description = "IT-opt s_plus_z on 1 GPU: all proteins run sequentially in one worker." + +# Run with: +# sampleworks-runs --preset it_opt_1gpu +# +# One job per protein, one worker, so 40 proteins is 40 jobs run back to back. The sibling +# it_opt_4gpu.toml is this same configuration with gpu_count = 4, which splits those jobs over +# four workers. +# +# The settings below are the defaults of it_opt_scratch/slurm_ensemble_and_score.sh, so a run +# here is comparable to the recorded trees. To sample a different mode, change which-latent +# (single = s, pair = z, both = s+z); for the unguided or coordinate-DPS references use +# scalers = "pure_guidance" with step-scaler-type = "none" or "noisespace" instead. + +[defaults] +DATA_DIR = "/data/inputs" +RESULTS_DIR = "/data/results/it_opt_1gpu" +PROTEINS_CSV = "${DATA_DIR}/proteins.csv" + +# ensemble-sizes and gradient-weights MUST stay single values. Their defaults are "1 2 4 8" and +# "0.01 0.1 0.2", which would expand one job per protein into twelve. gradient-weights only +# feeds the step scaler, which latent_opt does not use, so its value is arbitrary. +# align-to-input matches the scratch runner, which set it unconditionally; the grid defaults it +# off. bond-length-weight is named explicitly because the scratch runner defaulted it to 0.0. +[shared_args] +proteins = "${PROTEINS_CSV}" +model = "protenix" +scalers = "latent_opt" +which-latent = "both" +ensemble-sizes = "8" +gradient-weights = "0.0" +num-diffusion-steps = 200 +bond-length-weight = 5e-5 +align-to-input = true + +[[jobs]] +name = "s_plus_z" +env = "protenix" +gpu_count = 1 +output_subdir = "s_plus_z" +args = {} diff --git a/experiments/it_opt_4gpu.toml b/experiments/it_opt_4gpu.toml new file mode 100644 index 00000000..d8694685 --- /dev/null +++ b/experiments/it_opt_4gpu.toml @@ -0,0 +1,44 @@ +description = "IT-opt s_plus_z on 4 GPUs: proteins split across four workers, one per GPU." + +# Run with: +# sampleworks-runs --preset it_opt_4gpu +# +# Identical to it_opt_1gpu.toml except for gpu_count and jobs-per-gpu. run_grid_search sets +# max_workers = len(gpus) * jobs-per-gpu and gives worker i the slice jobs[i::max_workers], so +# 4 GPUs x 2 gives 8 workers and 40 proteins become 5 jobs each rather than 40 in a row. +# +# jobs-per-gpu = 2 matches the recorded run (--jobs-per-gpu 2 on 4 GPUs). Two jobs share a card, +# each holding its own copy of the weights: a measured single job peaked near 11 GiB of an 80 GiB +# H100, so two fit with wide margin. Raise it further only after checking VRAM on the largest +# protein in the set, since cost scales with atom count. +# +# Settings are the defaults of it_opt_scratch/slurm_ensemble_and_score.sh; see the sibling +# preset for how to switch mode. + +[defaults] +DATA_DIR = "/data/inputs" +RESULTS_DIR = "/data/results/it_opt_4gpu" +PROTEINS_CSV = "${DATA_DIR}/proteins.csv" + +# ensemble-sizes and gradient-weights MUST stay single values. Their defaults are "1 2 4 8" and +# "0.01 0.1 0.2", which would expand one job per protein into twelve. gradient-weights only +# feeds the step scaler, which latent_opt does not use, so its value is arbitrary. +# align-to-input matches the scratch runner, which set it unconditionally; the grid defaults it +# off. bond-length-weight is named explicitly because the scratch runner defaulted it to 0.0. +[shared_args] +proteins = "${PROTEINS_CSV}" +model = "protenix" +scalers = "latent_opt" +which-latent = "both" +ensemble-sizes = "8" +gradient-weights = "0.0" +num-diffusion-steps = 200 +bond-length-weight = 5e-5 +align-to-input = true + +[[jobs]] +name = "s_plus_z" +env = "protenix" +gpu_count = 4 +output_subdir = "s_plus_z" +args = { jobs-per-gpu = 2 } diff --git a/run_grid_search.py b/run_grid_search.py index 846ab2cf..51836488 100755 --- a/run_grid_search.py +++ b/run_grid_search.py @@ -19,7 +19,12 @@ from loguru import logger as log from sampleworks.utils.guidance_constants import GuidanceType, StructurePredictor -from sampleworks.utils.guidance_script_arguments import GuidanceConfig, JobConfig, JobResult +from sampleworks.utils.guidance_script_arguments import ( + add_latent_opt_args, + GuidanceConfig, + JobConfig, + JobResult, +) from sampleworks.utils.protein_input import ProteinInput @@ -267,12 +272,29 @@ def run_grid_search( successful = 0 failed = 0 - max_workers = len(gpus) - log.info(f"Running {len(jobs)} jobs with {max_workers} parallel workers") + # Workers per GPU (--jobs-per-gpu) lets two jobs share a card, which is worthwhile because a + # single job leaves the GPU idle while it featurizes and writes output. Two constraints shape + # the worker count: + # never more workers than jobs -- a surplus worker gets an empty queue, and the + # worker_job_queues[worker_num][0] lookup below then raises IndexError before any model + # loads. Easy to hit once jobs_per_gpu multiplies the count. + # keep it a whole multiple of the GPU count, so every card carries the same load. Clamping + # to the job count alone leaves e.g. 5 workers on 4 GPUs: one card runs two jobs + # concurrently while the other three run one and then idle. + max_workers = min(len(gpus) * args.jobs_per_gpu, len(jobs)) + if max_workers > len(gpus): + max_workers -= max_workers % len(gpus) + gpus_used = min(max_workers, len(gpus)) + log.info( + f"Running {len(jobs)} jobs with {max_workers} parallel workers " + f"({gpus_used} GPU(s) x {max_workers // gpus_used} jobs/GPU)" + ) - # Divide the job among the workers: + # Divide the job among the workers. The worker index is no longer the GPU ordinal, so the + # device wraps with i % len(gpus): workers 0..7 across 4 GPUs pair up as 0,1,2,3,0,1,2,3. + # Passing i directly would ask for cuda:4..cuda:7 on a 4-GPU box and fail immediately. worker_job_queues = [ - [build_args_for_process_pool(j, args, i) for j in jobs[i::max_workers]] + [build_args_for_process_pool(j, args, i % len(gpus)) for j in jobs[i::max_workers]] for i in range(max_workers) ] # we'll pickle each job queue separately and then execute each job queue in a separate process @@ -779,6 +801,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--augmentation", action="store_true", help="Enable augmentation") parser.add_argument("--align-to-input", action="store_true", help="Align to input structure") + # Latent optimization (IT-opt) arguments, used when --scalers includes latent_opt. + # Registered from the shipped adder rather than restated here so that flag names, types and + # defaults stay identical to the sampleworks-guidance CLI. + add_latent_opt_args(parser) + # RF3-specific arguments parser.add_argument( "--disable-chiral-features", @@ -803,6 +830,14 @@ def parse_args() -> argparse.Namespace: default="auto", help="Max parallel jobs (default: auto = number of GPUs)", ) + parser.add_argument( + "--jobs-per-gpu", + type=int, + default=1, + choices=(1, 2), + help="Concurrent jobs per GPU (default: 1). Each holds its own copy of the model weights, " + "so 2 needs roughly twice the VRAM; 2 is the tested ceiling", + ) parser.add_argument( "--dry-run", action="store_true", diff --git a/src/sampleworks/core/scalers/latent_optimization.py b/src/sampleworks/core/scalers/latent_optimization.py index ec2212df..32de5227 100644 --- a/src/sampleworks/core/scalers/latent_optimization.py +++ b/src/sampleworks/core/scalers/latent_optimization.py @@ -126,6 +126,70 @@ def __call__(self, latents: Sequence[Tensor], baselines: Sequence[Tensor]) -> Te return torch.stack(terms).sum() +class _PerMemberStepper: + """Denoise one ensemble member at a time, each with its OWN trunk latent. + + The it_opt "multiple-leaf" scheme gives every ensemble member its own latent, so ``s``/``z`` + carry a leading ``ensemble_size`` batch dim. Stock model diffusion modules instead take an + un-batched conditioning and broadcast it internally, so a batched latent either crashes + (Protenix) or mis-broadcasts (Boltz's ``multiplicity``). This adapter sidesteps that in a + model-agnostic way: for each member it slices that member's latent, runs the wrapped model's + normal un-batched ``step`` on that member alone, and stacks the results. Members stay + independent, so gradients stay per-member; the summed density reward couples them only through + the ensemble average. Cost is N forwards instead of one batched call, with an equivalent result. + + ``featurize`` / ``initialize_from_prior`` pass straight through to the wrapped model. + """ + + def __init__( + self, + model, + io: AttrLatentIO, + *, + optimize_single: bool, + optimize_pair: bool, + ensemble_size: int, + ): + self._model = model + self._io = io + self._optimize_single = optimize_single + self._optimize_pair = optimize_pair + self._ensemble_size = ensemble_size + + def step(self, x_t: Tensor, t, *, features: GenerativeModelInput) -> Tensor: + """Loop the wrapped model's ``step`` over ensemble members; stack the per-member results.""" + cond = features.conditioning + per_member: list[Tensor] = [] + for i in range(self._ensemble_size): + # Slice only the OPTIMIZED latents (they carry the ensemble batch dim); a non-optimized + # latent stays the shared un-batched baseline already on ``cond``. + cond_i = cond + if self._optimize_single: + cond_i = self._io.write_single(cond_i, self._io.read_single(cond)[i]) + if self._optimize_pair: + # read_pair returns None when the io addresses no pair rep. sample() only sets + # optimize_pair together with a pair_attr, so this is a misconfigured io reaching + # us directly; say so rather than failing on a None subscript mid-denoise. + pair = self._io.read_pair(cond) + if pair is None: + raise ValueError( + "optimize_pair is set, but the io addresses no pair representation." + ) + cond_i = self._io.write_pair(cond_i, pair[i]) + t_i = t + if isinstance(t, Tensor) and t.ndim >= 1 and t.shape[0] == x_t.shape[0]: + t_i = t[i : i + 1] + features_i = GenerativeModelInput(conditioning=cond_i) + per_member.append(self._model.step(x_t[i : i + 1], t_i, features=features_i)) + return torch.cat(per_member, dim=0) + + def featurize(self, *args, **kwargs): + return self._model.featurize(*args, **kwargs) + + def initialize_from_prior(self, *args, **kwargs): + return self._model.initialize_from_prior(*args, **kwargs) + + class LatentOptimization: """Trajectory scaler that optimizes the model's ``s``/``z`` latents (IT-opt). @@ -259,6 +323,16 @@ def sample( schedule = sampler.compute_schedule(num_steps=self.num_steps) grad_enabler = _GradEnablingScaler() + # Denoise per member so each uses its own latent: stock model diffusion modules take an + # un-batched conditioning, so the batched per-member latents can't go through in one call. + stepper = _PerMemberStepper( + model, + io, + optimize_single=self.optimize_single, + optimize_pair=self.optimize_pair, + ensemble_size=self.ensemble_size, + ) + # --- optional coordinate-space geometry penalty ------------------------- # BondGeometryReward penalizes stretched bonds and steric clashes in the denoised structure, # curbing the overshoot where an aggressive latent update trades geometry for density fit. @@ -279,7 +353,7 @@ def sample( for outer in tqdm(range(self.outer_steps)): optimizer = torch.optim.Adam(latents, lr=self.learning_rate) # a fresh, persistent Adam round_losses = self._optimize_one_round( - model=model, + model=stepper, sampler=sampler, reward=reward, features=features, @@ -307,7 +381,7 @@ def sample( # --- final clean sampling round with the optimized latents -------------- final_coords, trajectory, losses = self._sample_with_frozen_latents( - model=model, + model=stepper, sampler=sampler, reward=reward, io=io, @@ -341,8 +415,12 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): detached baselines (anchor targets), and per-latent anchor weights. Each leaf is a detached clone made ``requires_grad=True`` -- a true leaf severed from any trunk graph, so Adam updates it directly (leaves persist and are - updated in place across rounds and steps). Shapes are preserved (whatever - the wrapper caches), so no assumption is made about a batch dimension. + updated in place across rounds and steps). + + Each leaf gets a leading ``ensemble_size`` batch dimension -- one INDEPENDENT latent per + ensemble member, all cloned from the same trunk baseline (the it_opt "multiple-leaf" + scheme), so members can diverge rather than share one latent. The baseline kept for the + anchor stays un-batched and broadcasts across members. """ conditioning = features.conditioning latents: list[Tensor] = [] @@ -367,7 +445,14 @@ def _leaf_latents(self, features: GenerativeModelInput, io: AttrLatentIO): "conditioning does not expose it. Check the attribute names for this model." ) baseline = baseline.detach() - leaf = baseline.clone().requires_grad_(True) + # it_opt "multiple-leaf" scheme: give each ensemble member its OWN latent. We stack + # ensemble_size independent copies of the trunk baseline into a leading batch dim, so + # each member gets its own gradient and can diverge, instead of collapsing onto one + # shared latent. (The reference does the same via batch-expand-then-clone.) Stock + # diffusion modules take an un-batched conditioning, so _PerMemberStepper slices this + # batch dim back apart and runs one forward per member at denoise time. + member_copies = [baseline for _ in range(self.ensemble_size)] + leaf = torch.stack(member_copies).requires_grad_(True) # The conditioning is a frozen dataclass, so setattr would raise; replace() returns a # copy with this one field swapped and every sidecar field left untouched. conditioning = dataclasses.replace(conditioning, **{attr: leaf}) diff --git a/src/sampleworks/utils/cif_utils.py b/src/sampleworks/utils/cif_utils.py index 82989b9a..1e4b96e5 100644 --- a/src/sampleworks/utils/cif_utils.py +++ b/src/sampleworks/utils/cif_utils.py @@ -12,6 +12,7 @@ from loguru import logger from sampleworks.utils.atom_array_utils import ( + BLANK_ALTLOC_IDS, find_all_altloc_ids, save_structure_to_cif, select_altloc, @@ -247,6 +248,79 @@ def resolve_mixed_hetatm_atom_altlocs(cif_path: Path | str) -> Path: return tmp_path +def remap_altlocs_to_ab(cif_path: Path | str) -> Path: + """Relabel a residue's non-A/B alternate altloc(s) to fill the A/B slots. + + Some depositions label a residue's two conformers ``A``/``C`` (or ``A``/``D``, etc.) instead of + ``A``/``B``. Tools that assume the two altlocs are literally named A and B -- e.g. a + min-RMSD-to-altloc-A/altloc-B scorer -- then silently drop the second conformer. This is an + optional pre-processing step that, per residue, maps the alternate label onto the free ``B`` + slot: an atom with altloc ``A`` stays ``A`` and the other non-blank altloc becomes ``B``; if + ``A`` is absent, the two labels are assigned to ``A``/``B`` in sorted order. + + Only residues with **exactly two** non-blank altlocs that are **not already** ``{A, B}`` are + touched. Residues that are already A/B, have a single or blank altloc, or carry three or more + altlocs (ambiguous which two to keep) are left unchanged, so the metric's numbers do not move + for structures that were already A/B. + + A warning is logged for every remapped ``(chain, residue)`` position. + + Parameters + ---------- + cif_path + Path to the input CIF file. + + Returns + ------- + Path + Path to a remapped temporary CIF file if any position was changed, or the original + ``cif_path`` unchanged if none were. + """ + cif_path = Path(cif_path) + atom_array = load_any(cif_path, altloc="all", extra_fields=["occupancy", "b_factor"]) + if isinstance(atom_array, AtomArrayStack): + atom_array = atom_array[0] + if not hasattr(atom_array, "altloc_id"): + return cif_path # no altloc annotation -> nothing to remap + + altloc = atom_array.altloc_id.copy() # per-atom altloc characters, mutated in place below + chain_id = atom_array.chain_id + res_id = atom_array.res_id + remapped = 0 + + for chain in np.unique(chain_id): + for rid in np.unique(res_id[chain_id == chain]): + pos = (chain_id == chain) & (res_id == rid) + labels = sorted(set(altloc[pos].tolist()) - BLANK_ALTLOC_IDS) + if len(labels) != 2 or set(labels) == {"A", "B"}: + continue + if "A" in labels: + mapping = {next(x for x in labels if x != "A"): "B"} + else: + mapping = {labels[0]: "A", labels[1]: "B"} + for src, dst in mapping.items(): + sel = pos & (altloc == src) + altloc[sel] = dst + remapped += int(sel.sum()) + logger.warning( + f"Chain {chain}, residue {rid}: remapped altlocs {labels} -> A/B so the alternate " + "conformer is not dropped by A/B-only tooling." + ) + + if remapped == 0: + return cif_path + + atom_array.set_annotation("altloc_id", altloc) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".cif", prefix="sampleworks_altloc_ab_", delete=False + ) as tmp_file: + tmp_path = Path(tmp_file.name) + + save_structure_to_cif(atom_array, tmp_path) + logger.info(f"Wrote altloc-remapped CIF to temporary file: {tmp_path}") + return tmp_path + + def add_category_to_cif( ciffile: CIFFile, data: dict[str, Any], diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index 53e773e3..bccc75ea 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -178,6 +178,18 @@ def validate_model_checkpoint( return str(checkpoint_path) +# Inference-time latent optimization (IT-opt) tunables. Named once here so that from_cli() +# and populate_config_for_guidance_type(), the two independent paths that copy CLI values onto +# a config, cannot drift apart as flags are added. +_LATENT_OPT_ATTRS = ( + "which_latent", + "learning_rate", + "outer_steps", + "anchor_weight", + "max_grad_norm", + "bond_length_weight", +) + # Attributes set dynamically by add_*_args helpers that should be copied # from a parsed argparse.Namespace onto a GuidanceConfig instance. _DYNAMIC_ATTRS = [ @@ -193,12 +205,7 @@ def validate_model_checkpoint( "guidance_interval", # latent optimization (IT-opt) -- must be listed here or from_cli() drops the parsed # values and _run_guidance()'s getattr(args, ...) always sees the defaults (flags = no-ops). - "which_latent", - "learning_rate", - "outer_steps", - "anchor_weight", - "max_grad_norm", - "bond_length_weight", + *_LATENT_OPT_ATTRS, # model-specific "model_checkpoint", "method", @@ -404,6 +411,24 @@ def populate_config_for_guidance_type(self, job: JobConfig, args: argparse.Names self.fk_lambda = args.fk_lambda self.fk_resampling_interval = args.fk_resampling_interval self.ensemble_size = job.ensemble_size + elif job.scaler == GuidanceType.LATENT_OPT: + # IT-opt tunables are scalar flags rather than grid axes, so they come straight off + # the driver's namespace. This copy is what makes them work at all: the grid builds + # its GuidanceConfig directly and never calls from_cli(), so the _DYNAMIC_ATTRS loop + # that normally moves parsed values onto a config is unreachable from here. Without + # it every --which-latent / --learning-rate / --outer-steps flag is a silent no-op, + # left at the add_latent_opt_args default that __post_init__ already seeded. + for attr in _LATENT_OPT_ATTRS: + value = getattr(args, attr, None) + if value is not None: + setattr(self, attr, value) + # _run_guidance reads ensemble_size directly rather than via getattr, so it must + # be set here or a latent_opt job raises AttributeError. + self.ensemble_size = job.ensemble_size + # Deliberately no step_size / step_scaler_type here. add_latent_opt_args does not + # define them, so leaving them unset makes _run_guidance fall through to the same + # defaults the scratch runner hit (step_size 0.01, noisespace). Setting them would + # silently change what the recorded runs did. else: self.step_size = job.gradient_weight self.step_scaler_type = args.step_scaler_type diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index f62f2351..1f58443b 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -507,6 +507,10 @@ def _run_guidance(args: GuidanceConfig, guidance_type: str, model_wrapper, devic structure = annotate_structure_for_protenix( structure, + # Root Protenix's per-sample input dump (protenix_input*.json) under the job's output + # dir; otherwise out_dir falls back to the input id, resolves against the CWD, and + # leaves a stray / folder there on every run. + out_dir=str(Path(args.output_dir) / "protenix_input"), recycling_steps=recycling_steps, # Disable diffusion shared-vars cache for LATENT_OPT so gradients can # flow to z_trunk; cached tensors can otherwise become stale. diff --git a/tests/models/test_latent_optimization.py b/tests/models/test_latent_optimization.py index 3d5f258e..b00a63f2 100644 --- a/tests/models/test_latent_optimization.py +++ b/tests/models/test_latent_optimization.py @@ -92,7 +92,9 @@ def test_leaf_latents_makes_requires_grad_leaves(features): assert leaf.requires_grad and leaf.is_leaf # a true leaf Adam can update directly for leaf, base in zip(latents, baselines): assert not base.requires_grad # the anchor target is detached - torch.testing.assert_close(leaf.detach(), base) # leaf starts exactly at the baseline + # one leaf per ensemble member on a leading batch dim; the baseline stays un-batched + for member in leaf.detach(): + torch.testing.assert_close(member, base) # every member starts at the baseline # the rewritten conditioning holds the SAME leaf objects, so the model reads # the optimizable tensor rather than the original cached one assert new_features.conditioning.s is latents[0] diff --git a/tests/rewards/test_geometry.py b/tests/rewards/test_geometry.py new file mode 100644 index 00000000..ca2e0048 --- /dev/null +++ b/tests/rewards/test_geometry.py @@ -0,0 +1,15 @@ +"""A test case for testing the geometry reward + +This test case is designed to verify the reward function that evaluates the geometry of a +molecular structure. It checks whether the reward function correctly computes the reward +based on the provided atomic coordinates and other relevant parameters. +""" + +import gemmi +from sampleworks.core.rewards.geometry import _covalent_radius + + +def test_known_element_matches_gemmi(): + r = _covalent_radius("C") + assert r > 0, "Covalent radius for Carbon should be greater than 0" + assert r == gemmi.Element("C").covalent_r diff --git a/tests/runs/conftest.py b/tests/runs/conftest.py index 78204e7a..c0a492ff 100644 --- a/tests/runs/conftest.py +++ b/tests/runs/conftest.py @@ -41,6 +41,8 @@ def force_pixi_argv(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: "boltz2_md", "boltz2_xrd", "full_8gpu", + "it_opt_1gpu", + "it_opt_4gpu", "protenix", "protenix_dual", "protpardelle",