OpenDDE / docs /inference_instructions.md
aurekaresearch's picture
Upload 8 files
373e2bb verified
|
Raw
History Blame Contribute Delete
8.89 kB

Inference Instructions

Concise reference for installing OpenDDE, preparing runtime data, and running opendde commands. For Docker, see docker_installation.md.

Install

OpenDDE requires Python >=3.11. Install uv if needed:

curl -LsSf https://astral.sh/uv/install.sh | sh
source "$HOME/.local/bin/env"

GPU install:

uv venv --python 3.11
source .venv/bin/activate
uv pip install --torch-backend cu126 'opendde[gpu]'
opendde doctor

CPU install:

uv venv --python 3.11
source .venv/bin/activate
uv pip install --torch-backend cpu 'opendde[cpu]'
opendde doctor

Source install from a checkout:

uv venv --python 3.11
source .venv/bin/activate
uv pip install --torch-backend cu126 -e '.[gpu]'
opendde doctor

Runtime data

Set OPENDDE_ROOT_DIR to the directory that stores checkpoints and runtime data:

$OPENDDE_ROOT_DIR/
β”œβ”€β”€ checkpoint/opendde.pt
β”œβ”€β”€ common/
└── search_database/        # needed for local template/RNA-MSA search

Prepare data from a source checkout:

export OPENDDE_ROOT_DIR=/path/to/opendde_data
bash scripts/download_opendde_data.sh

For a protein-only prediction that disables MSA, template, and RNA-MSA features, search databases are not needed:

bash scripts/download_opendde_data.sh --skip-search-database

If you already have a checkpoint:

mkdir -p "$OPENDDE_ROOT_DIR/checkpoint"
cp /path/to/opendde.pt "$OPENDDE_ROOT_DIR/checkpoint/opendde.pt"

Alternatively, keep the checkpoint anywhere readable and pass it directly with opendde pred --load_checkpoint_path /path/to/opendde_abag.pt.

Released checkpoints:

Checkpoint Use case Download
opendde.pt General-purpose OpenDDE checkpoint. opendde.pt
opendde_abag.pt ABAG-optimized checkpoint for antibody-antigen complexes. opendde_abag.pt

Use opendde.pt with -n opendde_v1 as the default general-purpose checkpoint. To use the ABAG-optimized checkpoint, keep it as opendde_abag.pt and pass it with --load_checkpoint_path, for example opendde pred --load_checkpoint_path "$OPENDDE_ROOT_DIR/checkpoint/opendde_abag.pt".

Concrete setup:

export OPENDDE_ROOT_DIR=/path/to/opendde_data
mkdir -p "$OPENDDE_ROOT_DIR/checkpoint"

# General-purpose checkpoint used by default with -n opendde_v1.
curl -L \
  -o "$OPENDDE_ROOT_DIR/checkpoint/opendde.pt" \
  https://huggingface.co/aurekaresearch/OpenDDE/resolve/main/opendde.pt

# ABAG-optimized checkpoint, selected with --load_checkpoint_path.
curl -L \
  -o "$OPENDDE_ROOT_DIR/checkpoint/opendde_abag.pt" \
  https://huggingface.co/aurekaresearch/OpenDDE/resolve/main/opendde_abag.pt

Then run general-purpose inference without an explicit checkpoint path. For ABAG inference, add:

--load_checkpoint_path "$OPENDDE_ROOT_DIR/checkpoint/opendde_abag.pt"

Useful environment variables:

Variable Purpose
OPENDDE_ROOT_DIR Checkpoints, common files, search databases. Defaults to ~/.cache/opendde.
OPENDDE_DEPENDENCY_URL Override checkpoint download root.
OPENDDE_COMMON_URL Override common runtime file download root. Falls back to OPENDDE_DEPENDENCY_URL when set.
OPENDDE_SEARCH_DATABASE_URL Override template/RNA-MSA database download root.
LAYERNORM_TYPE LayerNorm backend; defaults to torch. Set to fast_layernorm to opt into the fused kernel.

Template/RNA-MSA preprocessing also needs HMMER. Template inference may need kalign:

apt-get update && apt-get install -y hmmer kalign

Input JSON

OpenDDE input is a top-level list of jobs:

[
  {
    "name": "tiny",
    "modelSeeds": [101],
    "sequences": [
      {
        "proteinChain": {
          "sequence": "ACDEFGHIK",
          "count": 1
        }
      }
    ]
  }
]

covalent_bonds is optional and may be omitted from a job; include it only to declare explicit covalent links between entities.

Full schema: infer_json_format.md.

Convert a structure file to JSON:

opendde json -i examples/7pzb.pdb -o ./output --altloc first
opendde json -i examples/2lwu.cif -o ./output --altloc first --assembly_id 1

Preprocess optional features

# Protein MSA
opendde msa -i examples/input.json -o ./output

# Protein MSA + template search
opendde mt -i examples/input.json -o ./output

# Protein MSA + template search + RNA MSA when RNA is present
opendde prep -i examples/input.json -o ./output

Notes:

  • Protein MSA uses the public ColabFold MMseqs2 API unless A3M paths are already present in the JSON.
  • Template and RNA-MSA search use local databases under $OPENDDE_ROOT_DIR/search_database/.
  • Updated JSON files are written next to the input JSON.

Details: msa_template_pipeline.md.

Run prediction

Standard run:

opendde pred -i examples/input.json -o ./output -n opendde_v1

Compatibility run with the standard step/cycle counts:

opendde pred \
  -i examples/input.json \
  -o ./output \
  -n opendde_v1 \
  --use_msa false \
  --use_template false \
  --use_rna_msa false \
  --sample 1 \
  --step 200 \
  --cycle 10

Inference defaults to fp32 and auto triangle kernels (PyTorch on CPU, cuEquivariance on a CUDA GPU), so neither needs to be set explicitly.

4-GPU Fold-CP inference

Note: the current Fold-CP path is provided as a distributed-inference demo. It verifies that OpenDDE can execute the four-GPU context-parallel path, including MSA-enabled inputs, but memory capacity and runtime performance are still being actively optimized. We plan to continue improving this path in collaboration with NVIDIA, including integration and tuning for acceleration libraries such as cuEquivariance where applicable.

Fold-CP distributes token-pair-heavy inference work over four GPUs. Launch it with torchrun and expose exactly the GPUs you want to use:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --standalone --nproc_per_node 4 \
  -m runner.batch_inference pred \
  -i examples/protein_200.json \
  -o ./output_cp4 \
  -n opendde_v1 \
  --use_msa false \
  --use_template false \
  --use_rna_msa false \
  --sample 1 \
  --step 200 \
  --cycle 10 \
  --foldcp_mode distributed \
  --foldcp_size_dp 1 \
  --foldcp_size_cp 4

Runtime notes:

  • --nproc_per_node 4 must match --foldcp_size_dp 1 times --foldcp_size_cp 4.
  • --foldcp_size_cp 4 creates a 2 x 2 context-parallel mesh.
  • The same input, model, dtype, cycle, step, sample, MSA, template, and kernel settings should be used when comparing single-GPU and Fold-CP outputs.
  • Outputs are written under the requested -o/--out_dir just like normal inference.
  • Optional --foldcp_metrics_jsonl path/to/metrics.jsonl records Fold-CP timing and memory metrics.

For single-GPU inference, omit the Fold-CP flags or set --foldcp_mode single --foldcp_size_cp 1.

Use prepared features:

opendde pred -i examples/examples_with_template/example_9fm7.json \
  -o ./output -n opendde_v1 \
  --use_msa true --use_template true

opendde pred -i examples/examples_with_rna_msa/example_9gmw_2.json \
  -o ./output -n opendde_v1 \
  --use_rna_msa true

Optional TFG Guidance

OpenDDE includes default-off Training-Free Guidance (TFG) for protein-ligand runs. TFG refines each sampled trajectory with geometry potentials while keeping the requested --sample count unchanged.

opendde pred -i examples/input.json -o ./output -n opendde_v1 \
  --use_tfg_guidance true

Outputs are written to:

<out_dir>/<job_name>/seed_<seed>/predictions/

Common flags

Flag Meaning
-n, --model_name Model name. Currently opendde_v1.
--load_checkpoint_path Explicit checkpoint path.
--seeds Comma-separated seeds, e.g. 101,102. Overrides the job's modelSeeds; if unset, modelSeeds are used, or a random seed when both are absent.
--use_msa Use/generate protein MSA features.
--use_template Use/generate template features.
--use_rna_msa Use/generate RNA MSA features.
--use_tfg_guidance Enable Training-Free Guidance.
--foldcp_mode single or distributed; use distributed with torchrun for four-GPU Fold-CP inference.
--foldcp_size_cp Number of context-parallel ranks. Four-GPU Fold-CP uses 4.
--foldcp_metrics_jsonl Optional JSONL path for Fold-CP timing and memory metrics.
--dtype bf16, fp16, or fp32.
--trimul_kernel, --triatt_kernel auto, cuequivariance, or torch.

Run opendde <command> --help for the full option list.