Add MPS support, upgrade PyTorch, refactor packaging, and fix CI + correctness/security bugs#66
Open
fnachon wants to merge 29 commits into
Open
Add MPS support, upgrade PyTorch, refactor packaging, and fix CI + correctness/security bugs#66fnachon wants to merge 29 commits into
fnachon wants to merge 29 commits into
Conversation
To simplify the installation
add mps as valid device
add mps device
switch from setup.py to .toml file
torch.linalg.cholesky is not implemented for mps. Do the calculation on the cpu and move the resulting tensor back to the gpu
torch.linalg.det is not implemented for mps. do the calculation on the cpu
torch.linalg.cummax and torch.linalg.cummin not implemented for mps. Move the calculation to the cpu then move the result back to the gpu.
torch.linalg.eigh not implemented for mps. Do in cpu
Update for running on Apple Silicon GPU with mps
- Apply attention mask before softmax (not after) so weights sum to 1; add nan_to_num guard for all-masked rows (ghost chains in mixed batches) - call_with_explicit_mps_cpu_fallback now catches RuntimeError in addition to UserWarning so linalg.svd on MPS falls back to CPU instead of crashing - Cache key includes tensor shape to prevent wrong MPS→CPU routing across different-sized calls - Unify three independent tensor-tree traversals into one (_iter_tensors) - Extract _with_mps_fallback_cache shared by both public MPS helpers - Add get_default_device() to utility/torch.py; Chroma.__init__ uses it instead of an inline ternary - Remove dead self.attention from MultiHeadAttention.__init__ Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Bumps the pinned torch version and regenerates both conda lockfiles. MPS remains available on osx-arm64 with the conda-forge cpu_generic build. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The blanket *.cif rule in .gitignore prevented tests/resources fixtures from ever being committed, so CI (which checks out a clean clone) failed with FileNotFoundError across the data-utility, layers, and models jobs. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
test_proclass_conditioner needs a registered API token to download proprietary model weights, which isn't available in CI; mark it integration/slow like the equivalent tests/models/test_chroma.py does so it's excluded from the quick-test run. test_multi_head_attention_matches_legacy_path compares a batched-einsum projection against a legacy per-head loop; they're only equivalent up to floating-point reduction order, which differs across BLAS backends and made the default torch.allclose tolerance too strict on CI's runner. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Found via a systematic review of chroma/layers, chroma/models, chroma/data, and chroma/utility. Each fix was verified with a targeted reproduction plus the existing test suite. Correctness: - Attention: fix double-scaling of logits in the gated Attention class (divided by d_k instead of sqrt(d_k)), and fix a crash in the gate path (torch.sigmoid doesn't take a dim kwarg). - system.py: AtomLocationView.defined() checked `is not None` but undefined locations are marked with NaN, so it always returned True and let NaN coordinates leak into exported CIF files; to_XCS() crashed past ~62 chains because an unused chain-letter fallback was computed unconditionally; ResidueView.get_backbone() passed atom-type strings where integer indices were expected; _read_pdb's `options` argument was silently dropped due to a positional/keyword mismatch; ARG NH1/NH2 disambiguation renamed a location object instead of the atom (added AtomView.rename to fix); closed file handles in from_CIF/from_PDB/to_CIF/to_PDB. - chroma.py: _protein_list_to_XCS ignored its device argument; removed _design_ar, a dead method referencing a never-assigned attribute. - graph_design.py: fixed a tuple/None bug that crashed loss() whenever sidechains=False. - symmetry.py: subsample() returned a chain map that didn't correspond to the sampled coordinates. - mvn.py: added the missing RRt_clamp_inverse buffer needed by multiply_inverse_covariance. - procap.py: a single padded sample in a batch silently dropped sequence conditioning for the whole batch instead of just that sample; re-enabled a swallowed IndexError in context reordering. - geometry.py: removed a byte-identical duplicate VirtualAtomsCA class. Security (chroma/utility/api.py): - The registered access token was sent as a query param to any URL passed as a weights path; added a host allowlist. - Weight downloads were written directly to a predictable, world-writable /tmp path with no atomic rename; now stream to a temp file and os.replace into place, and add a request timeout. - register_key wrote the token to a world-readable config.json; now created with 0600 permissions. Performance: - model.py: save_model raised NotImplementedError for s3: paths after already serializing the full state_dict; now checks first. - system.py: to_XCS called get_residue(i) in a loop, which re-walks the chain list from the start each time; now iterates residues() once. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR brings this fork's
mainbranch up to date with a series of changes made since diverging from upstream: Apple Silicon (MPS) support, a PyTorch upgrade, a packaging/test refactor, and a round of correctness/security bug fixes found during a manual code review.Apple Silicon (MPS) support
environment_Mac.yaml,pyproject.toml) to install/run on Apple Silicon.Upgrade PyTorch to 2.12.1
Packaging, tests, and locked environments refactor
tests/resources/*.cifwere unintentionally excluded by a blanket*.cifrule in.gitignore(left over from before these tests existed); added an explicit exception and committed the fixtures so CI can actually run these tests.Correctness bugs, security issues, and dead code (found via code review)
Verified with targeted reproductions plus the existing test suite:
Attentionclass (divided byd_kinstead ofsqrt(d_k)), and fixed a crash in the gate path (torch.sigmoiddoesn't take adimkwarg).chroma/data/system.py:AtomLocationView.defined()checkedis not None, but undefined locations are marked withNaN, so it always returnedTrueand letNaNcoordinates leak into exported CIF files;to_XCS()crashed past ~62 chains because a fallback chain-letter computation ran unconditionally even when unused;ResidueView.get_backbone()passed atom-type strings where integer indices were expected;_read_pdb'soptionsargument was silently dropped due to a positional/keyword mismatch; ARG NH1/NH2 disambiguation renamed a location object instead of the atom (addedAtomView.renameto fix); closed file handles infrom_CIF/from_PDB/to_CIF/to_PDB; replaced anO(n * chains)per-residue lookup into_XCSwith a single linear pass.chroma/models/chroma.py:_protein_list_to_XCSignored itsdeviceargument; removed_design_ar, a dead method referencing a never-assigned attribute.chroma/models/graph_design.py: fixed a tuple/Nonebug that crashedloss()wheneversidechains=False.chroma/layers/structure/symmetry.py:subsample()returned a chain map that didn't correspond to the sampled coordinates.chroma/layers/structure/mvn.py: added the missingRRt_clamp_inversebuffer needed bymultiply_inverse_covariance.chroma/models/procap.py: a single padded sample in a batch silently dropped sequence conditioning for the whole batch instead of just that sample; re-enabled a swallowedIndexErrorin context reordering.chroma/layers/structure/geometry.py: removed a byte-identical duplicateVirtualAtomsCAclass.chroma/utility/api.py(security): the registered access token was sent as a query parameter to any URL passed as a weights path, with no host check; added a host allowlist. Weight downloads were written directly to a predictable, world-writable/tmppath with no atomic rename; downloads now stream to a temp file andos.replaceinto place, plus a request timeout was added.register_keywrote the token to a world-readableconfig.json; now created with0600permissions.chroma/utility/model.py:save_modelraisedNotImplementedErrorfors3:paths only after already serializing the full state dict; now checks first, avoiding wasted work and an orphaned temp file.Test plan
pytest -m "not integration and not slow"acrosstests/data,tests/models,tests/layers,tests/utility(163 passed)tests/utility/test_api.py::test_api(integration, full network download + token flow)