High-Performance Computing Guide

This guide covers running Auto3D on HPC clusters, multi-GPU systems, and optimizing performance for large-scale workflows.

Multi-GPU Configuration

Basic Multi-GPU Usage

Specify multiple GPUs with gpu_idx:

from Auto3D import Auto3DOptions, main

config = Auto3DOptions(
    path="large_dataset.smi",
    k=1,
    use_gpu=True,
    gpu_idx=[0, 1, 2, 3],  # Use 4 GPUs
)
output = main(config)

Via CLI:

auto3d run dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3"

How Multi-GPU Works

Auto3D distributes molecules across GPUs:

Dataset is chunked based on available memory
Each GPU processes its assigned chunk
Results are merged into a single output file

For N GPUs, you get approximately N× throughput.

Selecting Specific GPUs

On shared systems, select available GPUs:

# Use only GPUs 2 and 3 (avoiding 0 and 1)
config = Auto3DOptions(
    path="input.smi",
    k=1,
    use_gpu=True,
    gpu_idx=[2, 3],
)

Or use environment variables:

export CUDA_VISIBLE_DEVICES=2,3
auto3d run input.smi --k=1 --gpu

Memory Management

Controlling Memory Usage

For large datasets, control memory allocation:

config = Auto3DOptions(
    path="huge_dataset.smi",   # 100K+ molecules
    k=1,
    memory=64,                  # Assign 64GB RAM
    capacity=50,                # Molecules per GB
    batchsize_atoms=1024,       # Atoms per batch per GB
)

Reducing GPU Memory

For limited GPU memory:

config = Auto3DOptions(
    path="input.smi",
    k=1,
    batchsize_atoms=512,        # Smaller batches
    use_gpu=True,
)

Or use single model instead of ensemble:

export AUTO3D_USE_ENSEMBLE=0
auto3d run input.smi --k=1 --gpu

Chunked Processing

Auto3D automatically chunks large datasets:

config = Auto3DOptions(
    path="million_molecules.smi",
    k=1,
    memory=128,      # Available RAM in GB
    capacity=42,     # Default: 42 molecules per GB
)

Each chunk is processed independently and results are merged.

SLURM Job Scripts

Basic SLURM Script

#!/bin/bash
#SBATCH --job-name=auto3d
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --gpus=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=32G
#SBATCH --time=04:00:00
#SBATCH --output=auto3d_%j.out
#SBATCH --error=auto3d_%j.err

# Load modules
module load cuda/11.8
module load anaconda3

# Activate environment
conda activate auto3d

# Run Auto3D
auto3d run molecules.smi --k=5 --gpu --engine=AIMNET

Multi-GPU SLURM Script

#!/bin/bash
#SBATCH --job-name=auto3d_multi
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --gpus=4
#SBATCH --cpus-per-task=16
#SBATCH --mem=128G
#SBATCH --time=12:00:00
#SBATCH --output=auto3d_%j.out

module load cuda/11.8
module load anaconda3
conda activate auto3d

# Use all 4 GPUs
auto3d run large_dataset.smi --k=1 --gpu --gpu-idx="0,1,2,3"

Array Jobs for Multiple Files

#!/bin/bash
#SBATCH --job-name=auto3d_array
#SBATCH --partition=gpu
#SBATCH --array=1-10
#SBATCH --gpus=1
#SBATCH --cpus-per-task=4
#SBATCH --mem=32G
#SBATCH --time=02:00:00

module load cuda/11.8
conda activate auto3d

# Process different input files
INPUT_FILE="batch_${SLURM_ARRAY_TASK_ID}.smi"
auto3d run "$INPUT_FILE" --k=1 --gpu

PBS/Torque Script

#!/bin/bash
#PBS -N auto3d
#PBS -l nodes=1:ppn=4:gpus=1
#PBS -l mem=32gb
#PBS -l walltime=04:00:00
#PBS -q gpu

cd $PBS_O_WORKDIR
module load cuda/11.8
source activate auto3d

auto3d run molecules.smi --k=5 --gpu

Performance Optimization

TF32 Acceleration

Enable TensorFloat-32 on Ampere+ GPUs:

config = Auto3DOptions(
    path="input.smi",
    k=1,
    allow_tf32=True,  # ~1.5x faster
    use_gpu=True,
)

torch.compile Optimization

For PyTorch 2.0+, enable compilation:

export AUTO3D_COMPILE_MODEL=1
auto3d run input.smi --k=1 --gpu --engine=ANI2x

Provides ~1.25x speedup after warmup.

Engine Selection for Speed

From fastest to slowest:

ANI2xt: Ultra-fast, good for screening
ANI2x: Very fast, well-validated
AIMNET single: Fast, most versatile (default)
AIMNET ensemble: Slower, highest accuracy

# Fastest for screening
auto3d run input.smi --k=1 --gpu --engine=ANI2xt

# Best accuracy
export AUTO3D_USE_ENSEMBLE=1
auto3d run input.smi --k=1 --gpu --engine=AIMNET

Optimal Batch Sizes

Tune batch size for your GPU:

GPU Memory	Recommended batchsize_atoms	Notes
8 GB	512	RTX 3070, etc.
16 GB	1024	RTX 3080, V100
24 GB	1536	RTX 3090, A5000
40+ GB	2048	A100, H100

Benchmarking

Measuring Throughput

import time
from Auto3D import Auto3DOptions, main

# Prepare test dataset
start = time.time()

config = Auto3DOptions(
    path="benchmark_1000.smi",
    k=1,
    use_gpu=True,
    verbose=False,
)
output = main(config)

elapsed = time.time() - start
molecules_per_second = 1000 / elapsed
print(f"Throughput: {molecules_per_second:.1f} molecules/second")

Comparing Configurations

import time
from Auto3D import Auto3DOptions, main

configs = [
    ("ANI2xt-single", {"optimizing_engine": "ANI2xt"}),
    ("AIMNET-single", {"optimizing_engine": "AIMNET"}),
    ("ANI2x-compile", {"optimizing_engine": "ANI2x"}),
]

for name, kwargs in configs:
    start = time.time()
    config = Auto3DOptions(
        path="benchmark.smi",
        k=1,
        use_gpu=True,
        **kwargs
    )
    main(config)
    elapsed = time.time() - start
    print(f"{name}: {elapsed:.1f}s")

Distributed Computing

Processing Across Nodes

For truly massive datasets, split across nodes:

#!/bin/bash
# split_and_run.sh

# Split input into chunks
split -l 10000 huge_dataset.smi chunk_

# Submit job for each chunk
for chunk in chunk_*; do
    sbatch --export=INPUT="$chunk" single_node.slurm
done

With single_node.slurm:

#!/bin/bash
#SBATCH --gpus=4
#SBATCH --mem=128G

auto3d run "$INPUT" --k=1 --gpu --gpu-idx="0,1,2,3"

Merging Results

After parallel processing, merge outputs:

from rdkit import Chem
from pathlib import Path

# Find all output files
output_files = list(Path(".").glob("chunk_*_out.sdf"))

# Merge
writer = Chem.SDWriter("merged_output.sdf")
for sdf_file in sorted(output_files):
    for mol in Chem.SDMolSupplier(str(sdf_file)):
        if mol is not None:
            writer.write(mol)
writer.close()

print(f"Merged {len(output_files)} files")

Troubleshooting

CUDA Out of Memory

Reduce batch size:

config = Auto3DOptions(batchsize_atoms=256, ...)

Use fewer GPUs with more memory each
Process in smaller chunks

Job Timeout

Increase time limit
Use faster engine (ANI2xt)
Reduce k or use window
Split into array jobs

Slow I/O

On shared filesystems:

Copy input to local scratch:

cp $SLURM_SUBMIT_DIR/input.smi $TMPDIR/
cd $TMPDIR
auto3d run input.smi --k=1 --gpu
cp *_out.sdf $SLURM_SUBMIT_DIR/

Use SSD scratch if available

Node Failures

Make jobs resumable:

#!/bin/bash
#SBATCH --requeue

# Check if output exists
if [ -f "output.sdf" ]; then
    echo "Output exists, skipping"
    exit 0
fi

auto3d run input.smi --k=1 --gpu

Monitoring

Track GPU Usage

# In a separate terminal
watch -n 1 nvidia-smi

# Or log to file
nvidia-smi --query-gpu=timestamp,utilization.gpu,memory.used --format=csv -l 5 > gpu_log.csv &

Memory Profiling

import torch

# At start
torch.cuda.reset_peak_memory_stats()

# After run
peak_memory = torch.cuda.max_memory_allocated() / 1e9
print(f"Peak GPU memory: {peak_memory:.2f} GB")

Best Practices

Test locally first: Validate on small dataset before HPC submission
Request appropriate resources: Don’t over-request GPUs/memory
Use scratch storage: Avoid slow shared filesystems for I/O
Enable checkpointing: For long jobs, process in resumable chunks
Monitor GPU utilization: Aim for >80% GPU usage
Log everything: Capture stdout/stderr for debugging
Clean up: Remove temporary files after successful runs