Deep learning for Korali

Korali is a framework for Bayesian uncertainty quantification, optimisation and reinforcement learning on HPC clusters. My master's thesis gave it neural networks.

Why not just call PyTorch

The obvious question, and the one that defines the problem.

Korali distributes work across cluster nodes with its own execution model — it owns the processes, the communication and the scheduling. A framework like PyTorch wants to own those things too. Running one inside the other means two runtimes each believing they control the same resources, which produces contention in the best case and deadlock in the worst.

So the neural network support had to live inside Korali's execution model rather than beside it: its own layers, its own forward and backward passes, dispatching to vendor libraries for the heavy linear algebra while Korali kept control of the parallelism.

Verification before performance

A deep-learning implementation that is fast and subtly wrong is worse than no implementation, because the error surfaces as a slightly worse model rather than a crash.

So the first milestone was not speed. It was gradient checking against reference implementations — forward passes matching PyTorch and TensorFlow within numerical tolerance, backward passes verified against both analytical gradients and finite differences, layer by layer, before any optimisation.

This is slow and unglamorous and it is the part that makes the rest trustworthy.

Backends

The implementation dispatches to:

• cuDNN for NVIDIA GPUs
• oneDNN for CPU execution
• fallback paths where neither applies

Benchmarking across these answers the question that actually matters to someone running a simulation: given my problem size and my cluster, which backend should I use? The answer is not uniform. Small networks with small batches spend more time moving data to the GPU than computing on it, and the CPU path wins outright. The crossover point is worth measuring rather than assuming — which is the general lesson, and the one I have used most since.

What I took from it

Framework boundaries are where the hard problems live. The individual layers are textbook. Making them cooperate with an execution model that was not designed for them is where the real work was.

Benchmarks are only useful with their conditions attached. "cuDNN is faster" is not a result. "cuDNN is faster above this batch size, on this hardware, for this layer configuration" is one.

Numerical verification is a feature. It is also the thing most likely to be skipped under time pressure, and the thing you most regret skipping.