Replication is an implementation fact
My public KataGo Transformer training stack has a JAX path that supports data parallelism across eight local TPU devices. Each device processes part of a batch, computes gradients, and participates in a collective reduction before updating its model replica.
That establishes an execution capability. It does not establish an eightfold speedup, better sample efficiency, or unchanged learning dynamics. Those claims need separate comparisons.
I find it useful to begin with a simpler question: what single-device computation should this distributed step reproduce? If the reference objective is unclear, a fast distributed run can optimize something slightly different without making that difference obvious.
The loss reduction defines the objective
For D devices, each holding b examples, an equally weighted mean loss has the form:
If every device computes the mean of its equally sized local batch, averaging its gradients reproduces the gradient of the global mean, up to floating-point differences. The repository implements the gradient collective with jax.lax.pmean.
The normalization matters when there are weights or masks. Averaging local losses that were each divided by a different sum of weights generally differs from dividing the global weighted sum by the global sum of weights. Both are definable objectives; they are not interchangeable.
In the current loss implementation, weighted loss terms are summed and divided by the local number of positions. With equal local batch sizes, averaging gradients corresponds to that per-position global objective. This is different from claiming that the implementation normalizes by the total nonzero target weight. Reading the denominator is part of reviewing the distributed algorithm.
State is larger than the parameter tree
A useful parity test holds one global batch fixed. Starting from the same parameters and optimizer state, split that batch across devices, perform one update, and compare against the reference update on the concatenated batch. I would compare gradients before comparing only the final scalar loss: different gradients can sometimes yield deceptively similar loss values.
I would also inspect optimizer moments, normalization statistics, moving loss statistics, random-number handling, and the count used by the learning-rate schedule. A synchronized parameter tree is insufficient if one of these other states evolves differently.
Exact continuation is an even stronger claim. Restoring model and optimizer arrays is valuable, but replaying an interrupted run also requires consistent sample order, data cursor, random state, and numerical settings. My repository provides automatic checkpoint resume; I would not infer bitwise-identical continuation from that feature name alone.
These are tests I would require before claiming full equivalence. They are not a report that every possible configuration in the repository has passed such a suite.
Time device work, not just Python dispatch
JAX dispatch is asynchronous. A Python call can return before the device finishes the work. The JAX documentation on asynchronous dispatch explains why a benchmark must synchronize on a result when measuring execution time.
There is also a difference between synchronizing every step and synchronizing at the end of a steady-state window. The first helps expose individual-step latency. The second can preserve the overlap the real workload relies on. I would label the timing mode, warm up compilation, and keep compilation time separate from the steady-state rate.
The training stack exposes both component profiling and end-to-end timing. Separately compiled components help locate a suspected bottleneck, but their timings are not guaranteed to sum to the compiled full step. Fusion and scheduling can change at the boundary.
Specify which scaling question you are asking
For a fixed global batch, increasing the device count is a strong-scaling experiment: each device gets less work and communication can become more prominent. Holding per-device batch constant instead increases the global batch. That can improve hardware utilization while changing the optimization problem faced by a fixed learning-rate schedule.
I would make the distinction explicit in a result table:
| Comparison | Held fixed | Main question |
|---|---|---|
| Fixed global batch | Total examples per update | How much of one step can be parallelized? |
| Fixed local batch | Examples per device | How does throughput grow when total work grows? |
| Fixed quality target | Evaluation criterion and tuning budget | Does the extra hardware reduce time to a useful model? |
The public README’s approximately 453 positions/s profile is a single-chip v6e measurement for a 24-layer, width-1024 model at batch 16. It should not be plotted as an eight-device result. I have not published a controlled eight-device scaling curve there, so this note does not supply one.
A note on the API
The published implementation uses pmap and local-device collectives. As checked on September 19, 2026, the current JAX documentation directs new work toward shard_map or related newer interfaces. This article explains the existing implementation; it is not a recommendation to begin every new distributed system with pmap.
The more durable questions survive an API migration: what objective is reduced, which state is synchronized, where the timing boundary lies, and what remains fixed as devices are added.
Sources and next test
- Training loop and loss implementation.
- Public configuration notes.
- JAX documentation for
pmean,pmap, and asynchronous dispatch.
The next result I would want is a one-step parity check followed by a repeated, fixed-global-batch scaling curve with input-pipeline and communication costs reported. That would turn “runs on eight devices” into an interpretable systems result.
Code snapshot reviewed: 90c12af. Later revisions may change implementation details.