Over the last seven weeks we have built models, written training loops, composed optimizers, bulletproofed our code with Chex, and learned to debug inside JIT-compiled functions. Every one of those articles used dummy data or a small built-in dataset, generated on the fly with a few lines of NumPy or tfds.load().
Real projects do not work that way. Real datasets are millions of images sitting in cloud storage, gigabytes of audio clips streamed from disk, or CSV files too large to load into memory at once. And here is the uncomfortable truth: it does not matter how fast your @nnx.jit training step is if your data loader cannot keep up with it.
This is one of the most common and least glamorous bottlenecks in machine learning engineering. You buy an expensive GPU or rent TPU time, write a beautifully optimized training loop, and then watch your utilization graph sit at 20 percent because the accelerator is idle, waiting for the next batch of data to arrive from a slow Python loop.
Today we fix that.
We are introducing Grain, a data loading library purpose built for JAX. Grain plays a role similar to torch.utils.data.DataLoader in PyTorch, but its design accounts for the things JAX cares about most: determinism, explicit randomness, and staying out of the way of your compiled training step.
By the end of this article you will understand:
Why naive Python data loading becomes the bottleneck in JAX training
The three building blocks of a Grain pipeline: DataSource, Sampler, and Transformations
How to write deterministic and reproducible data augmentation
How to parallelize data loading across multiple CPU workers
How to shard data correctly for distributed training
How to wire a Grain pipeline into the training loop we built in Week 4
Let's feed the beast.
Why data loading becomes the bottleneck
Picture a training loop where the model itself runs in under a millisecond per batch, thanks to @nnx.jit and XLA compilation. Now picture a plain Python function that reads an image from disk, decodes it, resizes it, and normalizes it. That single operation, done in pure Python, can easily take several milliseconds.
If your data loading is a simple for loop over a list of file paths, your training loop looks like this:
for i in range(num_steps):
# This line might take 5-10ms per image, done sequentially
batch = load_and_preprocess_batch(file_paths[i])
# This line might take 0.5ms thanks to JIT
loss = train_step(model, optimizer, batch)
The GPU finishes its work almost instantly and then sits idle while the CPU slowly prepares the next batch. You paid for a race car and you are driving it in first gear.
The fix is not to make the model slower to match the data loader. The fix is to make data loading fast, parallel, and asynchronous, so that while the GPU is busy computing batch N, the CPU is already preparing batch N+1 in the background.
That is exactly what Grain is built to do.
The three building blocks of Grain
A Grain pipeline is assembled from three components that work together: a DataSource that knows how to read individual records, a Sampler that decides the order records are visited in, and a list of Transformations that process each record or batch.
import grain.python as grain
data_loader = grain.DataLoader(
data_source=my_source,
sampler=my_sampler,
operations=my_transformations,
worker_count=4,
)
Let's build each piece.
The DataSource
A DataSource is the simplest possible abstraction: something with a length and a way to fetch a single item by index. It does not know anything about batching, shuffling, or augmentation. Its only job is: given an index, return the raw record.
import grain.python as grain
class MyImageSource(grain.RandomAccessDataSource):
def __init__(self, file_paths, labels):
self._file_paths = file_paths
self._labels = labels
def __len__(self):
return len(self._file_paths)
def __getitem__(self, index):
# Return raw, unprocessed data. Keep this fast and simple.
return {
'path': self._file_paths[index],
'label': self._labels[index],
}
source = MyImageSource(file_paths=my_paths, labels=my_labels)
Notice how minimal this is. It does not decode images or apply any transformations. That separation of concerns is intentional and it is what allows Grain to parallelize the expensive work later.
The Sampler
The Sampler decides which index to visit, in what order, for how many epochs, and with what random seed. This is where shuffling and reproducibility live.
sampler = grain.IndexSampler(
num_records=len(source),
shuffle=True,
num_epochs=None, # None means infinite, useful for continuous training
seed=42,
)
Because the seed is explicit, two runs with the same seed visit records in exactly the same order. This matters enormously when you are debugging a training run or trying to reproduce a result from a paper or a colleague's experiment.
Transformations are the actual processing steps applied to each record: decoding an image, normalizing pixel values, applying augmentation, or grouping records into batches.
There are two kinds of custom transforms in Grain, and the distinction matters.
Deterministic transforms inherit from MapTransform and implement a map method. Given the same input, they always produce the same output.
import numpy as np
class Normalize(grain.MapTransform):
def map(self, element):
element['image'] = element['image'].astype(np.float32) / 255.0
return element
Random transforms inherit from RandomMapTransform and implement random_map, which receives a numpy.random.Generator as an argument. This is the critical detail: you must use the provided generator, not np.random directly, or your augmentation will not be reproducible.
class RandomFlip(grain.RandomMapTransform):
def random_map(self, element, rng):
if rng.random() > 0.5:
element['image'] = np.fliplr(element['image'])
return element
You then assemble your transforms into an ordered list, finishing with grain.Batch to group individual records into batches:
operations = [
Normalize(),
RandomFlip(),
grain.Batch(batch_size=64, drop_remainder=True),
]
Notice drop_remainder=True. If you have been following along since Week 6, you will remember exactly why this matters: an inconsistent final batch size forces JAX to recompile your JIT function. Dropping the incomplete final batch keeps every shape identical and avoids the recompilation trap entirely.
Putting it together
Here is a complete pipeline combining all three pieces:
import grain.python as grain
import numpy as np
# DataSource: raw access to records
class MNISTSource(grain.RandomAccessDataSource):
def __init__(self, images, labels):
self._images = images
self._labels = labels
def __len__(self):
return len(self._images)
def __getitem__(self, index):
return {
'image': self._images[index],
'label': self._labels[index],
}
source = MNISTSource(images=train_images, labels=train_labels)
# Sampler: order and reproducibility
sampler = grain.IndexSampler(
num_records=len(source),
shuffle=True,
num_epochs=None,
seed=42,
)
# Transformations: processing pipeline
class Normalize(grain.MapTransform):
def map(self, element):
element['image'] = element['image'].astype(np.float32) / 255.0
return element
operations = [
Normalize(),
grain.Batch(batch_size=64, drop_remainder=True),
]
# Assemble the DataLoader
data_loader = grain.DataLoader(
data_source=source,
sampler=sampler,
operations=operations,
worker_count=4,
)
# Use it as an iterator
data_iterator = iter(data_loader)
first_batch = next(data_iterator)
print(f"Batch image shape: {first_batch['image'].shape}")
print(f"Batch label shape: {first_batch['label'].shape}")
This looks similar to a PyTorch DataLoader, and that is intentional. The difference is under the hood: Grain's sampler and random transforms are built around the same explicit, seedable randomness philosophy that JAX itself uses.
Parallelism: The worker_count parameter
The single most impactful performance knob in Grain is worker_count.
data_loader = grain.DataLoader(
data_source=source,
sampler=sampler,
operations=operations,
worker_count=0, # Sequential, single process
)
With worker_count=0, everything runs sequentially in your main Python process. This is useful for debugging, since you can set breakpoints inside your transforms and step through them normally, but it will not use your CPU efficiently.
data_loader = grain.DataLoader(
data_source=source,
sampler=sampler,
operations=operations,
worker_count=8, # Parallel, using multiprocessing
)
With worker_count > 0, Grain spins up multiple worker processes. Because these are separate processes rather than threads, they bypass Python's Global Interpreter Lock entirely, meaning your image decoding, resizing, and augmentation can genuinely run in parallel across CPU cores. Grain uses shared memory to transfer completed batches back to your main process efficiently, avoiding the overhead of pickling large arrays repeatedly.
A practical rule of thumb: start with worker_count roughly equal to the number of CPU cores available, then benchmark. If your data source is already in memory (for example, a NumPy array you loaded once at the start), you may not need many workers at all, since there is no disk I/O to overlap. In that case:
data_loader = grain.DataLoader(
data_source=source,
sampler=sampler,
operations=operations,
worker_count=0,
read_options=grain.ReadOptions(num_threads=0),
)
Disabling thread prefetching when your dataset already lives in memory avoids unnecessary overhead from a prefetching mechanism you do not need.
Sharding for distributed training
When you train across multiple JAX processes, whether that is multiple GPUs on one machine or multiple machines in a cluster, each process must see a distinct, non overlapping slice of the data. Otherwise you are wasting compute recomputing gradients on the same examples multiple times per step.
Grain handles this with ShardOptions, passed into your sampler:
try:
# Automatically detects jax.process_index() and jax.process_count()
shard_options = grain.ShardByJaxProcess(drop_remainder=True)
except ImportError:
# Fallback for a single process setup
shard_options = grain.ShardOptions(
shard_index=0, shard_count=1, drop_remainder=True
)
sampler = grain.IndexSampler(
num_records=len(source),
shard_options=shard_options,
shuffle=True,
num_epochs=None,
seed=42,
)
ShardByJaxProcess is the recommended default. It reads directly from the JAX runtime to figure out how many processes are running and which one this particular process is, then automatically hands each process its own disjoint slice of the dataset. You do not need to manually calculate offsets or worry about duplicated examples across workers.
We will return to this in much more depth in Week 10, when we cover sharding and parallelism for the model itself. For now, the important thing to remember is that data sharding and model sharding are two separate concerns, and Grain handles the data side cleanly.
Rebuilding a training loop with Grain
Let's take the training loop from week 4 and swap the dummy NumPy data generation for a real Grain pipeline.
The data pipeline
import grain.python as grain
import numpy as np
import jax
import jax.numpy as jnp
from flax import nnx
import optax
class ArrayDataSource(grain.RandomAccessDataSource):
"""Wraps in-memory NumPy arrays as a Grain data source."""
def __init__(self, images, labels):
self._images = images
self._labels = labels
def __len__(self):
return len(self._images)
def __getitem__(self, index):
return {
'image': self._images[index],
'label': self._labels[index],
}
class Normalize(grain.MapTransform):
def map(self, element):
element['image'] = element['image'].astype(np.float32) / 255.0
return element
def build_data_loader(images, labels, batch_size, seed):
source = ArrayDataSource(images, labels)
sampler = grain.IndexSampler(
num_records=len(source),
shuffle=True,
num_epochs=None,
seed=seed,
)
operations = [
Normalize(),
grain.Batch(batch_size=batch_size, drop_remainder=True),
]
return grain.DataLoader(
data_source=source,
sampler=sampler,
operations=operations,
worker_count=4,
)
Model, optimizer, and training step
This should look familiar from week 4:
class ArchiveDigitizer(nnx.Module):
def __init__(self, hidden_dim: int, num_classes: int, rngs: nnx.Rngs):
self.linear1 = nnx.Linear(784, hidden_dim, rngs=rngs)
self.linear2 = nnx.Linear(hidden_dim, num_classes, rngs=rngs)
def __call__(self, x):
x = self.linear1(x)
x = nnx.relu(x)
return self.linear2(x)
rngs = nnx.Rngs(jax.random.PRNGKey(42))
model = ArchiveDigitizer(hidden_dim=128, num_classes=10, rngs=rngs)
tx = optax.adam(learning_rate=0.005)
optimizer = nnx.Optimizer(model, tx=tx, wrt=nnx.Param)
@nnx.jit
def train_step(model, optimizer, batch_images, batch_labels):
def loss_fn(m):
logits = m(batch_images)
loss = optax.softmax_cross_entropy_with_integer_labels(
logits=logits, labels=batch_labels
).mean()
return loss
loss_val, grads = nnx.value_and_grad(loss_fn)(model)
optimizer.update(model, grads)
return loss_val
The training loop, powered by Grain
# Assume train_images, train_labels are loaded NumPy arrays, e.g. from MNIST
batch_size = 64
data_loader = build_data_loader(train_images, train_labels, batch_size, seed=42)
data_iterator = iter(data_loader)
num_steps = 2000
print("Starting training with Grain data pipeline...")
running_loss = 0.0
for step in range(num_steps):
batch = next(data_iterator)
x_batch = jnp.array(batch['image']).reshape(batch_size, -1)
y_batch = jnp.array(batch['label'])
loss = train_step(model, optimizer, x_batch, y_batch)
running_loss += loss.item()
if step % 200 == 0 and step > 0:
avg_loss = running_loss / 200
print(f"Step {step:4d} | Average Loss: {avg_loss:.4f}")
running_loss = 0.0
print("Training complete.")
The training step itself did not change at all from Week 4. Only the mechanism feeding it data changed, and that separation is exactly the point. Your model logic and your data pipeline should be independent, swappable pieces.
Reproducibility and checkpointing
One benefit of Grain's explicit, seeded design is that data pipelines are fully reproducible. Given the same seed and the same source, two runs will visit records in the identical order and apply the identical augmentations at each step.
For long running jobs, you also want to be able to resume training after an interruption without skipping or repeating data. Grain integrates with Orbax, the checkpointing library we will cover in Week 9, allowing you to save the exact position of your data iterator alongside your model weights and optimizer state. When training resumes, both the model and the data pipeline pick up exactly where they left off.
Common pitfalls
A few mistakes come up often enough to call out directly.
Using np.random instead of the provided generator. Inside a RandomMapTransform, always use the rng argument passed into random_map. Calling np.random.random() directly breaks reproducibility, since that global state is not controlled by Grain's seeding.
Forgetting drop_remainder. As covered in Week 6, an inconsistent final batch shape triggers expensive JIT recompilation. Set drop_remainder=True on grain.Batch unless you have a specific reason not to.
Overestimating worker_count. More workers is not always better. If your __getitem__ implementation is already fast, for example reading from an in-memory array, spinning up eight worker processes adds multiprocessing overhead without any benefit. Benchmark before assuming.
Doing heavy work in the DataSource instead of Transformations. Keep __getitem__ minimal. Expensive operations like image decoding and augmentation belong in your transformation list, where Grain can distribute them across workers. A bloated __getitem__ cannot be parallelized as effectively.
Exercises
Build a custom augmentation. Write a RandomMapTransform that randomly rotates an image by a small angle, using the provided rng generator. Verify that running the pipeline twice with the same seed produces identical rotations.
Benchmark worker counts. Run the same data pipeline with worker_count set to 0, 2, 4, and 8. Time how long it takes to iterate through 500 batches at each setting. Plot the results.
Simulate a distributed shard. Manually construct two ShardOptions objects representing shard 0 of 2 and shard 1 of 2. Confirm that the two resulting samplers produce completely disjoint sets of indices across a full epoch.
Trigger and fix a recompilation. Remove drop_remainder=True from your grain.Batch call. Train for a few epochs and use chex.assert_max_traces from Week 6 to confirm that the final, undersized batch triggers a second compilation. Then fix it.
Quick reference
import grain.python as grain
import numpy as np
# DataSource
class MySource(grain.RandomAccessDataSource):
def __len__(self):
return len(self._data)
def __getitem__(self, index):
return self._data[index]
# Sampler
sampler = grain.IndexSampler(
num_records=len(source),
shuffle=True,
num_epochs=None,
seed=42,
shard_options=grain.ShardByJaxProcess(drop_remainder=True),
)
# Deterministic transform
class MyMap(grain.MapTransform):
def map(self, element):
return element
# Random transform
class MyRandomMap(grain.RandomMapTransform):
def random_map(self, element, rng):
return element
# Assembling the pipeline
data_loader = grain.DataLoader(
data_source=source,
sampler=sampler,
operations=[MyMap(), MyRandomMap(), grain.Batch(64, drop_remainder=True)],
worker_count=4,
)
data_iterator = iter(data_loader)
batch = next(data_iterator)
What's next
Our training loop can now pull in real data at real speed. But every training run we have written so far lives and dies with the Python process running it. If your machine crashes, if you lose power, or if you simply need to stop and resume tomorrow, all of that progress vanishes.
Next week we solve that problem with Orbax, JAX's checkpointing library. We will learn how to save not just model weights, but the entire training state: optimizer momentum, step count, and even the data iterator position from Grain, so a training run can be paused and resumed without losing a single batch of progress.
Next week: Never lose your progress: Checkpointing with Orbax
