RosettaCommons / RosettaCommons/foundry

[BUG] Runtime TRITON error with RF3 on Collab T4 GPU

Open
#240 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
966
Forks
181
Avg merge
4d 4h
Merged PRs (30d)
2

Description

I am trying to run RF3 on a Google Collab T4 GPU following the example in the IPD Design Pipeline Collab but consistently get this error when I try to run prediction on a batch of sequences:

INFO:rf3.inference_engines.rf3:[rank: 0] Loading checkpoint from /root/.foundry/checkpoints/rf3_foundry_01_24_latest_remapped.ckpt...
WARNING:atomworks.ml:Using element type for atom names of atomized tokens.
INFO: Using bfloat16 Automatic Mixed Precision (AMP)
INFO:lightning.pytorch.utilities.rank_zero:Using bfloat16 Automatic Mixed Precision (AMP)
INFO:rf3.inference_engines.rf3:[rank: 0] Outputs will be written to /content/Predictions.
INFO:rf3.inference_engines.rf3:[rank: 0] Found 2 structures to predict!
INFO:rf3.inference_engines.rf3:[rank: 0] Predicting structure 1/2: seq0
WARNING:atomworks.ml:Cached data not found for ALA at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/A/ALA/ALA.pt
WARNING:atomworks.ml:Cached data not found for ARG at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/A/ARG/ARG.pt
WARNING:atomworks.ml:Cached data not found for ASN at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/A/ASN/ASN.pt
WARNING:atomworks.ml:Cached data not found for ASP at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/A/ASP/ASP.pt
WARNING:atomworks.ml:Cached data not found for ATP at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/A/ATP/ATP.pt
WARNING:atomworks.ml:Cached data not found for GLN at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/G/GLN/GLN.pt
WARNING:atomworks.ml:Cached data not found for GLU at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/G/GLU/GLU.pt
WARNING:atomworks.ml:Cached data not found for GLY at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/G/GLY/GLY.pt
WARNING:atomworks.ml:Cached data not found for HIS at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/H/HIS/HIS.pt
WARNING:atomworks.ml:Cached data not found for ILE at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/I/ILE/ILE.pt
WARNING:atomworks.ml:Cached data not found for LEU at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/L/LEU/LEU.pt
WARNING:atomworks.ml:Cached data not found for LYS at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/L/LYS/LYS.pt
WARNING:atomworks.ml:Cached data not found for MET at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/M/MET/MET.pt
WARNING:atomworks.ml:Cached data not found for PHE at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/P/PHE/PHE.pt
WARNING:atomworks.ml:Cached data not found for PRO at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/P/PRO/PRO.pt
WARNING:atomworks.ml:Cached data not found for SER at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/S/SER/SER.pt
WARNING:atomworks.ml:Cached data not found for THR at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/T/THR/THR.pt
WARNING:atomworks.ml:Cached data not found for TRP at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/T/TRP/TRP.pt
WARNING:atomworks.ml:Cached data not found for TYR at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/T/TYR/TYR.pt
WARNING:atomworks.ml:Cached data not found for VAL at /net/tukwila/lschaaf/datahub/MACE-OMOL-Jul2025/mace_embeddings/V/VAL/VAL.pt
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[/tmp/ipykernel_391/4213100572.py](https://localhost:8080/#) in <cell line: 0>()
      1 # Run RF3 prediction for all designed sequences
      2 # Set an output directory to save predicted strucutres into local file system
----> 3 rf3_all_outputs = inference_engine.run(inputs=rf3_inputs, out_dir="Predictions", annotate_b_factor_with_plddt=True)

50 frames
[/usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/compiler.py](https://localhost:8080/#) in make_llir(self, src, metadata, options, capability)
    339         if os.environ.get("TRITON_DISABLE_LINE_INFO", "0") == "0":
    340             passes.llvmir.add_di_scope(pm)
--> 341         pm.run(mod)
    342         # LLVM-IR (MLIR) -> LLVM-IR (LLVM)
    343         llvm.init_targets()

RuntimeError: PassManager::run failed 

After some googling it sounds like this is an internal GPU issue and not code related, but I find it odd that it happens every time I try to run RF3.

Is this a known bug, and do you have any advice for running RF3 on Collab?

Thanks, Ben

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the failure in the IPD Design Pipeline Colab on a T4, following the shown inference_engine.run(...) call; inspect the Triton compiler traceback and RF3/PyTorch environment versions. Done means determining whether this is a known compatibility issue and documenting a verified way to run RF3 on Colab, or confirming the details maintainers need to investigate it.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
cloud, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.