Kernel Crash (Exit Code 3221225477) when loading pLDDT model

[tangyuhao8]

Bug Report

Describe the bug
I am encountering a fatal kernel crash when running the sample.ipynb notebook. The issue occurs specifically during the model initialization phase, on the line that calls model_wrapper.from_pretrained_plddt_model() with the simplefold_1.6B model selected.

The kernel dies with the Windows exit code 3221225477. Notably, this code block runs successfully without any errors when simplefold_model is set tosimplefold_100M.

To Reproduce
Steps to reproduce the behavior:

1. Set up the conda environment as per the repository's instructions.
2. Open the sample.ipynb notebook in VS Code.
3. Configure the model parameters in the notebook cell as follows:

   simplefold_model = "simplefold_1.6B"
   ckpt_dir = "./artifacts"
   plddt = True  # The crash occurs when this is True
   backend = "torch"

4. Execute the cells in order.
5. The kernel crashes when executing the following block of code, specifically on the plddt_model = ... line:

   # initialize the folding model and pLDDT model
   model_wrapper = ModelWrapper(
       simplefold_model=simplefold_model,
       ckpt_dir=ckpt_dir,
       plddt=plddt,
       backend=backend,
   )
   device = model_wrapper.device
   folding_model = model_wrapper.from_pretrained_folding_model()
   plddt_model = model_wrapper.from_pretrained_plddt_model() # <-- Crash happens here

Expected behavior
The plddt_model should be loaded successfully without the kernel process terminating.

Actual behavior
The Jupyter kernel process dies immediately. The VS Code Jupyter log shows the following error:

[error] Disposing session as kernel process died ExitCode: 3221225477, Reason: 

System Environment (please complete the following information):

• OS: Windows 11
• Python Version: 3.10.18
• PyTorch Version:
  pytorch-lightning 2.5.5
  torch 2.8.0
  torchaudio 2.8.0
  torchmetrics 1.8.2
  torchvision 0.23.0
• GPU: CPU-only

Any help or insight into what might be causing this memory access violation would be greatly appreciated. Thank you!

[tangyuhao8]

Title: (Update) Kernel crashes at multiple memory-intensive steps

I have a further update on this issue. The kernel crash is not isolated to the plddt_model loading step.

After successfully loading the simplefold_1.6B model (by setting plddt=False to bypass the initial crash), I encounter the exact same kernel crash (ExitCode: 3221225477) at a later stage in the sample.ipynb notebook.

The crash now occurs when initializing the InferenceWrapper:

# This cell now causes the identical kernel crash
inference_wrapper = InferenceWrapper(
    output_dir=output_dir,
    prediction_dir=prediction_dir,
    num_steps=num_steps,
    tau=tau,
    nsample_per_protein=nsample_per_protein,
    device=device,
    backend=backend
)

This suggests the issue is not specific to a single function, but is a more general problem related to resource allocation or library conflicts (e.g., PyTorch/system drivers) when handling large models in my specific environment.

[yuyangw]

Hi, thanks for the interest in our work. This kernel crash is likely due to the out of memory when loading large model. Can you let us know what is the RAM memory of your machine? And also can you try loading smaller model (e.g., 360M/700M) without pLDDT module to see if the same issue happens?

[tangyuhao8]

Hi, thank you for the quick and helpful reply!

Regarding the RAM on my machine, my system settings show 16.0 GB of installed RAM (with 13.7 GB available). I think this is the memory you were referring to.

I followed your suggestion and tested simplefold_100M(while all other models like 360M, 700M will encounter Kernel Crash) with the pLDDT module disabled (plddt=False). I'm happy to report that the code block runs successfully and produces results without any kernel crash. This confirms your hypothesis that the issue is indeed related to an out-of-memory problem when loading the larger models on my system.

On a final note, I'd like to express my sincere appreciation for your work. Being able to run a protein structure prediction model locally on a consumer-grade PC with low configuration is a fantastic achievement and makes this technology much more accessible.

Thanks again for your help and for this amazing project!

[yuyangw]

Hi, thanks for the feedback. We'll update the document to include minimal hardware requirement to run models of different sizes. Meanwhile, we'll work on potential implementations to reduce maximal memory needed for running large SimpleFold models.