When the Batch Is Already One: ShardTensor and the Cost of Dividing the World

There is so much fluff out there telling us that we are nearing the limits of our current technology. I would argue that couldn’t be further from the truth. As an optimist about this technology, and someone very much rooting for the singularity, I think we have barely scratched the surface of this current renaissance. The more I learn, the more I find myself returning to one thing that seems to sit underneath so many of the problems we are trying to solve: communication.

As I continue trying to better understand parallel programming, specifically across GPUs, CPUs, and machine learning, I’ve become increasingly interested in what happens when a problem becomes too large to fit neatly inside one piece of hardware. More compute helps, but once work is distributed, the problem changes. Now the work has to be divided mindfully, the different pieces need enough information about one another to remain useful, and whatever we gain through parallelism cannot be completely swallowed by the cost of moving information around.

That brought me to ShardTensor: Domain Parallelism for Scientific Machine Learning.

Scientific machine learning has an continuous scaling problem. In many workloads, the model itself is not necessarily what overwhelms GPU memory. Weather models, fluid simulations, volumetric imaging, meshes, and other scientific workloads can operate over enormous spatial domains, and increasing resolution can cause memory requirements to grow incredibly quickly.

Traditional data parallelism helps by distributing different samples across different GPUs. GPU 0 might work on one sample while GPU 1 works on another. But eventually you reach a fairly unforgiving boundary: what happens when a single sample no longer fits onto one GPU?

You can’t divide the batch any further when the batch is already one.

Instead of distributing different samples across devices, the domain itself is divided. One large scientific input can be partitioned across several GPUs, with each device operating on only a portion of the overall tensor. ShardTensor brings this idea into PyTorch by extending distributed tensor abstractions to better support the kinds of spatial decomposition required by scientific ML.

The papers underlying idea is not some entirely new invention created for AI. Domain decomposition has existed in scientific computing for decades. Scientific ML is effectively returning to an old parallel-computing problem and asking how we can make that technique accessible inside modern machine-learning frameworks.

But dividing the domain does not make the rest of the domain disappear. A GPU working on one region may eventually need information belonging to another region. Certain operations can remain local, while others require communication across device boundaries. Suddenly memory capacity, communication bandwidth, synchronization, topology, and the placement of data all become part of the computation.

Parallelism solves one problem by creating another. We divide work because one processor, one GPU, or one memory space cannot efficiently carry the whole problem, but every division creates new relationships between the pieces. At some point, the quality of those relationships becomes just as important as the amount of compute available.

That also feels strangely connected to what I have been exploring recently with multi-agent systems. I don’t mean that GPU shards and autonomous agents are equivalent, they obviously are not. One is explicitly programmed distributed computation, while the other can involve agents making independent decisions under incomplete information. But both force us to confront a similar systems question: once something has been divided into independently operating parts, what information must move between those parts for the larger system to remain coherent?

Distributing a domain allows workloads to move beyond the memory limit of a single GPU, but communication itself is not free. For smaller workloads, the overhead can outweigh the benefit. For sufficiently large workloads, however, distributing the domain becomes the thing that makes the computation possible in the first place.

That distinction is something I want to understand much better. We talk constantly about having more compute, larger models, more GPUs, and more agents, but increasingly I’m becoming interested in a different question: how effectively can we coordinate the resources we already have?

Maybe the next major gains in AI will not come exclusively from making every individual component larger or faster. Some of them may come from getting much better at deciding where computation should happen, what information should move, and when independently operating pieces of a system actually need one another.