OpenDDE-MLX

OpenDDE-MLX is a rewrite of OpenDDE, the open-source all-atom biomolecular co-folding model, on top of MLX. It runs the released checkpoint natively on Apple Silicon with unified memory, without PyTorch, CUDA, cuEquivariance, or any multi-device machinery.

• Apple Silicon only. Metal GPU through MLX; no fallbacks for other platforms.
• Single machine. The Fold-CP context-parallel path and all distributed code are gone.
• Same weights, same outputs. Every module is verified against the PyTorch implementation with numerical parity tests; the released .pt checkpoint is converted once to .safetensors.
• Faster, and far smaller. 1.5-2.1x the speed of upstream on the same GPU, 4-7x the CPU path, in up to 5.5x less memory. See Benchmarks.

[!IMPORTANT] OpenDDE-MLX is a preview. The model, CLI flags, and JSON fields follow the upstream 1.1.x release.

Installation

Apple Silicon macOS with CPython 3.11, 3.12 or 3.13. We recommend uv. Choose one of the following methods.

Install from PyPI

uv venv --python 3.12
uv pip install --python .venv opendde-mlx

Converting a released .pt checkpoint additionally needs PyTorch (CPU):

uv pip install --python .venv "opendde-mlx[convert]"

Install from source

git clone git@github.com:aurekaresearch/OpenDDE-MLX.git
cd OpenDDE-MLX
uv venv --python 3.12
uv pip install --python .venv -e .

After either installation, verify the environment with:

uv run --no-project --python .venv opendde-mlx doctor

Model and runtime data

Checkpoints and runtime assets live under OPENDDE_ROOT_DIR (default ~/.cache/opendde):

$OPENDDE_ROOT_DIR/
├── checkpoint/opendde.safetensors   # converted weights (or opendde.pt)
├── common/                          # CCD components (downloaded automatically)
└── search_database/                 # only for local template / RNA-MSA search

opendde-mlx pred downloads the CCD assets and the released opendde.pt on first use and converts it to opendde.safetensors (the conversion needs the convert extra). To convert an existing checkpoint explicitly:

opendde-mlx convert ~/.cache/opendde/checkpoint/opendde.pt

Released checkpoints (same as upstream):

Quick start

opendde-mlx pred \
  -i examples/example_without_msa.json \
  -o ./output \
  --use_msa false

With MSA (uses the public ColabFold MMseqs2 server), an explicit dtype and a lighter schedule:

opendde-mlx pred \
  -i examples/example.json \
  -o ./output \
  --dtype bf16 \
  --cycle 4 \
  --step 50 \
  --sample 2

Outputs are written to output/<name>/seed_<seed>/predictions/: <name>_sample_<rank>.cif, <name>_summary_confidence_sample_<rank>.json, and <name>_full_data_sample_<rank>.json.

The CLI mirrors upstream OpenDDE — same commands (pred, json, msa, mt, prep), same flags, same JSON formats — so its documentation applies here. Dotted config overrides work through the module entry point:

python -m runner.inference \
  --input_json_path examples/input.json \
  --model.N_cycle 4 \
  --sample_diffusion.N_step 50 \
  --dtype bf16

Coming from OpenDDE

OPENDDE_ROOT_DIR (default ~/.cache/opendde) has the same layout in both projects, so CCD assets and checkpoints are reused, and your inference JSONs, commands and output files carry over unchanged. Two things differ:

• Weights are read as .safetensors. Parameter names are unchanged, so it is a one-time format conversion of the checkpoint you already have (see Model and runtime data); PyTorch is not needed afterwards.
• The default dtype is bf16 instead of upstream's fp32, and the flags for features this port drops are gone: --device, --foldcp_*, --trimul_kernel, --triatt_kernel, --enable_tf32, --deterministic, --use_default_params and --use_tfg_guidance.

Benchmarks

Three single-chain protein jobs of 100, 200 and 300 residues, prefixes of the same chain so that length is the only variable. No MSA, no templates. Four recycles, 30 diffusion steps, one sample, seed 42.

Each cell is model forward time · peak memory. The GPU rows use PyTorch's MPS backend and MLX's Metal backend on the same integrated GPU (18 cores on the M3 Pro, 32 on the M1 Max); memory is what the framework reports for its own allocations, or peak process RSS for the CPU rows. Apple Silicon has unified memory, so both come out of the same pool as the RAM column.

[!NOTE] These are deliberately small runs. A default job (10 cycles, 200 steps, 5 samples) on a real complex peaks near 15 GB, and the runner raises the Metal wired limit so those buffers stay resident. Macs with less than 24 GB of unified memory are not recommended: the job will still run, but the machine will be under heavy memory pressure and unpleasant to use meanwhile.

Everything needed to reproduce this, including the inputs and a one-command script, is in examples/inference_benchmark.

Speed-up of the fastest MLX GPU dtype

Each side uses whichever dtype was faster on that machine: for MLX that is bf16 on the M3 Pro and fp16 on the M1 Max.

A full production run

The table above uses a reduced schedule to keep every cell comparable. This one is what you actually get from opendde-mlx pred with the shipped defaults: 10 recycles, 200 diffusion steps, 5 samples, MLX GPU fp16.

Results from other Macs are welcome. Run the script and open a pull request adding your JSON to examples/inference_benchmark/results/:

python examples/inference_benchmark/run_benchmark.py \
  --dtypes fp32,bf16,fp16 \
  --out examples/inference_benchmark/results/<your_chip>.json

Performance notes

• The Pairformer trunk runs in bf16 by default; the diffusion module and confidence head stay in fp32. Add --fp32_diffusion false / --fp32_confidence false to run them reduced too, or --dtype fp32 for a full-precision trunk.
• --dtype fp16 is about 10% faster than bf16 on M1/M2, whose GPUs have no native bfloat; on M3 and later bf16 is native and marginally faster. fp16 needs two guards against its narrow exponent range: the triangle multiplication normalises both factors before the sum over tokens (exact, because the following LayerNorm is scale-invariant), and attention mask biases are clamped to a magnitude that stays finite.
• Diffusion pair biases for all 24 transformer blocks and the step-invariant atom conditioning are computed once per rollout when they fit in infer_setting.diffusion_bias_cache_gb (6 GB).
• Pair updates (triangle multiplication / attention, transitions) are chunked and double-buffered automatically for large complexes so peak memory stays bounded; attention always hits the fused Metal SDPA kernel.

Input format

See docs/infer_json_format.md for the JSON schema (proteinChain, dnaSequence, rnaSequence, ligand, ion, covalent_bonds) and docs/msa_template_pipeline.md for MSA/template preparation.

Citation

Please cite the OpenDDE technical report: OpenDDE: An Open-Source Drug Design Engine.

License

Apache-2.0, same as upstream OpenDDE.