Skip to content

config.md

Latest commit

 

History

History
298 lines (232 loc) · 9.31 KB

config.md

File metadata and controls

298 lines (232 loc) · 9.31 KB
Raw

Configuration

The configuration is organized into four top-level sections:

Section Purpose
optimization Batch mode, convergence criteria, charge & multiplicity
model Model backend, name, SevenNet modal, checkpoints, tokens, D3 correction
conformer_generation Conformer count and random seed
technical Device selection, memory & autobatcher knobs, logging

Unknown fields are preserved. Always use a config file for CLI and API calls.

Full Example (JSON, Fairchem UMA)

{
  "optimization": {
    "batch_optimization_mode": "batch",
    "batch_optimizer": "fire",

    "charge": 0,
    "multiplicity": 1,

    "force_convergence_criterion": 5e-2,
    "energy_convergence_criterion": null
  },

  "model": {
    "model_type": "fairchem",
    "model_name": "uma-s-1p2",
    "model_modal": null,

    "model_path": null,
    "model_cache_dir": null,

    "huggingface_token": null,
    "huggingface_token_file": "/home/hf_secret",

    "d3_correction": false,
    "d3_functional": "PBE",
    "d3_damping": "BJ"
  },

  "conformer_generation": {
    "max_num_conformers": 20,
    "conformer_seed": 42
  },

  "technical": {
    "device": "cuda",
    "max_memory_padding": 0.95,
    "memory_scaling_factor": 1.75,

    "max_atoms_to_try": 100000,
    "steps_between_swaps": 1,

    "logging_level": "INFO"
  }
}

Minimal Example (JSON, ORB-v3)

{
  "model": {
    "model_type": "orb",
    "model_name": "orb_v3_direct_omol"
  },
  "technical": {
    "device": "cuda"
  }
}

Minimal Example (JSON, SevenNet)

For multi-modal checkpoints such as 7net-omni or 7net-mf-ompa, a model_modal fidelity must be supplied. Single-modal checkpoints (e.g. 7net-0) ignore it.

{
  "model": {
    "model_type": "sevennet",
    "model_name": "7net-omni",
    "model_modal": "omol25_high"
  },
  "technical": {
    "device": "cuda"
  }
}

D3 Dispersion Correction

DFT-D3(BJ) dispersion correction is supported for the ORB-v3, Fairchem/UMA, and SevenNet backends. Internally ORB uses orb-models' native D3SumModel and Fairchem uses torch-sim's D3DispersionModel (via SumModel in batch mode); both share the same nvalchemiops GPU kernel. SevenNet uses its own native D3 implementation (SevenNetD3Model for the batch path, SevenNetD3Calculator for the single-structure path).

{
  "model": {
    "model_type": "orb",
    "model_name": "orb_v3_direct_omol",

    "d3_correction": true,
    "d3_functional": "PBE",
    "d3_damping": "BJ"
  },
  "technical": {
    "device": "cuda"
  }
}
{
  "model": {
    "model_type": "fairchem",
    "model_name": "uma-s-1p2",

    "d3_correction": true,
    "d3_functional": "PBE",
    "d3_damping": "BJ"
  },
  "technical": {
    "device": "cuda"
  }
}

Full Example (YAML)

optimization:
  batch_optimization_mode: batch
  batch_optimizer: fire

  charge: 0
  multiplicity: 1

  force_convergence_criterion: 5.0e-2
  energy_convergence_criterion: null

model:
  model_type: fairchem
  model_name: uma-s-1p2
  model_modal: null

  model_path: null
  model_cache_dir: null

  huggingface_token: null
  huggingface_token_file: /home/hf_secret

  d3_correction: false
  d3_functional: PBE
  d3_damping: BJ

conformer_generation:
  max_num_conformers: 20
  conformer_seed: 42

technical:
  device: cuda
  max_memory_padding: 0.95
  memory_scaling_factor: 1.75

  max_atoms_to_try: 100000
  steps_between_swaps: 1

  logging_level: INFO

Parameter Reference

optimization

Parameter Default Description
batch_optimization_mode "batch" "sequential" (ASE per structure) or "batch" (torch-sim GPU-accelerated)
batch_optimizer "fire" Optimizer for both single and batch modes: "fire", "gradient_descent", "lbfgs", or "bfgs". In single-structure (ASE) mode, maps to ASE FIRE/BFGS/LBFGS ("gradient_descent" falls back to FIRE). Defaults to "fire" if invalid or unset
charge 0 Total charge of the system (inferred from SMILES, not overridden)
multiplicity 1 Spin multiplicity of the system
force_convergence_criterion 5e-2 Force convergence threshold (eV/A). Used for both single and batch modes
energy_convergence_criterion null Energy convergence threshold. Batch mode only; force takes precedence if both set

model

Parameter Default Description
model_type "fairchem" Backend: "fairchem" / "uma", "orb" / "orb-v3", or "sevennet" / "7net"
model_name "uma-s-1p2" Model identifier (see Available Models)
model_modal null SevenNet multi-modal fidelity selector. Valid values depend on the checkpoint (7net-omni: e.g. "omol25_high", "omol25_low", "spice", "qcml"; 7net-mf-ompa: "mpa", "omat24"). Required for multi-modal checkpoints; ignored by single-modal checkpoints and by other backends
model_path null Local checkpoint path (Fairchem and SevenNet; overrides model_name)
model_cache_dir null Directory for cached model downloads
huggingface_token null HuggingFace token for model access
huggingface_token_file null File path to read the HF token from
d3_correction false Enable D3 dispersion correction. Supported for the ORB, Fairchem/UMA, and SevenNet backends
d3_functional "PBE" DFT functional for D3 correction (ORB/Fairchem use the orb-models BJ-damping table; SevenNet uses its own native D3)
d3_damping "BJ" Damping scheme for D3 correction ("BJ" or "BJM")

conformer_generation

Parameter Default Description
max_num_conformers 20 Maximum number of conformers to generate from SMILES
conformer_seed 42 Random seed for conformer generation

technical

Parameter Default Description
device "cuda" "cpu", "cuda", or "cuda:N" (e.g. "cuda:0"). Falls back to cuda:0 if the requested index doesn't exist
max_memory_padding 0.95 Fraction of GPU memory the autobatcher is allowed to fill during calibration. Lower = more headroom, smaller batches
memory_scaling_factor 1.75 Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. Must be > 1
max_atoms_to_try 100000 Upper bound on the autobatcher's calibration probe size (atoms)
steps_between_swaps 1 Optimization steps between batch swaps in the in-flight autobatcher. 1 is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower
logging_level "INFO" Logging verbosity: "DEBUG", "INFO", "WARNING", "ERROR"

You can also control the GPU with the CUDA_VISIBLE_DEVICES environment variable.

Available Models

Fairchem UMA (model_type: "fairchem")

Model Name Description
uma-s-1p2 UMA small, version 1.2 (default)
uma-s-1p1 UMA small, version 1.1
uma-m-1p1 UMA medium, version 1.1

Fairchem UMA models require a HuggingFace token for download.

ORB-v3 (model_type: "orb")

Use the underscored form as model_name.

Model Name Description
orb_v3_direct_omol ORB-v3 direct, omol (recommended)
orb_v3_conservative_omol ORB-v3 conservative, omol
orb_v3_direct_20_omat ORB-v3 direct, 20-layer omat
orb_v3_direct_inf_omat ORB-v3 direct, inf omat
orb_v3_conservative_20_omat ORB-v3 conservative, 20-layer omat
orb_v3_conservative_inf_omat ORB-v3 conservative, inf omat
orb_v3_direct_20_mpa ORB-v3 direct, 20-layer mpa
orb_v3_direct_inf_mpa ORB-v3 direct, inf mpa
orb_v3_conservative_20_mpa ORB-v3 conservative, 20-layer mpa
orb_v3_conservative_inf_mpa ORB-v3 conservative, inf mpa

SevenNet (model_type: "sevennet")

Multi-modal checkpoints (7net-mf-ompa, 7net-omni, 7net-omni-i8, 7net-omni-i12) also require a model_modal fidelity whose valid values depend on the checkpoint (e.g. omol25_high, omol25_low, spice, qcml for 7net-omni; mpa, omat24 for 7net-mf-ompa). A local checkpoint may instead be supplied via model_path.

Model Name Description
7net-0 SevenNet-0 (single-modal)
7net-0_22may2024 SevenNet-0, 22 May 2024 checkpoint
7net-l3i5 SevenNet l3i5
7net-mf-0 SevenNet multi-fidelity 0
7net-mf-ompa SevenNet multi-fidelity OMPA (multi-modal)
7net-omat SevenNet OMat
7net-omni SevenNet Omni (multi-modal)
7net-omni-i8 SevenNet Omni i8 (multi-modal)
7net-omni-i12 SevenNet Omni i12 (multi-modal)

D3 dispersion correction can be enabled for any ORB, Fairchem/UMA, or SevenNet model with d3_correction: true.

These model names are also available programmatically as gpuma.AVAILABLE_FAIRCHEM_MODELS, gpuma.AVAILABLE_ORB_MODELS, and gpuma.AVAILABLE_SEVENNET_MODELS.

See the examples/ folder for:

  • Single optimization (example_single_optimization.py)
  • Ensemble / batch optimization (example_ensemble_optimization.py)
  • Dispersion (D3) correction with Fairchem and ORB (example_dispersion.py)
  • Full multi-axis benchmark across models / optimizers / convergence (full_benchmark.py)

The example configs are sanitized (no tokens in plain text).