Scaling Molecular Simulations: Our GPU Infrastructure Journey

When Gimerny AI started in 2021, our entire compute infrastructure consisted of a single DGX A100 workstation with 8 GPUs, sitting under a desk in our Berlin office. Today, we run workloads across 2,000+ GPUs spanning AWS, Google Cloud, and NVIDIA DGX Cloud. This post traces that journey and shares the architectural decisions, mistakes, and optimizations that got us here.

Phase 1: The Single-Node Era (8 GPUs)

Our first models were small enough to train on a single node. A GNN with 2 million parameters for molecular property prediction trains in hours on 8 A100s. At this stage, we used PyTorch's DistributedDataParallel (DDP) for multi-GPU training and kept things simple: one training job at a time, manual experiment tracking in spreadsheets, and results files stored on a shared NFS mount.

This worked until it did not. As our models grew (the GimernyGenome foundation model has 1.2 billion parameters) and our data expanded (4.2 petabytes of sequencing data), single-node training became untenable. A training run that should take 2 weeks on 64 GPUs would take 4 months on 8.

Phase 2: Multi-Node Training on AWS (64-256 GPUs)

We moved to AWS p4d instances (8x A100 per node) and built our first multi-node training infrastructure. Key components included NVIDIA NCCL for inter-node GPU communication over 400Gbps Elastic Fabric Adapter (EFA) networking, PyTorch FSDP (Fully Sharded Data Parallel) for sharding model parameters, gradients, and optimizer states across GPUs, and a custom job scheduler built on Kubernetes that managed GPU allocation across training, inference, and simulation workloads.

The transition from single-node to multi-node surfaced several challenges. Communication overhead became significant: NCCL all-reduce operations on 256 GPUs consumed 30% of wall-clock time for our smaller models. We mitigated this with gradient accumulation (reducing communication frequency), gradient compression (top-k sparsification), and overlap of communication with computation.

Fault tolerance was essential at this scale. With 32 nodes, the probability of at least one node failing during a multi-day training run is substantial. We implemented automatic checkpointing every 30 minutes and a restart protocol that detects failed nodes, redistributes the workload, and resumes from the latest checkpoint without human intervention.

Phase 3: Heterogeneous Workloads (256-2,000 GPUs)

Our workloads are heterogeneous. Model training requires large blocks of tightly-coupled GPUs with fast interconnects. Molecular generation (running the trained diffusion model to produce molecules) is embarrassingly parallel: each molecule is generated independently. Molecular dynamics simulations have their own parallelization patterns, with spatial decomposition across GPUs.

We built a workload-aware scheduler that categorizes jobs into three tiers. Training jobs receive dedicated clusters with EFA networking and topology-aware placement (ensuring GPUs that communicate most are on the same switch). Generation jobs run on spot instances at 60-70% cost savings, with automatic retry on preemption. Simulation jobs use a custom spatial decomposition that maps molecular system partitions to GPU topology.

This heterogeneous approach improved our cost efficiency by 2.3x compared to our Phase 2 architecture, which treated all workloads identically.

Data Pipeline Optimization

GPU utilization is only as good as the data pipeline feeding it. For our multi-omics foundation model, the training data (4.2PB of sequencing files) vastly exceeds GPU memory and even local SSD capacity. We built a streaming data pipeline using a custom PyTorch DataLoader that prefetches data from S3 using multiple threads, decodes and tokenizes sequencing files on-the-fly using CPU workers, maintains a shuffle buffer in CPU memory for pseudo-random sampling, and transfers batches to GPU asynchronously via pinned memory and CUDA streams.

This pipeline sustains 95% GPU utilization during training, compared to 60% with a naive approach of downloading data to local storage before training.

Current Architecture

Our current platform runs across three cloud providers (AWS, GCP, and NVIDIA DGX Cloud) with a unified control plane. We use approximately 1,200 NVIDIA A100 GPUs (80GB) for training and generation, 600 NVIDIA H100 GPUs for our largest models and newest workloads, and 200 NVIDIA L40S GPUs for inference serving and demos.

Total monthly compute spend is approximately $1.8M, of which 40% is offset by reserved instances and committed use discounts. Spot and preemptible instances save an additional 25%.

Lessons for Other Startups

Three pieces of advice for startups scaling GPU infrastructure. Start with managed services. We wasted months building custom Kubernetes operators that cloud-managed solutions (SageMaker, Vertex AI, DGX Cloud) now provide out of the box. Invest in observability early. At 2,000 GPUs, a single misconfigured node can silently degrade training performance by 10%. We instrument every GPU with DCGM metrics (utilization, memory bandwidth, NVLink throughput, thermal throttling) and alert on anomalies. Negotiate aggressively. Cloud GPU pricing is highly negotiable for multi-year commitments. Our effective per-GPU-hour rate is 40% below list price.

The democratization of AI infrastructure means that startups can now access supercomputer-class resources that were previously available only to tech giants. The challenge is not access but efficient utilization.