DiPaCo: Distributed Path Composition

ArXiv: 2403.10616

🎯 Pitch

Training a 150M-parameter path through a modular network can match a 1B-parameter dense transformer while using 45% less wall-clock time, because DiPaCo replaces monolithic synchronization with rarely communicating modules trained in isolated silos. This means large-scale models can be trained on loosely connected, geographically distributed workers without the need to co-locate thousands of GPUs.

1. Executive Summary

This paper proposes DIstributed PAth COmposition (DiPaCo), a co-designed modular architecture and training algorithm that distributes computation by paths — sequences of shared modules that define input-output functions — rather than training a monolithic model. Using the C4 language modeling benchmark, the authors demonstrate that a DiPaCo model with 256 paths of 150 million parameters each, trained via a combination of coarse document-level routing (offline, once per sequence) and the low-communication DiLoCo optimization algorithm (local SGD with infrequent outer gradient averaging), matches the validation perplexity of a 1.3 billion-parameter dense transformer while reducing wall-clock training time by 45% and requiring only one-eighth the co-located devices per compute island. At inference time, routing more frequently (every 64 tokens) at test time allows the same 150M-parameter-per-path DiPaCo to match a 1B-parameter dense model using over 6× fewer parameters per forward pass, establishing that modular, communication-efficient distributed training can recover monolithic-model performance while decoupling model scale from co-location requirements.

2. Context and Motivation

The Core Problem: Scaling Requires Co-location, and Co-location Is Becoming Unsustainable

The central tension this paper addresses is a structural mismatch between how we want to scale machine learning and how current training paradigms require us to scale it. Modern neural network training relies on several forms of parallelism — data parallelism (replicating the model across devices, each processing different batches), model parallelism (splitting the model's layers across devices), and pipeline parallelism (streaming data through a chain of devices, each responsible for a segment of the model). All of these techniques share a common, unspoken assumption: the devices involved must be tightly interconnected, exchanging parameters, gradients, or activations at every training step.

This co-location requirement creates a specific, escalating bottleneck. As Section 1 of the paper states:

"State of the art models are still essentially monoliths, and their optimization requires exchanges of parameters, gradients, and activations at every step of the learning process."

The word monolith here is precise. A monolithic model is one where every parameter interacts with every input token through the full forward and backward passes, and gradient information must be aggregated globally before any parameter update. This forces a particular infrastructure setup: large clusters of identical accelerators in the same physical location, connected by high-bandwidth interconnects. The paper frames this as an engineering and infrastructure challenge that worsens as models grow — provisioning and managing thousands of tightly coupled devices for weeks or months of training is operationally complex.

But the problem extends beyond engineering logistics. The paper identifies what we might call organizational and collaborative scaling problems:

• Training run disposability: "The training process itself is often restarted for each new model release, essentially discarding much of the computation for training the last model." Because monoliths are trained end-to-end as unified artifacts, incremental improvements typically require retraining from scratch rather than updating or reusing components.

• Localized change is impossible: "it is difficult to localize the effects to the final model of changes to any step in the process." In a monolithic training pipeline, modifying data preprocessing, architecture, or hyperparameters affects everything downstream in ways that are hard to isolate.

• Community exclusion: The monolithic paradigm makes it "difficult to leverage the potential of the greater ML community, rather than a single organization." If only one organization can train the final model (because only they have the co-located cluster), then contributions from outside researchers — new modules, better routing strategies, domain-specific improvements — cannot be easily integrated.

These three challenges together suggest that monolithic training may become a hard scaling ceiling even if we solve the raw engineering of building larger clusters. The paper's thesis is that architecture and training algorithm must be co-designed to break this coupling between model scale and co-location requirements.

The Assumption Framework: Why Current Constraints May Be Artificial

Section 2.1 of the paper introduces two working assumptions that deliberately invert the current ML training paradigm:

1. Training compute (FLOPs) is relatively cheap
2. Communication is relatively expensive

These assumptions are not true today — in fact, the paper explicitly says so: "These assumptions are not realistic in the current ML training paradigm." Current infrastructure is built around expensive, co-located accelerators where communication is fast and abundant. So why structure a research program around assumptions that don't hold?

The paper's argument is forward-looking and partially normative. If model sizes (and therefore compute requirements) continue to grow faster than our ability to co-locate devices, then communication will become the relative bottleneck. This is a scaling trajectory argument: it may not be true at 100B parameters, but what about 10T? What about 100T? At some point, the physical constraints of power, cooling, and interconnect bandwidth will force distribution across multiple locations — at which point the assumptions flip.

Conversely, the paper suggests a feedback loop argument: because we have designed training algorithms assuming co-location, we have built infrastructure optimized for co-location; because infrastructure enables co-location, we have designed algorithms requiring it. This may be "a spurious local minimum where compute is more constrained than it needs to be." If we could train effectively across poorly connected, heterogeneous devices (different GPU types, different data centers, even different continents), we could tap into a much larger pool of aggregate compute — making FLOPs effectively cheaper by increasing supply, even if per-FLOP cost remained the same.

A third assumption rounds out the setting: "both during train and evaluation, we cannot instantiate models as large as we would like to have on any single compute island." This is the memory constraint. Even if communication were free, device memory limits how large a model a single island can host. The paper wants to build models whose total parameter count exceeds what any individual worker can hold in memory.

The Two Obstacles to Modular Distributed Training

Given the goal of training large models across poorly connected compute islands, what specifically prevents us from simply splitting a model into pieces and distributing the pieces? The paper identifies two technical obstacles that prior work has not jointly addressed:

Obstacle 1: Routing granularity. The dominant approach to sparsity in large language models is token-level Mixture of Experts (MoE), where a router decides which expert(s) to activate for each token at each routed layer. This is the design used in GShard (Lepikhin et al., 2021) and Switch Transformers (Fedus et al., 2021), which the paper cites as state-of-the-art for training FLOP efficiency. The problem with token-level routing for the distributed setting is that it requires constant module swapping: as a sequence is processed token-by-token, different tokens may activate different experts, meaning the parameters for those experts must be available at each step. This forces all experts to be co-located (or at least accessible with extremely low latency), defeating the goal of distributing across distant, poorly connected workers.

Obstacle 2: Synchronization frequency. Standard distributed training synchronizes gradients after every optimization step — all workers compute gradients on their data shard, then an all-reduce operation aggregates them before the next step. If workers are geographically distributed, this per-step synchronization becomes intolerably slow because of communication latency. Existing work on federated learning addresses this by reducing communication frequency (e.g., FedAvg, FedOpt), but these methods typically consider the setting of independent models trained on separate clients' data — they don't address how to train a single large model that is structured to share parameters across workers in a way that makes infrequent synchronization effective.

DiPaCo's core insight is that these two obstacles can be solved together by making two coordinated design choices: route coarsely at the document level (eliminating the need for within-sequence expert swapping) and use infrequent outer optimization (DiLoCo, which averages parameter differences across workers every few hundred steps rather than every step). The coarse routing enables pre-sharding — the training data is partitioned by path before training begins, so each worker knows exactly which data it will process with which path. The DiLoCo optimization enables module sharing across paths — when multiple paths use the same module, their local copies of that module's parameters are periodically averaged, keeping them in approximate sync without per-step communication.

Prior Approaches and Where They Fall Short

The paper situates itself against several lines of prior work, each of which addresses part of the problem but not the whole:

Token-level Mixture of Experts (Lepikhin et al., 2021; Fedus et al., 2021; Shazeer et al., 2017). These models replace feed-forward layers in transformers with sparsely-activated expert pools, routing each token to a subset of experts. They achieve impressive training FLOP efficiency — the total parameters can be enormous while the per-token activated parameters remain modest. However, as noted above, they require all experts to be co-located because different tokens in the same sequence activate different experts, and gradient computation for a single sequence may touch many experts within one training step. As the paper states:

"Token-MoE are currently state of the art with respect to training FLOP efficiency, but require even more co-located accelerators than the equivalent-activated dense model for training."

This is a crucial point: token-level MoEs actually increase the co-location requirement relative to dense models, because the total parameter set is larger (all experts must be available) even though per-token compute is lower. They solve the wrong problem — they optimize for FLOP efficiency under co-location, not for relaxing the co-location requirement itself.

Document-level routing with independent experts (Gross et al., 2017; Gururangan et al., 2023). These approaches route entire documents to specific experts and train those experts independently on their assigned data shards. This is essentially the "Flat MoE" described in Section 2.6.3 of DiPaCo — a set of completely independent networks, each specialized to a different data domain. The routing is typically done via unsupervised clustering (e.g., k-means on document features) rather than learned discriminatively. While this approach satisfies the distribution requirement (experts are independent, so they can be trained on separate devices with no communication), it has a fundamental limitation: capacity and overfitting. Each expert sees only its own data shard. If the number of experts grows large relative to the total data, each shard becomes too small, and experts overfit — a pattern the paper demonstrates empirically in Table 2, where Flat MoE degrades from 256 paths with independent modules. There is no mechanism for sharing learned representations across experts, so each expert must learn everything from scratch on its own (often small) data partition.

DiLoCo (Douillard et al., 2023). The DiLoCo algorithm is the immediate precursor to this work and provides the outer optimization mechanism that DiPaCo extends. DiLoCo showed that a dense model can be trained across multiple workers with drastically reduced communication: workers perform local SGD on their own data shards for H steps, then send parameter differences to a central server which averages them using an outer optimizer (Nesterov momentum), and redistributes the updated parameters. This works for dense models — all workers share the same architecture and eventually converge to similar parameters. However, DiLoCo on its own doesn't address the memory constraint: the dense model must still fit on each worker, limiting the maximum model size. DiLoCo also doesn't provide any mechanism for heterogeneous specialization — all workers train the same function on different data, but the architecture doesn't adapt to data domains.

Federated Mixture of Experts (Reisser et al., 2021). This line of work combines MoE architectures with federated learning, where different clients may have different local experts. However, these approaches typically still use token-level or per-client routing without the hierarchical module sharing structure that DiPaCo introduces, and they don't address the scaling questions (how many paths, how much sharing, how to route at test time) that DiPaCo investigates.

Branch-Train-Merge (Li et al., 2022). This approach trains separate expert models on domain-specific data shards, then merges them (e.g., by averaging parameters) into a single model. The paper acknowledges this as closely related to their flat MoE baseline. The limitation relative to DiPaCo is that branch-train-merge doesn't maintain the modular structure after merging — it collapses experts back into a monolith. DiPaCo keeps modules separate and routes between them, preserving the distributed deployment advantage (only one path needs to be materialized at inference).

Pathways (Dean, 2021; Barham et al., 2022). The Pathways vision is explicitly cited as sharing motivations with DiPaCo — the idea of a modular, asynchronous, multitask system where different components can be trained independently and composed at inference time. However, Pathways is described as a general framework for building such systems, not a specific architecture or training algorithm. As the paper states:

"Unlike the Pathways framework which supports training of general modular multimodal multitask asynchronous systems, we propose a particular instantiation of a modular system that supports such kind of distributed training."

DiPaCo can be seen as one concrete realization of the Pathways vision, focusing on the specific technical mechanisms (coarse routing + DiLoCo-style outer optimization) needed to make modular distributed training work for language modeling.

Crowdsourced Training (Ryabinin and Gusev, 2020; Borzunov et al., 2022). These projects explore training large models using volunteer-contributed compute (e.g., PETALS for inference and fine-tuning). They share the motivation of decoupling model scale from centralized infrastructure but focus on the systems and incentive challenges of volunteer computing rather than the architectural innovations needed to make modular training effective.

How DiPaCo Positions Itself

DiPaCo's contribution is not any single algorithmic innovation in isolation, but the combination of design choices that together make modular, communication-efficient training of large models practical:

1. Coarse, offline routing solves the within-sequence expert swapping problem, enabling data pre-sharding and distributed batch computation.

2. DiLoCo applied per-module (rather than per-model) enables parameter sharing across paths with infrequent communication, preventing the overfitting that plagues fully independent experts while maintaining distribution.

3. Hierarchical path structure (multiple levels of modules, where a path is a composition of one module per level) provides a flexible tradeoff between capacity (more independent modules) and transfer (more shared modules). Section 2.6.1 shows how path-specific modules can increase capacity without requiring communication, while shared modules enable generalization across related data domains.

4. Test-time routing frequency is decoupled from training-time routing. During training, routing must be coarse (once per document) to enable pre-sharding. During inference, the paper shows that routing more frequently — every 64 or 128 tokens — recovers substantial performance (Table 3: +0.74 perplexity from routing every 128 tokens vs. once per sequence). This means the training-time constraint doesn't permanently limit inference quality.

The paper explicitly frames its contribution as a first prototype toward a less synchronous, more modular paradigm, not as a fully optimized system. Section 6 acknowledges that DiPaCo is not FLOP-efficient compared to dense compute-optimal models, and the abstract calls it "a first prototype towards a new paradigm of large-scale learning." This framing is important: the paper is arguing for a direction — architectural and algorithmic co-design for distributed training — rather than claiming to have achieved a Pareto-optimal system.

The key tension the paper navigates is between specialization (each path becoming expert on its data shard) and generalization (sharing modules so that learning transfers across paths). Too much specialization (flat MoE with many independent experts) leads to overfitting (Table 2). Too much sharing (all paths share all modules, i.e., DiLoCo on a single dense model) fails to exploit the capacity of additional parameters. DiPaCo's hierarchical module structure with partial sharing is the mechanism for navigating this tradeoff.

3. Technical Approach

This is primarily a systems and architecture paper whose core idea is that by co-designing the model architecture (modular paths through shared components) and the training algorithm (infrequent communication via DiLoCo applied per-module), one can train a large language model across poorly connected, heterogeneous compute islands without ever instantiating the full model on any single device.

3.1 Reader Orientation

What the system does in plain language: DiPaCo is a way to train a very large language model by splitting it into many smaller "paths" — each path is a complete, functioning neural network built from a sequence of shared building blocks (modules). Different paths specialize on different types of text, and they communicate only occasionally to share what they've learned. At test time, any given piece of text is routed through exactly one path, so the system uses far fewer parameters per prediction than a monolithic model of equivalent quality.

The problem it solves and the shape of the solution: The fundamental obstacle is that training large models today requires all devices to be co-located and tightly synchronized, because gradients must be exchanged at every step. DiPaCo solves this by making two coordinated design choices: (1) route entire documents to specific paths before training begins, so each worker knows exactly which data it will process and never needs to swap parameters mid-sequence, and (2) use DiLoCo's outer optimization loop applied per-module, meaning workers train their local paths independently for hundreds of steps, then periodically average their parameter differences only for the specific modules they share — with no global synchronization required. The result is that the "model" exists as a distributed collection of paths, each small enough to fit on a modest compute island, but collectively matching the performance of a much larger dense model.

3.2 Big-Picture Architecture (Diagram in Words)

The DiPaCo system has five major components, organized into a training pipeline that runs in repeated phases:

1. Data Pre-Sharder (offline, before training): Takes the full training dataset and, using a coarse router (k-means or discriminative classifier applied to the first 32 tokens of each document), assigns every document to exactly one path index j ∈ {1, …, P}. Output: P disjoint (or optionally overlapping, via top-2 assignment) data shards {D₁, …, D_P} where shard D_i is the subset of data assigned to path π_i.

2. Path Architecture (never fully materialized): The model is defined as a composition of L levels, where each level l has K_l candidate modules (expert parameter sets). A path π_{j₁,…,j_L} is a specific choice of one module per level. The total number of possible paths is P = ∏_{l=1}^L K_l. For example, a 16 × 16 DiPaCo has L = 2 levels, each with K₁ = K₂ = 16 modules, yielding P = 256 paths. Each path has size 150M parameters. Some modules may be shared across multiple paths (providing transfer learning); others may be path-specific (providing capacity).

3. Training Workers (one per path, potentially fewer via task queue): Each worker is responsible for training exactly one path on its assigned data shard. During an inner optimization phase, a worker performs τ steps of standard SGD with AdamW on its local data, updating the parameters of the modules that constitute its path. Workers are completely independent during this phase — no communication, no synchronization. They save checkpoints when finished.

4. Outer Optimizer Executors (sharded by module): After all workers finish an inner phase, separate executor processes — one per module (l, e) — gather the local parameters of that module from every path that uses it, compute the average parameter difference (outer gradient) across those paths, and apply an outer optimizer (Nesterov momentum) to update a global copy of the module parameters. The updated global parameters are then redistributed to the workers for the next inner phase. Critically, the outer optimization is parallelized per-module, and no single executor ever holds the full model.

5. Router (discriminative, optionally updated between phases): After an initial training phase using a generative (k-means) router, a discriminative router is trained: all paths score a small held-out "router dataset" of documents, the best-scoring path for each document becomes a classification target, and a logistic regression classifier is trained to predict path assignments from the same prefix features (first 32 tokens' hidden states). This classifier then re-shards the entire dataset, and training continues. At test time, routing can occur more frequently — every W tokens — using a sequence transduction model trained to predict the best path for each token window.

Information flow during one training phase: Router assigns documents to shards → Training workers pull tasks from a task queue (path i, shard i, starting checkpoint) → Each worker runs τ local optimization steps independently → Workers save checkpoints to distributed file system → Checkpoint metadata written to Spanner database → Outer optimizer executors poll Spanner for ready checkpoints → Each executor loads the relevant module from all paths using it, computes outer gradient, updates global parameters → Updated module checkpoints saved → New phase begins with updated global parameters redistributed to workers.

3.3 Roadmap for the Deep Dive

• First, the formal notation and architecture definition (Section 2.3 of the paper): blocks, modules, paths, and the router — establishing the precise vocabulary and indexing scheme that everything else builds on. This is essential because DiPaCo introduces a specific hierarchical parameter indexing that differs from standard dense model notation.

• Second, the coarse routing mechanism (Section 2.4): how documents are assigned to paths before training, including both generative (k-means) and discriminative (classifier-based) approaches, and why coarse routing is the critical enabler for distributed training without per-step communication.

• Third, the DiLoCo optimization algorithm as applied to DiPaCo (Sections 2.5, 2.6): the inner-outer loop structure, how outer gradients are computed per-module rather than per-model, and the specific optimization techniques (outer gradient norm rescaling, loss reweighting, early stopping) that make the system work at scale.

• Fourth, the architectural flexibility mechanisms (Sections 2.6.1–2.6.3): how DiPaCo navigates the capacity-sharing tradeoff through path-specific modules, hierarchical routing, and the flat MoE baseline, which connects the approach to prior work on document-level expert models.

• Fifth, the infrastructure design (Section 3): the task queue system, worker pool, sharded outer optimization executors, and fault-tolerance mechanisms that enable training across hundreds of paths on heterogeneous, preemptible hardware — this is not mere engineering detail but a co-designed component of the system that makes the architectural assumptions viable.

3.4 Detailed, Sentence-Based Technical Breakdown

This paper introduces a system architecture and training algorithm where the model is never a single monolithic network but rather a collection of composable paths, each an independently functional neural network built from shared modules. The training algorithm decomposes into two nested loops: an inner loop where each path trains independently on its own data shard using standard SGD, and an outer loop where modules shared across multiple paths are periodically synchronized via parameter averaging with Nesterov momentum. The routing mechanism that assigns data to paths operates coarsely (once per document) and offline, enabling data pre-sharding and eliminating the need for within-sequence communication.

Formal Notation and Architecture Definition

The paper establishes a precise notation system to define the modular architecture, which is necessary because the parameter indexing is more complex than in a dense model — there are multiple levels of parameter sets, and the same parameters may appear in multiple paths.

Base model and parameter partitioning. Start with a base model architecture with parameters θ. Partition the parameter indices into L subsets B_l for l ∈ [1, …, L]. Each B_l represents a contiguous sub-network that can act as an input-output mapping — in the transformer experiments, a level corresponds to several consecutive transformer blocks. For each level l, there are K_l distinct parameter choices called "modules" or "experts." A module at level l and expert index e is denoted (l, e) where e ∈ [1, K_l].

Paths as module compositions. A path π_{j₁,…,j_L} is defined by choosing one module per level — j_l indexes the chosen module at level l. The total number of possible paths is:

$P = \prod_{l=1}^{L} K_l$

where L is the number of levels and K_l is the number of modules at level l.

For example, a 16 × 16 DiPaCo has L = 2, K₁ = 16, K₂ = 16, yielding P = 256 paths. The paper also collapses the tuple [j₁, …, j_L] into a single index j ∈ {1, …, P} when convenient.

What this nesting achieves: each path is a complete neural network — the entire base architecture with specific module choices at each level. A path can process an input from start to finish without ever needing parameters from other paths. This is the property that enables distributed training: a worker only needs the parameters for its assigned path, which is a fraction of the total model size.

Router definition. The router r maps an input x to a path selection. For hard routing (the only type used in this work), r(x) produces a deterministic tuple [j₁, …, j_L] choosing exactly one module per level. The paper specifically does not use soft routing, where the router produces a distribution over paths — the choice is single-valued. The subset of training data routed to path j is called its "shard" D_j, and there is a one-to-one association: path π_i is associated to shard D_i.

Parameter indexing across paths and time. Because paths are not fully synchronized during training, the notation must distinguish local (per-path) and global (synchronized) parameters:

• θ(l, e)^t_i: the local copy of parameters for module (l, e) in path i at outer step t. This is what worker i actually uses and updates during its local training.
• θ(l, e)^t: the global copy of parameters for module (l, e) at outer step t — after synchronization, this value is identical across all paths using that module, so the path subscript is omitted.
• θ^t_i: the collection of all parameters used by path i at outer step t — the concatenation of one module per level.
• Δ(l, e)^t_i and Δ(l, e)^t: the local and global outer gradients (parameter differences) for module (l, e).

What this indexing enables operationally: the distinction between θ(l, e)^t_i and θ(l, e)^t is not mere bookkeeping — it captures the fact that between outer optimization steps, different paths' copies of the same module diverge as each trains on its own data. The outer gradient Δ(l, e)^t measures the average divergence, and the outer optimizer uses it to update the global parameters. The paper uses "iteration t" to index outer steps (communication rounds) throughout Algorithm 1.

Block structure flexibility. The paper notes that while in practice B_l determines contiguous sub-networks of the base architecture (e.g., transformer blocks 0–5 vs. blocks 6–11), this is not required by the formalism. One could partition parameters arbitrarily — for example, B₁ could be all biases and B₂ all linear weights — but the contiguous block structure maps naturally to transformer layers and is what the experiments use.

Coarse Routing: How Documents Are Assigned to Paths

The routing mechanism is the critical enabling design choice for distributed training. The paper's key insight is that routing at the document level (rather than per-token) allows data to be pre-sharded before training begins, which in turn allows each path to be trained independently on its own data shard without any within-sequence communication.

Why routing must be coarse during training. Token-level routing, as used in GShard or Switch Transformers, requires that different experts be available at different positions within a single sequence. This means the parameters for all possible experts must be accessible with low latency during sequence processing, which effectively requires co-locating them. By routing once per document, every token in a document is processed by the same path, so a worker only needs the parameters for that single path — and the worker can batch all tokens of the sequence together without any module swapping.

What "coarse" means operationally. The router uses the first 32 tokens of each document as a "context" or "prefix" to decide which path should process the entire document. The remaining tokens are used for training (at train time) or perplexity evaluation (at test time). All methods, including dense baselines, are evaluated the same way — perplexity computed on all but the first 32 tokens — so the prefix overhead is consistent across comparisons.

Generative Routing (k-Means)

Generative routing makes no use of the language modeling task itself; it simply clusters documents based on feature similarity and assigns each cluster to a path.

Feature extraction. The feature z for a document is the average hidden state from the last transformer block of a pretrained base language model, computed over the first 32 tokens of the document. This produces a fixed-dimensional vector representation for each document, capturing its semantic and stylistic properties as encoded by the pretrained model.

k-Means clustering. Given P paths, run k-means (with k = P) on the features z of all training documents. Let {c₁, …, c_k} be the learned cluster prototypes. The router assigns a document with prefix feature z to shard D_{r(z)} via:

$r(z) = \arg\min_{i \in [1,k]} \|z - c_i\|^2$

where z is the feature vector for the document and c_i is the i-th cluster centroid.

What it computes: for each document, find the cluster prototype that minimizes the squared Euclidean distance to the document's feature representation, and assign the document to the corresponding path.

Why this form: Euclidean distance is the standard k-means objective — it finds a Voronoi tessellation of the feature space where each document is assigned to its nearest prototype. This is purely "generative" because it ignores the downstream task (language modeling perplexity) and only optimizes feature reconstruction. The advantage is simplicity: k-means is computationally cheap even for millions of documents, and it produces balanced shards by default (each prototype attracts roughly equal numbers of documents). The disadvantage is that feature similarity may not align with which path would actually model a document best — two documents could be similar in their first 32 tokens but require very different language modeling capabilities for their subsequent tokens.

Discriminative Routing

Discriminative routing directly optimizes the routing decision for the downstream task: it assigns each document to whichever path models it best (lowest perplexity), then trains a classifier to predict that assignment.

The alternating procedure. Discriminative routing approximates Expectation-Maximization (EM), alternating between:

1. E-step (update latent variables): Given the current paths, determine which path is best for each document in a small held-out "router dataset" (0.5% of C4, set aside for this purpose).
2. M-step (update model parameters): Train paths on the data shards produced by the current router.

In this work, only one full alternation is performed: paths are first trained with a generative (k-means) router, then a discriminative router is trained and used to re-shard the entire dataset, and paths continue training from there.

Score computation. For each document in the router data (length L tokens) and each path f_i for i ∈ {1, …, K}, compute the per-token auto-regressive log-likelihood, producing an n × L × K array S_{ijp} where:

$S_{ijp} = \log p_{f_i}(t_{l,j} \mid t_{l-1,j}, \ldots, t_{1,j})$

where S_{ijp} is the log-probability assigned by path f_i to token t_{l,j} — the j-th token of document l — given all preceding tokens. The p subscript indexes the path.

What it computes: for every document and every path, the per-token log-likelihood of the true next token under that path's language model. Summing these across tokens gives the total log-likelihood (negative perplexity) of the document under that path.

Label generation. The target label for document i is the path that maximizes total log-likelihood:

$T_i = \arg\max_p \sum_{j=1}^{L} S_{ijp}$

where T_i is the index of the best path for document i. This is the "hard EM" approximation — instead of using the full posterior distribution over paths (which would have a soft assignment), the method picks the single best path as a hard target.

Classifier training. A K-class linear logistic regression classifier is trained with T_i as targets and g(document_i) — the average hidden state from the initial LM over the first 32 tokens — as features. This is a simple linear model on top of frozen features; the paper does not finetune the feature extractor. The classification uses a cross-entropy objective (standard for multi-class logistic regression).

Bias correction for imbalanced assignments. A practical problem arises with many paths: the paths that receive the fewest documents under argmax_p Σ_j S_{ijp} (the oracle assignment) may receive even fewer (or zero) under the trained classifier, because the classifier has limited capacity and tends to favor majority classes. To remedy this, the paper trains an additional bias term specifically to match the target document-to-path distribution — ensuring that each path receives documents in proportion to its oracle assignment frequency. Without this correction, some paths would be "empty" under the classifier and never receive training data, effectively wasting capacity.

Why discriminative routing improves over generative. The generative k-means router clusters documents by feature similarity, which is only a proxy for "which path will model this document well." Discriminative routing directly optimizes the routing criterion — assign each document to the path that achieves lowest perplexity on it — and learns to predict this assignment from features. The gain is larger with more paths (Figure 10) because the margin between random assignment and optimal assignment grows with the number of choices. Table 5 shows an absolute improvement of 0.7 perplexity points for an 8 × 8 DiPaCo with discriminative vs. generative routing.

Why this form (hard EM): the hard assignment is an approximation to full EM, which would compute a soft posterior over all paths for each document and weight training data accordingly. Hard assignment is simpler to implement with the pre-sharding infrastructure (each document goes to exactly one shard) and avoids the computational cost of training every document on every path. The paper notes that this makes the overall process an "approximation" to EM, and that more alternations yield diminishing returns (Figure 11: "14.0 → 13.38 → 13.36 → 13.25 PPL for phases 0, 1, 2, 3 respectively").

Routing More Frequently at Test Time

During training, routing occurs once per document to enable pre-sharding. During evaluation, the paper shows that routing more frequently — splitting a sequence into chunks and potentially re-routing at chunk boundaries — recovers substantial performance, closing the gap between DiPaCo and dense models.

The test-time routing procedure (Figure 3). At test time, a sequence of length S tokens is divided into chunks of W consecutive tokens (e.g., W = 128). The router, given the i-th chunk, predicts the best path to use for the (i+1)-th chunk. No model parameters from different paths need to be co-located — when a re-route occurs, text is simply communicated to the newly selected path's worker, which processes the next chunk.

Router training for sequence-level prediction. The token-level router is trained as a sequence transduction model. Given documents in the router data with per-token scores S_{ijp} (log-likelihood of each path at each position), the target output at position j is:

$T_{ij} = \arg\max_p \sum_{j}^{j'} S_{ijp}$

where j' = min(document length, j + L - 1), and L is a window size (chosen as 1024 — the full sequence length — for the results in Table 3, though the authors note that choosing L equal to the re-routing frequency was marginally better). The target T_{ij} is the path that maximizes total log-likelihood over the next L tokens starting from position j.

What it computes: for each token position, determine which path would have been optimal for the following window, and train a transformer to predict this optimal path index from the input tokens up to that position. The trained router can then make path predictions at any token boundary during evaluation.

Infrastructure implications. The cost of switching paths is small because the router runs infrequently (every 64–128 tokens, not every token) and in scoring mode — predicting token probabilities — not auto-regressive generation mode. Only text needs to be communicated to the router and the newly selected path; KV-caches are not shared across paths, so a re-route requires re-computing the KV-cache for the new path, which the paper acknowledges as a deployment inefficiency not addressed in this work (Section 6).

Results of frequent routing (Table 3). For a 16 × 16 DiPaCo with 256 paths of 150M parameters each:

• Routing once per sequence: 12.39 PPL (with early stopping: 12.22)
• Routing every 128 tokens: 11.48 PPL
• Routing every 64 tokens: 11.38 PPL
• Routing every 32 tokens: 11.31 PPL
• Routing every 16 tokens: 11.26 PPL

Routing every 64 tokens allows the 150M-parameter-per-path DiPaCo to match the 1B dense baseline (11.41 PPL) using over 6× fewer parameters per forward pass. The improvement from once-per-sequence to every-128-tokens is 0.74 perplexity, which is the largest single gain. Further increases in routing frequency yield consistent but diminishing returns.

The DiLoCo Optimization Algorithm Applied Per-Module

DiLoCo is the training algorithm that enables infrequent communication across paths. In the original DiLoCo work (Douillard et al., 2023), the algorithm trained a single dense model across k workers: all workers start from the same parameters, each trains independently on its own data shard for H inner steps, then parameter differences are averaged on a central server using an outer optimizer, and the updated global parameters are redistributed. DiPaCo extends this by applying DiLoCo per-module rather than per-model — different modules may be updated with gradients from different subsets of workers, depending on which paths use which modules.

The inner optimization loop. During an inner phase, each worker i (associated with path π_i) performs τ local optimization steps on its own data shard D_i. Starting from the synchronized global parameters received at the beginning of the phase (θ_i^t = θ_i^{t-1}, copying the global parameters into local copies), the worker:

1. Samples a batch x ~ D_i from its shard.
2. Computes the language modeling loss L = f(x, θ_i^t).
3. Updates local parameters via InnerOpt: θ_i^t ← InnerOpt(θ_i^t, ∇_L).

The inner optimizer is AdamW with the hyperparameters listed in Table 4 (Appendix): peak learning rate 4 × 10⁻⁴, cosine schedule, batch size 512, sequence length 1024 tokens. These are standard transformer LM training hyperparameters, applied identically to each path.

What happens to shared modules during local training. If two paths π₁ and π₂ both use module (l, e), they each have their own local copy of θ(l, e). During the inner phase, the two copies diverge because they are trained on different data shards (and possibly different loss surfaces, since the other modules in the paths may differ). The divergence accumulates over τ steps.

The number of inner steps. In the experiments, τ = 150 for most configurations. This choice balances two considerations: more steps per phase means less frequent communication (good for the distributed setting) but also more parameter divergence between workers (which may hurt the quality of the outer update, since the outer gradient is an average of potentially very different local updates).

The outer optimization loop (lines 15–20 of Algorithm 1). After all workers complete their inner phases, the outer optimization proceeds module-by-module. For each level l ∈ [1, …, L] and each module index e ∈ [1, …, K_l]:

1. Identify the set of paths that use module (l, e). Let P_{l,e} be the number of such paths.
2. Gather the local final parameters θ(l, e)^t_i from each such path i.
3. Compute the outer gradient — the average parameter difference between the global parameters from the previous phase and the local parameters after local training:

$\Delta(l, e)^t \leftarrow \frac{1}{P_{l,e}} \sum_{i=1}^{P_{l,e}} (\theta(l, e)^{t-1} - \theta(l, e)_i^t)$

where Δ(l, e)^t is the outer gradient for module (l, e) at outer step t, P_{l,e} is the number of paths using this module, θ(l, e)^{t-1} is the global parameter value from the previous outer phase, and θ(l, e)_i^t is worker i's local copy after inner training.

What it computes: for each shared module, the outer gradient is the average of how much each path's local training "moved" the module parameters away from the synchronized starting point. If all paths moved the parameters in similar directions, the average is large and coherent. If paths moved in different directions, the average cancels out — indicating that the module may benefit from being split into path-specific variants.

Why the difference direction: the computation subtracts (θ_global - θ_local) rather than averaging the local parameter values directly. This formulation is equivalent to averaging the local parameter moves (since all workers started from the same θ(l, e)^{t-1}) but is numerically more stable and is the standard formulation in the DiLoCo framework. The subtraction order means the outer gradient points from the local parameters back toward the global starting point — when added to the global parameters with the outer optimizer, it moves them in the direction the local training took them on average.

4. Apply the outer optimizer to produce updated global parameters:

$\theta(l, e)^t \leftarrow \texttt{OuterOpt}(\theta(l, e)^{t-1}, \Delta(l, e)^t)$

The outer optimizer is Nesterov momentum with outer learning rate 0.7 and outer momentum 0.9, following the recipe from the original DiLoCo paper. These values were not heavily tuned; the authors "searched over relatively few hyper-parameters: mainly learning rate and value of Nesterov momentum" (Section 4).

What Nesterov momentum does in this context. Standard momentum accumulates past outer gradients and adds them to the current one, smoothing the update trajectory. Nesterov momentum additionally looks ahead — it computes the gradient at the extrapolated position (current parameters plus the momentum term), which tends to reduce oscillation and speed convergence. In the DiPaCo setting, Nesterov momentum on the outer gradients serves a dual purpose: it smooths out noise from the stochastic inner training (different batches, different data distributions across shards) and it provides a form of "inertia" that prevents the global parameters from jumping too far based on a single inner phase.

Why AdamW for inner, Nesterov for outer. The paper uses different optimizers for the two loops. AdamW (inner) is the standard choice for language model training — it provides per-parameter adaptive learning rates, which handles the varying gradient scales across different weight matrices in transformers. Nesterov momentum (outer) is a simpler optimizer with a single global learning rate, chosen because the outer loop operates on parameter differences (which have already been locally optimized by AdamW) and needs to aggregate them across potentially varying numbers of paths. The outer loop learns less frequently (every 150 steps) and sees "pre-processed" gradient information, so it benefits from momentum's smoothing more than Adam's adaptivity.

The overall DiPaCo algorithm (Algorithm 1). The full procedure repeats for T outer steps:

For outer step t = 1…T:
    1. (Optional, done once in this work) Discriminatively re-shard data
    2. For each worker i in parallel:
        a. θ_i^t = θ_i^{t-1}  (copy global parameters to local)
        b. For τ inner steps:
            - Sample batch x ~ D_i
            - Compute loss L = f(x, θ_i^t)
            - θ_i^t ← InnerOpt(θ_i^t, ∇L)
    3. For each module (l, e) in parallel:
        a. Gather local parameters from all P_{l,e} workers using this module
        b. Compute outer gradient Δ(l, e)^t
        c. Update global: θ(l, e)^t ← OuterOpt(θ(l, e)^{t-1}, Δ(l, e)^t)

What happens for modules used by only one path (path-specific modules). When P_{l,e} = 1 — the module is used by exactly one path — there is no averaging of outer gradients. However, the paper still applies the outer optimization step (line 18 of Algorithm 1) rather than simply keeping the local parameters, because "empirically it improves convergence over the default optimizer" (Section 2.6.1 footnote). This is a non-obvious design choice: even for path-specific modules, the Nesterov momentum outer optimizer applied to the single local update provides a form of momentum-based smoothing that outperforms directly using the AdamW-local-trained parameters.

Parameter synchronization comparison: DiLoCo vs. full synchronization. An important ablation (Section 4.5) compares DiPaCo trained with the DiLoCo partial synchronization against a fully synchronous version where all paths compute gradients on their own data shards, then gradients are aggregated module-by-module, and a single AdamW step is taken with the aggregated gradient (i.e., standard distributed data-parallel training, but with path-specific routing). The results are striking:

• 2 × 2 DiPaCo: DiLoCo actually outperforms full synchronization by 0.3 PPL
• 4 × 4 DiPaCo: DiLoCo outperforms by 0.6 PPL
• 8 × 8 DiPaCo: Full synchronization is only 0.1 PPL better, despite communicating hundreds of times more frequently

This is a counterintuitive result. The paper's interpretation (Section 4.5): "This suggests that DiLoCo is an effective distributed optimization algorithm for DiPaCo." A possible explanation, though not developed in detail, is that the outer Nesterov momentum provides beneficial smoothing of the optimization trajectory that per-step synchronization lacks, and that the local training phases allow each path to make more progress on its specialized data before being pulled back toward the global consensus.

Advanced Optimization Techniques

The paper introduces three techniques that improve training stability and convergence when scaling to many paths with uneven data distributions.

Outer Gradient Norm Rescaling. Different modules are used by different numbers of paths. A module at a shared early level (e.g., the first few transformer blocks) might be used by all 256 paths in a 16 × 16 DiPaCo, while a path-specific module at a later level is used by only one path. When computing outer gradients, the average over many paths produces a gradient with a different norm (typically smaller, due to averaging effects) than the update from a single path.

The paper uses the intuition that averaging across n paths is akin to using an n× larger batch size, which would typically produce a gradient with smaller variance (and thus potentially smaller norm). To compensate, the outer gradient norm is rescaled by the square root of the number of paths:

$\Delta(l, e)^t \leftarrow \sqrt{P_{l,e}} \cdot \frac{1}{P_{l,e}} \sum_i (\theta(l, e)^{t-1} - \theta(l, e)_i^t)$

where the factor √P_{l,e} rescales the averaged gradient to have a norm comparable to what a single-path update would have.

Why the square root: in standard SGD, if n independent gradients are averaged, the variance scales as σ²/n, so the standard deviation scales as σ/√n. Rescaling by √n restores the original scale. The paper does not prove that this is optimal for the DiPaCo setting (where paths are not independent — they started from the same parameters and train on related data), but presents it as empirically motivated.

Loss Reweighting. Data shards for different paths may have very different sizes due to the routing mechanism (k-means tends to produce balanced clusters, but discriminative routing may produce imbalanced shards based on path specialization). If all paths are trained with equal weight in the outer gradient averaging (line 17 of Algorithm 1), then paths with smaller shards would be over-represented — their local parameter updates would contribute equally to the outer gradient despite being trained on less data.

To compute an unbiased estimate of the gradient over the full dataset, outer gradients are weighted proportionally to shard size:

$\Delta_{l,e}^{(t)} \leftarrow \sum \alpha_{l,e} (\theta_{l,e}^{(t-1)} - \theta_{l,e}^{(t)})$

with:

$\alpha_{l,e} = \frac{|\mathcal{D}_{l,e}|}{\sum |\mathcal{D}_{l',e'}|}$

where |D_{l,e}| is the number of tokens in the shard for the path using module (l, e) and the denominator sums over all paths using that module.

What this corrects: without reweighting, a path with 1M tokens would influence the outer gradient as much as a path with 100M tokens, despite the latter having more statistical evidence for its parameter moves. The weighting ensures that the outer gradient reflects the data distribution across shards.

Early Stopping (Path-Specific). A small subset of training examples in each shard is set aside as a path-specific validation set. For each path, the parameters that yield the lowest loss on this validation set are selected, rather than always taking the final parameters after τ inner steps.

Why this matters: when the number of paths is large and some shards are small, paths may begin to overfit their shard-specific data distribution within a single inner phase. Path-specific early stopping detects this and selects an earlier checkpoint. The paper notes (Section 2.7): "We found that early stopping improves generalization on small shards, e.g., when the number of shards is large."

Architectural Flexibility: Path-Specific Modules, Hierarchical Structure, and Flat MoE

The paper emphasizes that DiPaCo provides a flexible framework for trading off between sharing (which enables transfer learning and prevents overfitting on small shards) and capacity (which enables specialization and makes use of additional parameters).

Path-Specific Modules (Section 2.6.1). Some levels can have modules that are used by exactly one path. In the extreme, a level can have K_l = P modules — one per path, with no sharing. Figure 5 illustrates this: level 3 modules are path-specific. This is "a particularly easy way to increase parameter count in DiPaCo" because it requires no communication for those modules' outer gradients (there is only one path contributing). However, the outer optimization step (Nesterov momentum) is still applied even for single-path modules, as noted above.

What path-specific modules enable. In the experiments (Figure 9), adding path-specific modules consistently improves perplexity — for a given path count, unsharing some of the parameters increases total model capacity. The specific configuration used in experiments shares transformer blocks 0–5 and 6–11 (as well as the embedding matrix) across some paths, while other blocks are path-specific, providing a middle ground between full sharing (all paths identical) and full independence.

Hierarchical DiPaCo Structure (Section 2.6.2). The model can be scaled by increasing both the number of levels and the number of modules per level. A 16 × 16 DiPaCo has 2 levels and 256 paths. A 32 × 32 × 32 DiPaCo would have 3 levels and 32,768 paths. The paper notes that finding enough devices for 32,768 paths is difficult, so a sampling approach is proposed: at each start of inner optimization, sample a subset of the paths to train, cycling through paths over outer phases. This is not implemented in the current experiments but is presented as a path to arbitrary scaling.

Flat Mixture of Experts (Section 2.6.3). The extreme form of capacity increase is to have every path be a completely independent network — no parameter sharing at all. This corresponds to L = 1 (only one level), B = B₁ is the entire network, and K = K₁ is the number of independent paths. This is equivalent to the approach in Gross et al. (2017) and Gururangan et al. (2023), except the paper uses discriminative routing rather than k-means.

Why Flat MoE is a strong baseline for compositional DiPaCo. For a small number of paths (relative to total data), fully independent paths can each specialize deeply on their data domain without interference from shared parameters, potentially outperforming a comparable DiPaCo with shared modules. Table 1 shows that for P = 64, "Using 64 unshared paths we can approach a 16 × 16 DiPaCo that has no unsharing." However, this advantage flips when the number of paths grows large: Table 2 shows that Flat MoE with 256 paths degrades due to overfitting (each shard becomes too small), and even with top-2 overlapping shards and early stopping, the Flat MoE reaches only 13.6 PPL at 64K steps while continuing to overfit, whereas "there is no overfitting at 256 paths with overlapping shards with a 16 × 16 DiPaCo."

Why the compositional structure prevents overfitting. In a 16 × 16 DiPaCo, even though there are 256 paths, each module is used by 16 paths (since there are 16 modules per level). This means each module is trained on 16 × (shard_size) tokens, compared to shard_size tokens in Flat MoE. The shared modules see effectively 16× more data, gaining statistical strength from data assigned to different paths. This prevents the overfitting that occurs when each path must learn its entire parameter set from a single, potentially small shard.

What overlapping shards provides. When using generative or discriminative routing, each document can be assigned to its top-2 closest clusters (Equation 1, extended to top-2) rather than the top-1. This means each document appears in two shards. At training time, this limits the specialization of each path (since it now sees data that also appears in another path's shard) and increases storage, but "otherwise does not have a computational cost, as each path is still trained independently on its shard." The 16 × 16 DiPaCo uses top-2 overlapping shards at training time. Overlapping at evaluation time — forwarding through multiple paths and aggregating scores — would increase inference cost and is not used in this work.

Infrastructure Design

The infrastructure described in Section 3 is not mere engineering detail but a co-designed component of the system: the DiPaCo architecture and training algorithm assume specific infrastructure capabilities (independent workers, asynchronous task completion, fault-tolerant task queues), and the infrastructure is purpose-built to realize DiPaCo's assumptions efficiently.

The five-component infrastructure (Figure 6). The training workflow proceeds through six stages per phase:

1. Task scheduling: A training task scheduler (purple in Figure 6) adds training tasks to a training task queue at the start of each phase. Each task is a (path_id, shard_id) pair specifying which path to train on which data shard, along with a starting checkpoint.

2. Training execution: Workers in the worker pool (orange) pull tasks from the queue. When a worker becomes available, it fetches the next task, loads the checkpoint and shard, performs inner optimization (τ steps of AdamW) on accelerators, saves a checkpoint to the distributed file system (GFS), and records the checkpoint path with metadata (path ID, outer step ID, phase ID) in a Spanner database table (blue). The worker then becomes available for the next task.

3. Evaluation triggering: When a checkpoint is saved, evaluation tasks for that checkpoint are added to an eval task queue (yellow), consumed by evaluation workers in the same worker pool.

4. Outer optimization: A sharded outer optimizer task scheduler (light blue) distributes outer optimization tasks to sharded outer optimization executors (red). Each executor is responsible for a subset of modules (e.g., one module or a collection). Each executor loads training checkpoints containing its assigned modules as they become available in the Spanner database, performs parameter averaging and outer optimization (Nesterov momentum), and saves updated module checkpoints.

5. Phase completion and restart: After all training tasks finish, the current phase concludes. The next phase can begin as soon as each module finishes its outer optimization — training tasks in the new phase can start immediately for paths whose modules are ready, without waiting for all modules.

6. Health monitoring: A job status monitor (green) periodically checks the health of all workers and task queue servers, restarting them if unresponsive.

Worker pool design. The producer-consumer pattern decouples task assignment from task execution. Training tasks are independent — no synchronization or communication among workers. If a worker fails or is preempted, the fault-tolerant task queue server returns the incomplete task to the queue for reassignment to another worker. The worker pool can contain "heterogeneous types of devices across different regions," and can auto-scale pool size according to resource availability.

Multi-host synchronization. In multi-host training where SPMD (Single Program, Multiple Data) parallelism is used across hosts within a single training worker, a naive task queue client on each host would pull different tasks. The solution is a synchronized task queue client: only the first JAX host creates a real client; all other hosts get the response via a blocking all-gather operation across hosts. This ensures all hosts within a worker process the same task.

Outer optimization efficiency optimizations. As the number of paths scales to hundreds, the outer optimization becomes a potential bottleneck. Three specific optimizations keep average phase time under 2 minutes:

• Online parameter gradient averaging: Rather than waiting for all paths to finish before starting averaging, each executor loads and accumulates checkpoints into a running sum as soon as each becomes available. This overlaps computation with training completion.

• Sharded outer optimization executor (Figure 7): Different modules' parameter averaging is independent, so it is distributed across multiple servers. This reduces memory per executor (only one or a few modules' parameters rather than the full model) and speeds processing. Critically, "the overall model is never materialized in a single location but always split across several servers."

• Asynchronous checkpoint gathering: Checkpoints may be stored on servers geographically distant from the outer optimizer executor. An Effingo process (Google's internal copy service) is launched in the background to bring checkpoints closer before loading, reducing transfer latency.

Additional micro-optimizations include caching outer optimizer parameters, reusing just-in-time compiled optimizer code, and loading multiple paths in parallel via multi-threading queues.

The backup pool (Section 3.4). Ideally, there would be one training worker per path — 256 workers for a 16 × 16 DiPaCo. With each worker using 16 A100 GPUs, this would require 4,096 GPUs, which may exceed availability. The backup pool solution: spawn as many workers as there are paths across multiple accelerator types using low-tier (preemptible) priority. When an accelerator is available, it is allocated to a worker and runs one inner optimization phase (which takes roughly a couple of minutes). If preempted, the task returns to the queue. This allows training to proceed with fewer than P dedicated workers by doing multiple rounds per phase, cycling workers through paths.

Design Choices Summary and Their Justifications

• Document-level routing (not token-level): enables data pre-sharding, which is the prerequisite for independent path training without per-step communication. Token-level routing would require expert parameters to be available at each token boundary, defeating distribution.
• Offline routing (pre-computed before training): allows each worker to know its entire data shard in advance, enabling efficient batch construction and eliminating online routing overhead during training.
• Hard routing (single path per document): simpler than soft routing (distribution over paths), which would require weighting training data across multiple paths and complicate the pre-sharding logic. Hard assignment makes the training infrastructure straightforward.
• DiLoCo applied per-module (not per-model): allows different subsets of paths to synchronize at different granularities — shared modules synchronize across all paths using them, path-specific modules don't synchronize at all. This aligns communication with the architectural structure.
• Nesterov momentum outer optimizer (not Adam): the outer loop operates on pre-processed gradient information from local AdamW training. Nesterov provides effective smoothing across phases with a simple global learning rate, avoiding the complexity of per-parameter adaptive outer learning rates.
• Discriminative routing over generative: directly optimizes the routing criterion (which path models each document best) rather than using a proxy (feature similarity). Gains are larger with more paths (Figure 10) because the gap between proxy-optimal and task-optimal assignment grows.
• Bias-corrected logistic regression for discriminative routing: prevents path starvation where some paths receive zero training data under the learned classifier, which would waste model capacity.
• Top-2 overlapping shards at training (not evaluation): increases the data seen by each path (reducing overfitting) without increasing inference cost. At evaluation, only one path is used per chunk.
• Path-specific early stopping: detects and mitigates overfitting on small shards, which is critical when the number of paths is large.
• Frequent test-time routing (decoupled from training): recovers the performance gap from coarse training-time routing without changing the training procedure. During training, coarse routing enables distribution; during inference, frequent routing enables quality — the two are independently optimized.

4. Key Insights and Innovations

Innovation 1: Co-Design as a First-Class Principle — Architecture and Optimization Are Not Independent Choices

The most distinctive intellectual move in DiPaCo is not any single algorithm or architectural pattern, but the methodological claim that architecture and distributed optimization must be designed together to achieve training across poorly connected devices. This sounds obvious at first glance — of course architecture matters for distribution — but the paper shows that the field's dominant approaches have systematically decoupled these choices, and that decoupling imposes a hard scaling ceiling.

What the field did before. The standard approach to large-scale training, as the paper documents in its related work (Section 5), treats architecture and distribution as separable concerns. Token-level MoE models (Shazeer et al., 2017; Lepikhin et al., 2021; Fedus et al., 2021) design architectures for per-token sparsity — replacing FFN layers with expert pools to increase parameter count without proportionally increasing FLOPs — but assume all experts are co-located and accessible with low latency. The optimization is standard synchronous SGD with all-reduce gradient aggregation. Conversely, federated learning and distributed training methods (Reddi et al., 2021; Douillard et al., 2023) design optimization algorithms for reduced communication but assume a dense, monolithic architecture that fits on each worker. Neither side addresses the joint problem: what architecture would make infrequent communication sufficient for training a model that no single worker can host?

What makes DiPaCo's co-design distinctive. The paper makes two specific architectural choices that are directly motivated by the optimization constraints, and vice versa:

• Coarse routing is an architectural choice driven by the optimization need for pre-sharding. If communication happens only every ~150 steps (the DiLoCo outer loop frequency), then during those 150 steps, each worker must be entirely self-sufficient — it cannot request parameters from other workers, swap experts, or synchronize gradients. This requirement implies that routing must happen at a granularity coarser than a training step, preferably before training begins. The paper's choice of document-level, offline routing is not an arbitrary design decision — it is the necessary architectural counterpart to infrequent communication. If DiPaCo used token-level routing (as in GShard), the infrequent communication would be impossible because different tokens in the same batch would activate different experts, requiring those experts' parameters to be available locally.

• Module-level DiLoCo is an optimization choice driven by the architectural need for partial sharing. The DiPaCo architecture has modules shared across subsets of paths. If the optimization treated all paths identically (as in standard DiLoCo applied to a dense model), it would either oversynchronize (forcing path-specific modules toward a consensus they shouldn't share) or undersynchronize (failing to transfer learning across shared modules). The per-module outer gradient computation — averaging only over paths that actually use each module — is the optimization counterpart to the architectural structure. The architecture says "these paths share this module"; the optimization says "then average only their updates for that module."

This co-design claim is not merely a system engineering observation. It implies that the search space for distributed training solutions is the Cartesian product of architecture and optimization choices, not the union. Prior work explored architecture innovations (token MoE) under the assumption of full synchronization, and optimization innovations (DiLoCo, FedOpt) under the assumption of dense models. DiPaCo argues that the interesting solutions lie at the intersection: architectures designed around the communication budget, optimizers designed around the sharing structure.

Evidence anchoring the claim. The paper's comparison of DiPaCo with fully synchronous training (Section 4.5) provides the clearest evidence that the co-design matters. A 16 × 16 DiPaCo trained with full synchronization — gradients aggregated module-by-module at every step — performs only 0.1 PPL better than the DiLoCo version, despite communicating hundreds of times more frequently. This is not a claim that DiLoCo is "as good as" full synchronization — it is evidence that the DiPaCo architecture is specifically well-suited to infrequent communication, and that adding more communication provides negligible benefit because the architecture already structures parameter sharing appropriately. A dense model, by contrast, would degrade substantially if communication were similarly reduced.

Significance beyond raw performance. This co-design framing is a conceptual contribution that reframes the distributed training problem. Rather than asking "how can we make existing architectures train with less communication?" (a compression/approximation problem), DiPaCo asks "what architecture would be naturally amenable to low communication?" This shifts the research agenda from communication-efficient optimization (which attempts to salvage synchronous training under constraints) to communication-aware architecture design (which assumes constraints and builds around them). The difference is fundamental: the former treats distribution as a tax on a pre-existing design, while the latter treats it as a design parameter.

Innovation 2: The Sharing-Capacity Tradeoff as an Architectural Spectrum, Not a Binary Choice

DiPaCo introduces a nuanced, continuous tradeoff between parameter sharing (which enables transfer learning and statistical efficiency) and capacity (which enables specialization and model scale), and shows that the optimal point on this spectrum depends on the number of paths and the amount of training data. Prior work on modular architectures treated this as a binary choice: either share everything (dense models, DiLoCo) or share nothing (independent expert models, Flat MoE). DiPaCo's hierarchical module structure with configurable per-level sharing ratios makes this tradeoff an explicit architectural hyperparameter.

What the field did before. The prior landscape offered two extremes with nothing in between:

• Full sharing / dense models: All parameters are used for all inputs. This maximizes transfer (every training example influences every parameter) but provides no mechanism for specialization — the model must learn a single function that works for all data. DiLoCo (Douillard et al., 2023) falls here: it trains a dense model across workers with infrequent communication, but every worker uses the identical architecture. The model capacity is bounded by what fits on a single worker.

• No sharing / independent experts: Each expert is a completely separate model trained on its own data shard. This maximizes specialization (each expert can adapt fully to its data domain) but provides no transfer — learning on one shard does not benefit learning on another. This is the approach in Gross et al. (2017), Gururangan et al. (2023), and the branch-train-merge paradigm (Li et al., 2022). The limitation, as DiPaCo empirically demonstrates (Table 2), is that as the number of experts grows relative to the data, each shard becomes too small and experts overfit.

Why the in-between matters. DiPaCo's hierarchical composition — L levels, each with K_l modules, paths formed by taking one module per level — creates a space of partial-sharing configurations that neither extreme explores. In a 16 × 16 DiPaCo, each module at level 1 is used by 16 paths, and each module at level 2 is used by 16 paths (different subsets), meaning each path shares some but not all parameters with other paths. The degree of sharing is controlled by the product structure: if a path shares module e at level l with P_{l,e} other paths, those paths are "related" along dimension l but may differ at other levels. This creates a form of structured transfer — paths that share early modules learn common representations (e.g., syntax, basic semantics) while paths that differ at later modules specialize (e.g., domain-specific knowledge).

What makes this a genuine innovation rather than an obvious generalization of MoE is the empirical discovery that the optimal sharing ratio depends on the total number of paths. For small numbers of paths (e.g., P = 64), Flat MoE (no sharing) approaches the performance of 16 × 16 DiPaCo (Table 1): "Using 64 unshared paths we can approach a 16×16 DiPaCo that has no unsharing." This suggests that when shards are large enough, the statistical benefit of sharing is minimal — each path has sufficient data to learn independently. But for large numbers of paths (e.g., P = 256), Flat MoE overfits while DiPaCo continues to improve (Table 2): "there is no overfitting at 256 paths with overlapping shards with a 16×16 DiPaCo." The sharing structure becomes critical precisely when the number of paths pushes against the data-per-path limit.

Evidence anchoring the claim. Figure 9 shows monotonic improvement as path-specific modules are added to DiPaCo — for a fixed number of paths, unsharing some modules (increasing capacity) consistently improves perplexity. This demonstrates that the sharing-capacity tradeoff is real and tunable. Table 2 provides the complementary evidence: at 256 paths, Flat MoE (zero sharing) degrades to 14.1 PPL at 10K steps and continues to overfit at 64K steps (reaching only 13.6 PPL even with overlapping shards and early stopping), while the 16 × 16 DiPaCo with partial sharing reaches lower perplexity without overfitting. The shared modules in DiPaCo are each trained on 16× more data than Flat MoE modules (since each is used by 16 paths), giving them statistical strength that prevents overfitting while still allowing the composite model to reach 256× the total parameter count of a single path.

Significance beyond raw performance. This insight reframes modularity from a binary design choice (share or don't share) to a resource allocation problem: given a total parameter budget, a set of compute islands, and a dataset, how should parameters be allocated between shared and path-specific modules to maximize generalization? The answer, per DiPaCo's evidence, depends on the ratio of total parameters to total data and on the semantic structure of the data domains. This connects modular architecture design to long-standing questions in transfer learning, multi-task learning, and domain adaptation — but at a scale (hundreds of paths, billions of parameters) where these questions have not been systematically studied.

Innovation 3: Decoupling Training-Time and Inference-Time Routing Granularity

DiPaCo introduces an architectural principle with broad implications beyond this specific system: the routing granularity used during training need not match the routing granularity used during inference. The paper demonstrates that coarse routing (once per document) during training — necessary for pre-sharding and distributed independent path training — can be compensated for by finer-grained routing (every 64–128 tokens) at inference time, recovering most of the performance gap to dense models without changing the training procedure.

What the field did before. In token-level MoE models (Shazeer et al., 2017; Lepikhin et al., 2021), the routing granularity is identical at training and inference — both operate at the per-token level. The router is trained jointly with the model parameters via gradient descent, and the same router is used at inference. There is no concept of a training-time routing constraint that differs from an inference-time routing capability. In document-level expert models (Gururangan et al., 2023), routing is once per document at both training and inference, and there is no mechanism for finer-grained inference-time routing.

What makes the decoupling distinctive. DiPaCo's key observation is that the training-time routing constraint (coarse, offline, pre-sharded) is motivated entirely by the distribution requirement — it enables independent path training with no per-step communication. But this constraint is irrelevant at inference, where only one path is active per prediction and text can be cheaply re-routed between chunks. This creates an asymmetry: training must be coarse, but inference can be fine.

The paper exploits this asymmetry by training a separate sequence-level router specifically for inference (Section 7.2.2). This router predicts, for each token position, which path would be optimal for the following W tokens. It is trained on a held-out router dataset using path likelihood scores as supervision, independent of the main training loop. At inference, the router can re-route the sequence to a potentially different path every W tokens — W can be chosen independently of the training-time routing granularity.

The magnitude of the effect is the insight. Table 3 shows that routing every 128 tokens instead of once per sequence improves perplexity by 0.74 points for a 16 × 16 DiPaCo — this is the difference between trailing and matching the 1B dense baseline. Routing every 64 tokens adds another 0.10 points, and routing every 16 tokens adds 0.12 more. The marginal benefit of finer granularity diminishes, suggesting that document-level routing captures most of the relevant domain structure, but there is non-trivial within-document variation that coarse routing misses. This within-document variation is what the finer-grained inference router recovers.

Evidence anchoring the claim. Table 3 is the primary evidence, but the implicit evidence is stronger: the entire training procedure uses once-per-document routing, yet the final model quality (after frequent inference routing) matches a 1B dense model that never had any routing constraint. This means the training-time coarse routing did not permanently limit what the model could learn — the paths learned representations that support fine-grained routing decisions even though they were never trained with fine-grained assignments. The discriminative router training (Section 2.4.2) and the token-level router training (Section 7.2.2) are performed on a small held-out set (0.5% of C4) after the paths are already trained, confirming that the path specialization emerges from the coarse training and is sufficient to support fine-grained inference routing.

Significance beyond DiPaCo. This decoupling principle has implications for any system where training and inference face different computational constraints. For example: models trained with block-sparse attention patterns for memory efficiency could use denser attention at inference; models trained with quantization could use higher precision at inference; models trained with truncated backpropagation through time could use longer contexts at inference. The general principle is that training-time approximations motivated by resource constraints need not be baked into the inference procedure, and a lightweight post-hoc adaptation (like DiPaCo's sequence-level router) can recover the quality lost by the approximation. This flips the default assumption — that training and inference must use the same computational graph — and opens a design space where training-time efficiency and inference-time quality are optimized separately.

Innovation 4: A Unified Infrastructure Architecture for Distributed Modular Training

While this innovation is partially about engineering, the paper presents it as a co-designed component whose properties (fault tolerance, heterogeneous device support, elastic scaling) are direct consequences of the modular architecture, not bolted-on afterthoughts. The infrastructure described in Section 3 and Figure 6 is not merely an implementation of the DiPaCo algorithm — it is an argument that modular architectures enable qualitatively different training infrastructure than monolithic models permit.

What the field did before. Standard distributed training infrastructure (data parallelism, model parallelism, pipeline parallelism) assumes all devices are homogeneous, co-located, and available for the entire training duration. Fault tolerance typically involves checkpointing the entire model state and restarting all workers from the last checkpoint — because all workers are interdependent (gradients must be aggregated globally), a single worker failure can stall the entire training run. Heterogeneous hardware is generally not supported because the parallelism strategies assume identical per-device compute and memory.

What makes DiPaCo's infrastructure distinctive. The task queue architecture with independent workers (Section 3.1) is only possible because each path's training is completely independent during the inner optimization phase. There is no per-step synchronization, no gradient all-reduce, and no inter-worker dependency. This enables three capabilities that monolithic training cannot provide:

• Fault tolerance via task re-queueing. If a worker fails or is preempted, the task queue server returns its task to the queue and reassigns it to another worker. No other worker is affected, and no global checkpoint restoration is needed. The paper states this explicitly: "The key advantage of this design is that, thanks to the complete independence among workers, the system can continue making progress even if some workers become unavailable, as long as the worker pool is not empty."

• Heterogeneous device support. Workers can use "heterogeneous types of devices across different regions" because each worker only needs enough compute and memory to train a single path (150M parameters in the experiments). A worker with 8 V100 GPUs and a worker with 16 A100 GPUs can both participate — they just complete their tasks at different speeds, and the task queue implicitly load-balances by giving faster workers more tasks over time.

• Elastic resource utilization. The backup pool design (Section 3.4) — spawning workers on low-priority, preemptible accelerators — means the system can absorb transient compute availability. When spare capacity exists, more workers join the pool and training accelerates. When capacity is reclaimed, tasks are re-queued to the remaining workers. This is a form of elastic scaling that is impossible in tightly synchronized training, where the slowest worker determines the step time and adding/removing workers mid-training requires complex reconfiguration.

These capabilities are not incremental improvements over existing distributed training infrastructure — they represent a qualitatively different operational model. The paper is arguing, implicitly, that the modular architecture enables this infrastructure, and that the infrastructure's properties (fault tolerance, heterogeneity, elasticity) are first-class design goals that should influence architecture choices.

Evidence anchoring the claim. The paper reports training 256 paths with "average time per phase for outer update under 2 minutes" (Section 3.3) using these infrastructure optimizations — sharded outer optimization executors, online parameter gradient averaging, asynchronous checkpoint gathering. This is presented as a scaling result: the infrastructure can handle hundreds of paths without the outer optimization becoming a bottleneck. The backup pool (Section 3.4) is described as enabling training with fewer dedicated workers than paths — "we do multiple rounds of training within an outer iteration step until all paths have been trained" — which is how the system operates when 4,096 GPUs are not simultaneously available.

Significance beyond DiPaCo. This innovation reframes infrastructure as an architectural requirement rather than an implementation detail. If the community moves toward modular, composition-based models, the training infrastructure must support independent module training with asynchronous synchronization — and this infrastructure looks fundamentally different from the synchronous, homogeneous-cluster infrastructure that dominates today. DiPaCo provides a concrete reference design for what such infrastructure requires: task queues, metadata databases for checkpoint tracking, sharded outer optimization, and elastic worker pools. This is a systems contribution that is tightly coupled to the architectural contribution — you cannot have one without the other.

Innovation 5: Verifier-Free Adaptive Routing via Discriminative Re-Sharding

DiPaCo introduces a routing mechanism — discriminative re-sharding — that directly optimizes the routing criterion (which path minimizes perplexity for each document) using the paths' own evaluation scores as supervision, without requiring a separately trained verifier or quality estimator. This is distinct from both learned routing (where a router is trained jointly with the model via gradient descent) and heuristic routing (where hand-crafted or unsupervised features determine assignments).

What the field did before. Token-level MoE models train the router jointly with the model parameters via gradient descent through the gating network (Shazeer et al., 2017; Lepikhin et al., 2021). This works because the router is differentiable and all experts are co-located, so the routing decision can be optimized end-to-end with the language modeling loss. Document-level expert models (Gururangan et al., 2023; Gross et al., 2017) use unsupervised clustering (k-means on tf-idf or neural features) to assign documents to experts — the routing is not optimized for the downstream task, but for feature reconstruction.

DiPaCo's discriminative routing occupies a distinctive middle ground. It is not end-to-end differentiable (the paths are pre-trained, then the router is trained on path evaluation scores, then the paths are re-trained on re-sharded data), so it avoids the co-location requirement of joint training. But it is also not purely unsupervised — it uses the paths' own perplexity scores to optimize the routing assignment, making it task-aware in a way that k-means is not.

What makes it distinctive. The discriminative routing procedure is a form of self-supervised routing: the model evaluates its own performance on held-out data to decide which path should handle which data. This requires no external verifier, no human labels, and no pre-defined domain categories. The supervision signal — "which path achieves lowest perplexity on this document" — is produced by the model itself after an initial generative routing phase.

The key insight is that this self-supervision loop can be iterated. Section 7.2.1 describes alternating between re-training the router (E-step) and re-training the paths on re-sharded data (M-step), showing that performance improves with alternations (Figure 11: 14.0 → 13.38 → 13.36 → 13.25 PPL for phases 0–3 of a 16-path flat MoE). The diminishing returns after the first alternation explain why the main experiments use only one discriminative phase — but the fact that alternations help at all demonstrates that the routing and the model training are mutually informative. Better paths enable better routing; better routing enables better path specialization.

Why this matters beyond DiPaCo. The discriminative re-sharding mechanism is a specific instance of a more general principle: models can self-organize their own modularity through iterative evaluation and reassignment. This has implications for continual learning (new data can be routed to the most appropriate existing module, or trigger creation of a new module), for multi-task learning (tasks can be dynamically assigned to module compositions based on performance), and for model merging (the optimal decomposition of a monolithic model into modules could be discovered by evaluating candidate decompositions on held-out data). The paper does not develop these implications, but the discriminative routing procedure provides a template: train modules, evaluate them on a held-out set to produce target assignments, train a classifier to predict those assignments from input features, and re-shard.

Evidence anchoring the claim. Table 5 shows that discriminative routing yields a 0.7 PPL improvement over generative k-means routing for an 8 × 8 DiPaCo, and Figure 10 shows that this gap widens with more paths (from 8 to 64 paths, the discriminative advantage grows). The bias correction mechanism — training a bias term to match the target document-to-path distribution — is an important practical detail that prevents path starvation, and its necessity emerges from the observation that the logistic regression classifier tends to underestimate the frequency of minority paths.

Significance beyond raw performance. The discriminative routing innovation is significant not because it achieves better perplexity (unsurprising, since it directly optimizes a routing criterion aligned with the LM loss), but because it demonstrates that the routing problem can be solved with self-generated supervision at scale. This removes the need for pre-defined domain ontologies, human-labeled document categories, or differentiable routing networks — all of which are either expensive or incompatible with fully distributed training. The method is simple (linear logistic regression on frozen features), scalable (requiring only a small held-out set for evaluation), and improves with model quality (better paths produce better routing targets). This is a minimal, effective solution to the routing problem that arises naturally from the modular architecture itself.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All experiments use the C4 dataset (Raffel et al., 2020), a large corpus derived from Common Crawl web data. The text is tokenized with a SentencePiece tokenizer (Kudo and Richardson, 2018) using a vocabulary size of 32,000 tokens. The paper evaluates on validation perplexity computed on held-out C4 data, using the first 32 tokens of each sequence for routing (for DiPaCo models) and evaluating perplexity on the remaining tokens, a protocol applied identically to all methods including dense baselines.

• Base model(s). The primary architecture is a 150M-parameter decoder-only transformer with 12 blocks, 896-dimensional hidden states, and 16 attention heads. This serves as both the single-path baseline and as the per-path architecture for all DiPaCo configurations. The main comparison target is a 1.3B-parameter dense transformer with 24 blocks, 2048-dimensional hidden states, and 16 heads. A 1B-parameter dense model also appears in the final evaluation comparison (Table 3: 11.41 validation perplexity). The 150M size is chosen to be representative of a model that fits comfortably on a modest compute island, while the 1.3B model represents a monolith that requires substantially more co-located resources.

• Metrics. The primary metric is validation perplexity (PPL) — the exponentiated average negative log-likelihood per token — computed on held-out C4 data using sequences of 2,048 tokens during evaluation (vs. 1,024 tokens during training). Lower perplexity is better. The paper reports perplexity against number of weight update steps used at training time, arguing that this is a "close proxy for wall-clock time if all computations are done on the same accelerator type" (Section 4). All DiPaCo paths are trained for 88,000 weight update steps total across all inner optimization phases, matching the step count of the dense baselines. However, the paper explicitly warns: "this comparison is not standard in the literature, as weight updates for DiPaCo see more tokens and use more FLOPs when the number of paths is larger" — a DiPaCo with 256 paths processes 256 shards in parallel, so each training step represents 256× more tokens processed and 256× more FLOPs consumed than a single dense model step.

• Baselines. The paper compares against four distinct baselines, each cited from prior work:

  • 150M dense model: A 150M-parameter transformer trained for 88,000 steps on the full C4 dataset. This represents the performance achievable with the same per-path compute budget but no modularity or distribution.
  • 1.3B dense model: A 1.3B-parameter transformer trained for 88,000 steps on the full C4 dataset. This is the primary quality target — the paper aims to match or approach this model's perplexity using distributed modular training.
  • 1B dense model: Referenced in Table 3 as achieving 11.41 validation perplexity, used as the comparison point for the test-time frequent routing results (where 150M-per-path DiPaCo matches it with 6× fewer parameters per forward pass).
  • DiLoCo (Douillard et al., 2023): The base DiLoCo algorithm applied to a dense 150M model — all workers share identical parameters and synchronize via the same outer optimization procedure, but there is no modularity or routing. This isolates the benefit of DiLoCo's infrequent communication from the benefit of DiPaCo's modular architecture.
  • Flat MoE: A mixture of completely independent 150M-parameter paths with no parameter sharing (equivalent to Gross et al., 2017; Gururangan et al., 2023). This isolates the benefit of module sharing from the benefit of simply training more parameters independently.

• Generation budget / compute accounting. The paper's compute accounting is unconventional and requires careful interpretation. The x-axis in the main convergence plots (Figures 8, 9) shows "number of weight update steps" — the total number of inner optimization steps each path executes. For the dense baselines, this is standard: one step = one batch processed. For DiPaCo, one step means one step of AdamW on one path's shard, and 256 paths each perform this step independently in parallel (for 16 × 16 DiPaCo). Therefore, per "step" on the x-axis, DiPaCo processes approximately 256× more tokens and consumes 256× more FLOPs than the dense 150M baseline. The paper does not plot FLOP-matched curves; it instead argues that wall-clock time is roughly equivalent between the 150M dense baseline and DiPaCo (since the paths train in parallel), while the 1.3B dense model is 45% slower in wall-clock time (Section 4.1). This means DiPaCo's performance advantage is in wall-clock efficiency (matching 1.3B quality in less time), not in total compute efficiency (which would favor dense models due to their more efficient use of parameters). The inner optimization uses a batch size of 512 and sequence length of 1,024 tokens (training) or 2,048 tokens (evaluation). The inner optimizer is AdamW with peak learning rate 4 × 10⁻⁴ and cosine schedule; the outer optimizer is Nesterov momentum with learning rate 0.7 and momentum 0.9.

• Cross-validation / statistical protocol. The paper reports single-run results without cross-validation or statistical error bars. There is no mention of multiple random seeds, confidence intervals, or significance testing. This is typical for large-scale language model training papers (where training multiple replicates is prohibitively expensive) but limits the ability to assess whether differences between configurations are statistically reliable or within run-to-run variance. The discriminative routing uses a held-out "router dataset" (0.5% of C4, Section 7.2.1) to train the router classifier, which is a form of validation set for the routing component but does not constitute cross-validation of the overall training procedure.

Main Quantitative Results

Convergence vs. 1.3B Dense Baseline (Figure 8, Section 4.1)

The headline result is that a 16 × 16 DiPaCo with 256 paths of 150M parameters each — trained with one discriminative routing phase and document-level routing during training — closely approaches the performance of a 1.3B dense transformer while reducing wall-clock training time by approximately 45%. The convergence plot in Figure 8 shows three curves:

• 150M dense baseline (purple): Starting from a pretrained 150M model (trained for 24K steps, reaching approximately 18–19 PPL), this model continues training and converges to roughly 14–15 PPL at 88K total steps. This is the "you didn't add any modularity or distribution" lower bound.

• 16 × 16 DiPaCo with document-level routing (red, "1x routing"): This configuration, which makes one routing decision per document at both training and test time, substantially outperforms the 150M dense baseline, reaching approximately 12.4 PPL at 88K steps. The improvement over the 150M baseline demonstrates that the modular architecture with routing and parameter sharing successfully leverages the additional capacity (256 paths × 150M parameters = 38.4B total parameters, though only 150M are activated per forward pass) and the distributed data (each path sees a specialized shard).

• 1.3B dense baseline (orange): Trained for 88K steps on the full C4 dataset, reaching approximately 11.4 PPL. This is the target performance.

The 16 × 16 DiPaCo at 12.4 PPL leaves a gap of approximately 1.0 perplexity to the 1.3B dense baseline. The paper closes this gap through two mechanisms described subsequently:

1. Early stopping (path-specific): Reduces the gap modestly (Table 3: "once per sequence" PPL drops from 12.39 to 12.22 with early stopping).
2. Frequent test-time routing (Table 3): The largest gain. Routing every 128 tokens instead of once per sequence reduces perplexity from 12.22 to 11.48. Routing every 64 tokens further reduces to 11.38, matching the 1B dense model's 11.41 PPL. The full sequence of improvements is documented in Table 3 and represents a cumulative reduction of 1.01 perplexity points from the document-level routing baseline, bringing DiPaCo's performance to within 0.03 PPL of the 1.3B dense baseline (if one extrapolates the trend from the explicit 1B comparison in the paper text) or matching the explicitly stated 1B baseline.

Critically, the paper notes that "DiPaCo with 256 paths of size 150M partially sharing parameters, and one routing decision per document, is nearly sufficient to match the performance of a dense model of size 1.3B" — the gap is "nearly" closed before frequent routing, and fully closed with frequent routing at test time.

Wall-clock comparison. The paper states that DiPaCo's training "wall-clock time is 45% less than the 1.3B dense counterpart, and roughly equivalent to the 150M parameter single-path counterpart" (Section 4). This is because the 256 paths train in parallel on 256 independent compute islands, and the outer optimization overhead is kept under 2 minutes per phase (Section 3.3). The dense 1.3B model, by contrast, requires all devices to be co-located and synchronized at every step, making each step slower and the total training wall-clock longer.

Token and FLOP comparison. Crucially, this advantage is in wall-clock time, not total FLOPs or tokens. The paper is explicit: "this comparison is not standard in the literature, as weight updates for DiPaCo see more tokens and use more FLOPs when the number of paths is larger" (Section 4.1). A 16 × 16 DiPaCo at step 88,000 processes approximately 256× more tokens than the 150M dense baseline at the same step count — each path processes its own shard of the data, and the total tokens across all shards far exceeds what the dense model sees. DiPaCo is resource-intensive in total compute but efficient in wall-clock time, which is the intended tradeoff given the paper's assumption that "training compute (FLOPs) is relatively cheap" while "communication is relatively expensive" (Section 2.1).

Scaling Number of Paths and Parameters (Figure 9, Section 4.2)

Figure 9 investigates how DiPaCo's validation perplexity improves as the number of paths and total parameters increase, holding the per-path size constant at 150M parameters. The x-axis is training steps; the y-axis is validation perplexity. The figure shows a clear monotonic improvement with more paths:

• 2 × 4 DiPaCo (8 paths, L = 2, K₁ = 2, K₂ = 4): Reaches approximately 14.0 PPL at convergence.
• 4 × 4 DiPaCo (16 paths): Reaches approximately 13.5 PPL.
• 8 × 8 DiPaCo (64 paths): Reaches approximately 12.8 PPL.
• 16 × 16 DiPaCo (256 paths): Reaches approximately 12.4 PPL (before frequent routing and early stopping).

All DiPaCo variants substantially outperform the 150M dense baseline (which plateaus around 14–15 PPL) and approach the 1.3B dense baseline (around 11.4 PPL) as the number of paths grows to 256.

The paper additionally experiments with path-specific modules (Section 2.6.1) — modules at certain levels that are not shared across paths — to increase total parameter count without increasing the number of paths. Figure 9 includes curves for DiPaCo configurations with path-specific modules ("with path-specific modules" in the legend), showing that for a given number of paths, unsharing some modules further improves perplexity. The specific configuration uses path-specific modules at transformer blocks 0, 5, 6, 11 and the embedding matrix, while other blocks are shared (Section 4.2: "For all experiments with path-specific modules, the transformer blocks 0, 5, 6, 11, and the embedding matrix are not communicated across paths"). The precise perplexity values for these configurations at convergence are not enumerated in the text, but the curves in Figure 9 show visible downward shifts relative to the fully-shared versions.

In all cases, the inner optimization consists of τ = 150 steps per outer phase. The fact that performance improves monotonically with both more paths and more path-specific modules demonstrates that DiPaCo's quality scales with total parameter count (both through more paths and through less sharing), even though each path remains fixed at 150M parameters and the communication frequency remains low.

Parameter Sharing: DiPaCo vs. Flat MoE vs. DiLoCo (Table 1, Section 4.3)

Table 1 compares three approaches at the same wall-clock time and total inner optimization steps, isolating the effect of parameter sharing:

• DiLoCo (Douillard et al., 2023): All paths share all parameters — effectively a single 150M model trained with DiLoCo's outer optimization across multiple workers, each seeing a different data shard. This model sees the same total FLOPs and tokens as the larger DiPaCo models (since it processes data from all shards) but has only 150M parameters. It reaches approximately 13.8 PPL — better than the standard 150M baseline (which only sees the full unshuffled C4 dataset once) but substantially worse than any modular DiPaCo variant. The interpretation: the 150M model lacks the capacity to absorb the additional tokens and FLOPs; adding more parameters (via modular paths) is necessary to leverage the distributed data.

• Flat MoE (P = 64 fully independent paths): Each of the 64 paths is a completely independent 150M model trained on its own shard with no parameter sharing. This reaches approximately 12.5 PPL — competitive with the 16 × 16 DiPaCo (256 paths with partial sharing, ~12.4 PPL) despite using fewer total paths. The paper interprets this as: when shards are large enough (64 shards out of the total C4 dataset), independent experts can specialize effectively without needing shared parameters to transfer learning. The absence of sharing is not harmful because each expert sees sufficient data.

• 16 × 16 DiPaCo (256 paths with partial sharing): Reaches approximately 12.4 PPL, as in Figure 8.

The key finding is that for P = 64, Flat MoE (fully independent) nearly matches the compositional DiPaCo, suggesting that parameter sharing is most beneficial when the number of paths is large relative to the data. When each shard is small, sharing modules across paths provides a form of data pooling — a module used by 16 paths is trained on 16× more tokens than a module used by only one path.

Table 2 confirms this pattern. When the number of paths increases to 256 with Flat MoE (fully independent), validation perplexity degrades sharply: the model reaches 14.1 PPL after only 10K steps and continues to overfit as training progresses. Even with top-2 overlapping shards and early stopping, the Flat MoE reaches only 13.6 PPL at 64K steps, while the 16 × 16 DiPaCo with overlapping shards "has no overfitting" (Section 4.3) at 256 paths. The shared modules in the 16 × 16 DiPaCo effectively see 16 × (shard_size) tokens, giving them sufficient statistical strength, while the Flat MoE modules each see only shard_size tokens, which is too small for 256-way specialization.

Table 5 further isolates routing quality, showing that for an 8 × 8 DiPaCo (P = 64), discriminative routing yields 0.7 perplexity improvement over generative k-means routing (absolute values: roughly 13.5 vs. 14.2 PPL, interpolating from the paper's discussion). This confirms that task-aware routing (assigning documents to the path that models them best) is superior to feature-similarity-based routing, and the gap widens with more paths (Figure 10).

Routing Frequency at Evaluation (Table 3, Section 4.4)

The paper's most practically significant result for deployment is in Table 3: for a 16 × 16 DiPaCo (P = 256), routing more frequently at test time progressively closes the gap to dense models:

The reference comparison: the 1B dense model achieves 11.41 PPL (explicitly stated in Table 3's caption: "we match the performance of a dense 1B model (11.41)"). The 150M-per-path DiPaCo matches the 1B dense model when re-routing every 64 tokens, using over 6× fewer parameters per forward pass (150M vs. 1B). Further re-routing (every 16 tokens) reduces perplexity to 11.26, approaching but not quite reaching the 1.3B dense model (approximately 11.4 in Figure 8).

Several observations from these numbers:

• The largest single gain comes from moving from once-per-sequence to every-128-tokens re-routing: 12.22 → 11.48, a reduction of 0.74 perplexity. This suggests that the dominant within-document variation in optimal path assignment occurs at a granularity of roughly 128 tokens — consistent with topical shifts within multi-paragraph documents.
• Gains diminish with finer re-routing: 0.10 PPL from 128→64 tokens, 0.07 from 64→32, 0.05 from 32→16. This diminishing returns pattern suggests that most of the benefit of adaptive routing is captured at moderate granularity (64–128 tokens).
• The gap from the once-per-sequence baseline (12.39) to the every-64-tokens result (11.38) is 1.01 PPL — this is the total performance recovered by decoupling training-time routing (coarse) from inference-time routing (fine). This is a striking demonstration of the decoupling principle (Innovation 3 in Section 4 of this analysis).

The paper explicitly notes that the model "may choose to route to the exact same path that it selected previously" when re-routing — the frequent routing does not force path changes; it provides the option. The router makes independent decisions at each re-routing boundary based on the preceding token window, and many consecutive windows may naturally select the same path if the text remains topically coherent.

Table 3 also shows the effect of early stopping alone: moving from 12.39 to 12.22, a 0.17 PPL improvement. This is a smaller effect than frequent routing but still non-trivial. The combination of early stopping and routing every 128 tokens produces the majority of the total gain (12.39 → 11.48).

Optimizer Comparison: DiLoCo vs. Full Synchronization (Section 4.5)

The paper runs a critical ablation: for DiPaCo architectures at various scales, compare training with DiLoCo's partial synchronization (communicating every τ = 150 steps) against fully synchronous training where all paths compute gradients on their own data shards, gradients are aggregated module-by-module at every step, and a single AdamW step is taken with the aggregated gradient. This is standard distributed data-parallel training adapted to the DiPaCo architecture.

The results (Section 4.5):

• 2 × 2 DiPaCo (P = 4): DiLoCo outperforms full synchronization by 0.3 PPL.
• 4 × 4 DiPaCo (P = 16): DiLoCo outperforms by 0.6 PPL.
• 8 × 8 DiPaCo (P = 64): Full synchronization is only 0.1 PPL better, despite communicating hundreds of times more frequently.

These are exact numbers from the text: "DiPaCo trained with DiLoCo slightly outperforms their fully-synchronously-trained version by 0.3 and 0.6 perplexity points when using a 2×2 and 4×4 architecture, respectively. At 8×8 DiPaCo trained fully synchronously reaches better perplexity by only 0.1 perplexity."

This is a genuinely surprising result. The expectation would be that more frequent communication — providing each path with up-to-date information about what other paths are learning — should always improve optimization, even if marginally. That DiLoCo outperforms full synchronization for smaller path counts suggests that the infrequent communication provides a form of beneficial regularization: by allowing each path to train independently for 150 steps before being pulled back toward the global consensus, the model explores more diverse specializations that collectively improve generalization. The fact that full synchronization barely wins at P = 64 (and by only 0.1 PPL) strongly supports the paper's claim that DiLoCo is "an effective distributed optimization algorithm for DiPaCo" (Section 4.5) — adding hundreds of times more communication provides almost no benefit, confirming that the DiPaCo architecture is intrinsically well-suited to low-communication training.

Caveat. These experiments were run in the fully collocated setting (Section 4.5: "In the setting where all devices are actually collocated and there is no constraint in terms of communication"), meaning the DiLoCo variant did not experience the communication delays that would motivate it in practice. The fact that it matches or exceeds full synchronization even without those constraints is stronger evidence for its effectiveness.

Ablation Studies and Robustness Checks

• Inner optimizer steps per outer phase (τ): The paper fixes τ = 150 for most experiments but does not ablate this value. The original DiLoCo paper (Douillard et al., 2023) showed that τ can range from hundreds to thousands of steps with stable training, but DiPaCo does not explore whether more or fewer inner steps would improve performance. This is a notable omission, as τ is the primary knob controlling the communication-to-computation ratio.

• Number of discriminative routing alternations (Figure 11): The paper investigates whether alternating between router training and path training multiple times (full EM-style alternation) continues to improve performance. For a 16-path flat MoE, three discriminative phases produce diminishing returns: PPL improves from 14.0 (phase 0, generative k-means) → 13.38 (phase 1, first discriminative) → 13.36 (phase 2) → 13.25 (phase 3). The gain from phase 0 to phase 1 is 0.62 PPL; phase 1 to phase 2 is only 0.02; phase 2 to phase 3 is 0.11. The paper concludes that "one step of training the discriminative router leads to a significant improvement in validation perplexity, more alternating steps lead to further minor improvements" (Section 7.2.1). All main experiments use exactly one discriminative phase, which captures the majority of the benefit while minimizing computational overhead.

• Sharding method comparison (Table 5): For an 8 × 8 DiPaCo (P = 64), generative routing (product k-means) vs. discriminative routing yields a 0.7 PPL difference in favor of discriminative routing. "Longer training runs yield even greater gain from the discriminative gating." This confirms that task-aware routing matters, and the effect compounds with training duration — as paths become more specialized, the optimal assignment becomes more discriminable.

• Overlapping shards (Section 2.4.4, Table 2): The 16 × 16 DiPaCo uses top-2 overlapping shards at training time: each document appears in two shards (its top-2 closest clusters). The paper notes that this "does not make training slower in wallclock or FLOPs, although it does increase the size of the shards" (Section 4). For Flat MoE at 256 paths, overlapping shards with early stopping improves PPL from 14.1 (at 10K steps, degrading further) to 13.6 (at 64K steps, but still overfitting). The overlapping shards mitigate overfitting by giving each path more training data, but the core issue — each expert sees too little data — remains for Flat MoE at large path counts.

• Path-specific modules without outer optimization: The paper states (Section 2.6.1, footnote) that for path-specific modules (where P_{l,e} = 1), the outer optimization step is still applied despite there being no gradient averaging: "we still apply the outer optimization of Algorithm 1 because empirically it improves convergence over the default optimizer." This is an empirical finding — Nesterov momentum applied to the single-path outer gradient provides better convergence than simply keeping the AdamW-trained local parameters. The paper does not provide ablation results quantifying how much this improves over skipping outer optimization entirely for path-specific modules.

• Backup pool and elastic scaling: The infrastructure section describes the backup pool design (Section 3.4) — spawning workers on low-priority preemptible accelerators to maximize training throughput — and notes that this design enables training with fewer dedicated workers than paths. While not a model performance ablation, this is a systems-level robustness check: the paper demonstrates that the training procedure tolerates worker heterogeneity, preemption, and dynamic pool sizing, which are essential for the "poorly connected and heterogeneous workers" setting that motivates the work.

• No learning rate or optimizer hyperparameter sensitivity analysis: The paper states (Section 4) that "we have searched over relatively few hyper-parameters: mainly learning rate and value of Nesterov momentum." No learning rate sweep results, sensitivity curves, or failure modes are reported. The outer optimizer uses fixed values (outer LR = 0.7, outer momentum = 0.9) following the DiLoCo recipe. Whether DiPaCo is sensitive to these choices — particularly the outer learning rate, which controls how strongly the global parameters are pulled toward the average local update — is unexplored.

• No path sampling ablation: The paper proposes (Section 2.6.2) that when the number of paths exceeds available devices, one could sample a subset of paths to train at each outer phase. This is not implemented or evaluated. The scaling behavior of DiPaCo under path subsampling — which would further reduce communication and device requirements — remains unknown.

• No FLOPs-matched comparison against dense baselines (only wall-clock and step-matched): All comparisons use step-matched or wall-clock-matched baselines, never FLOPs-matched. A DiPaCo with 256 paths at step t consumes roughly 256× more training FLOPs than the 150M dense baseline at the same step count. The paper does not answer: if the 150M or 1.3B dense models were given the same total FLOPs as DiPaCo (by training for more steps or with larger batches), how would they perform? Section 6 explicitly acknowledges this: "DiPaCo is significantly less FLOP efficient per evaluation perplexity than a standard dense compute optimal model." This means the claimed wall-clock advantage must be weighed against the substantially higher total energy and compute cost, which may or may not be acceptable depending on the deployment context.

Critical Assessment

The paper makes three central claims that should be evaluated against the experimental evidence:

Claim 1: "DiPaCo exceeds the performance of a 1 billion-parameter dense transformer language model by choosing one of 256 possible paths, each with a size of 150 million parameters" (Abstract).

This claim is supported, but with important qualifications. Table 3 shows that a 16 × 16 DiPaCo achieves 11.38 PPL when routing every 64 tokens at test time, matching the 1B dense model's 11.41 PPL. The 16 × 16 DiPaCo uses 150M parameters per forward pass (since only one path is active), comparing favorably to the 1B dense model's 1B parameters per forward pass. However:

• Matching requires test-time routing every 64 tokens. Without frequent test-time routing, DiPaCo achieves only 12.22 PPL (with early stopping, once-per-sequence routing), which is substantially worse than the 1B model. The claim as stated in the abstract silently includes the frequent test-time routing optimization, which is not implied by the phrase "by choosing one of 256 possible paths." In fact, the model does choose one path, but it re-chooses every 64 tokens within a sequence, not once per input. This is a meaningful qualification for deployment, as Section 6 acknowledges that frequent re-routing requires re-computing the KV-cache for each newly selected path, which has latency and memory implications.
• The comparison is against a 1B model, not the 1.3B model that DiPaCo was primarily designed to match. The 1.3B dense model achieves lower perplexity (approximately 11.4, interpolating from Figure 8) than the 1B model (11.41). DiPaCo approaches but does not clearly exceed the 1.3B model even with 16-token re-routing (11.26 PPL, which is lower/better than 11.4 but the paper does not explicitly claim superiority).
• FLOP efficiency is not part of the claim. DiPaCo consumed substantially more total training FLOPs than the 1B dense model to achieve this result (since 256 paths each processed their shards). The abstract's phrasing — "for the same amount of training steps but less wall-clock time" — carefully frames the comparison in wall-clock terms. This is honest but easily misread as claiming total compute parity.

Claim 2: "Our approach facilitates training across poorly connected and heterogeneous workers, with a design that ensures robustness to worker failures and preemptions" (Abstract).

This claim is partially supported by the infrastructure design (Section 3) and the reported operational experience ("the average time per phase for outer update under 2 minutes," Section 3.3; the backup pool design enabling elastic scaling, Section 3.4). However, the experimental evidence is qualitative and anecdotal, not quantitative. The paper does not report:

• Measured communication bandwidth between workers in the experiments — were workers actually "poorly connected" or were they in standard datacenter configurations?
• Failure rates, recovery times, or training progress under preemption — the backup pool design is described but no data shows how it performs under realistic failure scenarios.
• Training with deliberately heterogeneous hardware — the paper mentions support for "heterogeneous types of devices across different regions" but does not report experiments where, e.g., half the paths train on GPUs and half on TPUs.
• Scaling to truly wide-area (cross-continental) communication — the Effingo process for asynchronous checkpoint gathering is described, but whether training actually occurred across geographically distant sites (as suggested by the "across the world" language in Figure 1 and the introduction) is not reported.

These omissions are understandable — this is a first prototype, and rigorous fault-tolerance benchmarking would be a separate systems paper — but they mean the robustness claim remains an architectural aspiration supported by design arguments rather than empirical demonstration.

Claim 3: "At inference time, only a single path needs to be executed for each input, without the need for any model compression" (Abstract).

This is true for the per-token or per-64-token forward pass: only one 150M-parameter path processes each chunk. It is not true for the full sequence: Table 3 shows that to achieve competitive perplexity, the model re-routes every 64 tokens, potentially activating a different path for each chunk. The total parameters used to process a full sequence could be up to sequence_length / 64 different 150M-parameter paths, though in practice many consecutive chunks likely route to the same path. The inference cost per token is roughly equivalent to a 150M model (since only one path runs per token), but the system must have all selected paths available — if every chunk routes to a different path, up to 32 paths (for a 2,048-token sequence with 64-token re-routing) must be loaded from storage, which has memory and I/O implications the paper does not model.

Methodological strengths:

• The experiments cleanly isolate the effects of modularity (DiPaCo vs. DiLoCo), parameter sharing (DiPaCo vs. Flat MoE), routing quality (discriminative vs. generative), and test-time routing granularity (Table 3 sweep from once-per-sequence to every-16-tokens).
• The counterintuitive result that DiLoCo matches or outperforms full synchronization (Section 4.5) is robustly demonstrated across three architectural scales.
• The overfitting analysis for Flat MoE (Table 2) clearly shows the boundary condition where parameter sharing becomes necessary, providing actionable guidance for practitioners.
• The paper is transparent about the unconventional compute accounting (Section 4: "this comparison is not standard in the literature") and explicit about FLOP inefficiency (Section 6).

Methodological weaknesses and missing experiments:

• No FLOPs-matched comparison. The most significant gap. All claims of efficiency are in wall-clock terms, but total FLOPs consumed by DiPaCo are 1–2 orders of magnitude higher than the dense baselines at the same x-axis position. The paper acknowledges this limitation (Section 6) but does not attempt any FLOP-efficiency optimization, leaving open the question of whether DiPaCo-like architectures can ever be competitive on FLOPs-per-perplexity with compute-optimal dense training.

• Single dataset (C4) and single scale (150M per path). All experiments use C4 for language modeling at the 150M-per-path scale. The paper does not investigate whether the benefits of modular training are specific to the domain diversity of C4 (which naturally creates domain-specific shards via routing) or generalize to other tasks and datasets. The proposed scaling to 32 × 32 × 32 DiPaCo (32,768 paths) with path sampling is not implemented.

• No latency measurements for inference with frequent re-routing. Table 3 shows perplexity improvements from re-routing every 64 tokens, but the latency cost of loading new paths (potentially from different workers), re-computing KV-caches, and communicating between the router and path workers is not measured or modeled. Section 6 acknowledges that "if we were to naively route more frequently during deployment … we would need to re-compute the KV-cache for each query after each re-route," but no latency data is provided.

• No comparison against token-level MoE at equivalent total parameters. The paper contrasts DiPaCo against token-level MoE conceptually (Section 5.2: "require even more co-located accelerators than the equivalent-activated dense model") but does not train a token-level MoE baseline for C4 at comparable scale (e.g., a 38B total parameter token-MoE with 150M activated parameters per token). This would be the most direct architectural comparison — both are sparsely activated, but token-MoE uses per-token routing with full synchronization while DiPaCo uses document-level routing with infrequent communication. Without this comparison, we cannot assess whether DiPaCo's communication efficiency comes at a cost in quality relative to the current state-of-the-art in sparse architectures.

• No hyperparameter sensitivity analysis. The learning rate, τ, outer optimizer settings, and batch size are fixed and described as "searched over relatively few hyper-parameters" (Section 4). Whether DiPaCo is robust to these choices or requires careful tuning per architecture configuration is unknown.

• No investigation of when DiPaCo breaks. The paper shows a clean scaling trend from 8 to 256 paths, but does not push to failure modes — how many paths before routing quality degrades? How small can shards become before overfitting overwhelms the benefit of shared modules? How large can τ be before the outer gradient averaging becomes incoherent? These failure boundaries would be informative for practitioners.

Experiments that would have strengthened the paper:

1. A FLOPs-matched comparison: train the 1.3B dense model for enough additional steps to consume the same total FLOPs as the 16 × 16 DiPaCo, and compare perplexity. This would reveal whether wall-clock advantage survives when total compute is held constant.

2. A token-level MoE baseline at equivalent scale (total parameters, activated parameters, training FLOPs) to quantify the quality cost of DiPaCo's reduced communication.

3. Training on a dataset without natural domain structure (e.g., a uniformly random subset of C4 with the same token count but no topical coherence) to test whether DiPaCo's routing advantage depends on the existence of meaningful domain clusters.

4. Ablation of τ (inner steps per outer phase): sweep from τ = 10 to τ = 1000 to find the optimal communication frequency for DiPaCo architectures of different sizes. This is the primary knob controlling the communication-computation tradeoff.

5. Inference latency benchmarking with frequent re-routing: measure end-to-end time-per-token for DiPaCo with re-routing every 64 tokens vs. a dense 1B model, including path loading and KV-cache recomputation costs.

Overall assessment: The experiments convincingly demonstrate that the DiPaCo architecture, combined with coarse routing and DiLoCo optimization, can train a large modular model across independent compute islands and achieve competitive quality (in wall-clock time) with a much larger dense model. The scaling trends (more paths → better perplexity, frequent test-time routing → closes the gap) are clear and well-ablated. However, the paper's contributions are primarily about a direction — modular, communication-efficient distributed training — rather than about achieving a Pareto-optimal point in the quality-cost space. The FLOP inefficiency and lack of latency measurements for inference mean that DiPaCo, in its current form, is not directly competitive with dense training for applications where total compute cost or inference latency dominate. The paper is honest about these limitations and frames itself as a prototype, which is appropriate. The strongest and most robust findings are the difficulty of scaling fully independent experts (Flat MoE overfitting) and the effectiveness of DiLoCo for modular architectures (matching full synchronization), both of which have clear practical implications for anyone attempting modular distributed training.

6. Limitations and Trade-offs

6.1 FLOP Efficiency: DiPaCo Is Substantially More Compute-Intensive Than Dense Training

The assumption or constraint. DiPaCo is designed under the explicit assumption that "training compute (FLOPs) is relatively cheap" while "communication is relatively expensive" (Section 2.1). The paper acknowledges this assumption is "not realistic in the current ML training paradigm" but argues it may become realistic if model sizes grow beyond what can be co-located. The consequence is that the system was never optimized for FLOP efficiency, and the paper is transparent about the result:

"In this work we made no effort to optimize its FLOP efficiency; and in the results presented, DiPaCo is significantly less FLOP efficient per evaluation perplexity than a standard dense compute optimal model." (Section 6)

The consequence. Every comparison in the paper is wall-clock-matched or step-matched, never FLOPs-matched. A 16 × 16 DiPaCo at training step t processes approximately 256× more tokens and consumes roughly 256× more FLOPs than the 150M dense baseline at the same step count (each of the 256 paths processes its own full shard of the C4 dataset independently). The headline claim that DiPaCo "exceeds the performance of a 1 billion-parameter dense transformer" (Abstract) uses 45% less wall-clock time but almost certainly uses an order of magnitude more total compute. The paper does not answer the obvious counterfactual: if the 1.3B dense model were given the same total FLOPs budget as DiPaCo (by training longer or with larger batches), would it outperform DiPaCo? The paper also does not compare against a token-level MoE (Switch Transformer or GShard) at equivalent scale, which achieves better training FLOP efficiency than dense models while activating similar per-token parameter counts. A practitioner evaluating whether to adopt DiPaCo must weigh the wall-clock savings against the substantially higher total energy cost and hardware time — a tradeoff that the paper's "compute is cheap" assumption sidesteps but that most real deployments cannot ignore.

What evidence exists in the paper. Section 4 explicitly warns: "this comparison is not standard in the literature, as weight updates for DiPaCo see more tokens and use more FLOPs when the number of paths is larger." Section 6 reiterates: "DiPaCo is significantly less FLOP efficient per evaluation perplexity than a standard dense compute optimal model." However, no FLOPs-matched comparison appears anywhere in the paper — not against dense baselines, not against token-level MoE, not even as a projected calculation. Figure 8 and Figure 9 use weight update steps on the x-axis, which understates DiPaCo's compute consumption by a factor of P (the number of paths) relative to the 150M baseline. Table 1, which compares DiPaCo against DiLoCo and Flat MoE, is at matched wall-clock time but unmatched FLOPs. The paper provides no information about how DiPaCo's perplexity-vs-FLOPs curve compares to alternatives.

Mitigation status. The paper acknowledges this as "the most salient limitation" (Section 6, first sentence) and suggests that "there are several straightforward design choices that could increase FLOP efficiency, for example allowing some paths to co-locate" (Section 7). However, no experiments explore these mitigations, and no FLOPs numbers are reported even for the existing system. The limitation is acknowledged but entirely unaddressed — it is deferred to future work. The paper frames this as acceptable because it is a prototype aimed at demonstrating a new paradigm, not a production-ready system. For a practitioner, this means DiPaCo in its current form would likely cost substantially more in total compute (and thus energy and hardware budget) than training a dense model of equivalent quality, even if it finishes faster in wall-clock time. Whether that tradeoff is acceptable depends on whether the deployment scenario values speed over total resource consumption — a question the paper does not help answer since it provides no FLOPs data.

6.2 Inference Latency: Frequent Test-Time Routing Incurs Unmodeled Overhead

The assumption or constraint. The paper's most practical result — matching a 1B dense model using 150M-per-path DiPaCo with over 6× fewer parameters per forward pass — depends on re-routing the sequence to a potentially different path every 64 tokens at test time (Table 3). The paper acknowledges a specific cost of this design:

"if we were to naively route more frequently during deployment … we would need to re-compute the KV-cache for each query after each re-route." (Section 6)

However, the paper provides no latency measurements, no memory modeling, and no analysis of how routing frequency affects end-to-end inference time.

The consequence. The KV-cache in a transformer stores the key and value vectors for all previously processed tokens, enabling autoregressive generation without recomputing attention over the full history at each step. When DiPaCo re-routes to a different path, the new path has not seen the previous tokens and has no cached keys/values for them. Computing the KV-cache from scratch for each re-routed chunk means re-processing the entire prefix through the new path — for a 2,048-token sequence with re-routing every 64 tokens and conservative path changes, this could mean re-computing attention over thousands of token pairs multiple times. The total FLOPs per sequence during inference could be substantially higher than the per-token parameter count suggests (since the 150M model is run multiple times with overlapping computation), and the wall-clock latency could be dominated by KV-cache recomputation rather than the forward pass itself.

Additionally, even in scoring mode (used for perplexity evaluation in this paper, as opposed to autoregressive generation), the need to potentially switch paths and load their parameters from storage introduces I/O overhead. If paths are served from different workers (as the distributed deployment vision suggests), each re-route incurs network communication to send the text chunk to the new worker and receive the output. The paper provides none of these measurements. The "6× fewer parameters per forward pass" figure (Section 4.4) accounts only for the parameters actively computing one chunk; it does not account for the cost of loading different paths, recomputing KV-caches, or communicating between router and workers.

What evidence exists in the paper. Table 3 shows perplexity improvements from finer-grained re-routing, establishing that re-routing every 64 tokens is necessary to match dense model quality — but reports no latency or throughput data. Section 6 explicitly names the KV-cache recomputation problem. Section 4.4 notes that "the cost of switching paths is small, since the router is run infrequently and in scoring (as opposed to auto-regressive generation) mode" — but this claim is qualitative, not quantitative, and does not address the KV-cache issue. The paper evaluates in scoring mode (computing perplexity over pre-existing text) rather than generation mode, which is a less latency-sensitive setting. No experiment measures wall-clock inference time for sequences with re-routing.

Mitigation status. The paper does not attempt to mitigate this limitation beyond acknowledging it. Potential solutions — such as sharing KV-caches across paths (if early layers are shared), caching recent KV states for frequently-used paths, or limiting re-routing to path transitions where the same path is selected — are not explored. The limitation is entirely deferred to future work. A practitioner deploying DiPaCo for latency-sensitive inference (e.g., interactive chat, real-time translation) would need to solve the KV-cache recomputation problem or accept substantially higher latency than the per-forward-pass parameter count implies. For batch scoring applications (like the perplexity evaluation used in the paper), this limitation is less severe but still means total inference compute is higher than a naive parameter-count comparison suggests.

6.3 Narrow Empirical Validation: Single Dataset, Single Scale, Single Model Family

The assumption or constraint. All experiments use a single dataset (C4) for a single task (language modeling evaluated by perplexity) at a single per-path scale (150M parameters) within a single model family (decoder-only transformers with a specific architecture). The paper does not study scaling laws, does not vary the per-path model size, and does not test on tasks other than language modeling perplexity. The authors acknowledge:

"This work also does not study scaling laws of DiPaCo. We work at a single path-size scale, and on one dataset." (Section 6)

The consequence. Several important questions about DiPaCo's generality are left unanswered:

• Domain structure dependence. C4 is a diverse web crawl corpus containing documents from many domains (news, forums, academic papers, code, etc.). The routing mechanism — whether generative k-means or discriminative — relies on there being meaningful topical or stylistic clusters in the data that paths can specialize to. On a dataset without such structure (e.g., a uniformly homogeneous corpus, or a task like mathematical reasoning where all problems share similar surface forms but differ in reasoning depth), the routing might degenerate — either all paths would receive similar data (making the modular architecture redundant) or the router would make arbitrary assignments (leading to worse specialization than a monolithic model). The paper provides no evidence about whether DiPaCo's gains are due to exploiting natural domain clusters in C4 or due to a more fundamental property of modular architectures.

• Per-path scale. All experiments fix each path at 150M parameters (12 transformer blocks, 896-dimensional hidden states, 16 heads). This is a relatively small model by contemporary standards. It is unknown whether DiPaCo's benefits — particularly the overfitting resistance from shared modules and the effectiveness of infrequent communication — persist at larger per-path scales (e.g., 1B parameters per path). At larger scales, each path might have sufficient capacity that overfitting on small shards is less of a concern, potentially changing the optimal sharing ratio. Conversely, larger paths might diverge more during local training phases, making the outer gradient averaging less coherent.

• Task generality. The paper evaluates only on language modeling perplexity. DiPaCo's modular design is motivated by general distributed learning, not language-specific properties. However, the routing mechanism (using the first 32 tokens to assign documents) depends on the prefix being informative about the remaining content — a property that holds for coherent text but may not hold for other modalities (images, audio) or tasks (classification, structured prediction). The paper does not test on any non-language task.

• Scaling behavior. The paper shows monotonic improvement from 8 to 256 paths but does not establish a scaling law. Does improvement continue linearly with more paths? Does it saturate? Is there a point where routing quality degrades because discriminative scoring of all paths becomes infeasible (the paper notes this: "discriminative scaling by scoring all paths cannot scale to large numbers of paths," Section 7.3)? Without scaling law data, a practitioner cannot estimate how many paths would be needed to match a specific target model quality.

What evidence exists in the paper. The paper's stated scope is explicitly narrow: "We consider this approach as a first prototype" (Abstract). All empirical claims are qualified with respect to the C4/150M-per-path setting. The paper provides no ablations varying the per-path model size, no experiments on other datasets, and no attempted scaling laws. Figure 9 shows scaling with path count for 150M paths only. The discriminative routing with "all paths scoring all documents" is noted as not scaling to large path counts (Section 7.3), but no alternative is evaluated.

Mitigation status. The limitation is honestly acknowledged (Section 6) but not addressed. The paper positions itself as a proof-of-concept demonstrating feasibility, not as a comprehensive empirical study. The scaling law gap is explicitly deferred to future work. A practitioner considering DiPaCo for a domain other than web-text language modeling, or at a substantially different scale, has no empirical guidance from this paper about whether the approach would transfer. The overfitting results for Flat MoE (Table 2) provide one transferable insight — parameter sharing becomes necessary when shards become small — but the specific thresholds (how small is too small, how much sharing is sufficient) are dataset- and scale-dependent and are not characterized beyond the C4/150M setting.

6.4 The Discriminative Router Does Not Scale to Very Large Path Counts

The assumption or constraint. The discriminative routing procedure (Section 2.4.2), which provides the best routing quality (Table 5: 0.7 PPL improvement over generative routing), requires evaluating every path's perplexity on every document in a held-out router dataset to generate classification targets. For P paths and N router documents, this is O(P × N) forward passes — feasible for P = 256 but prohibitive for the scaling trajectory the paper envisions. Section 7.3 explicitly notes:

"discriminative scaling by scoring all paths cannot scale to large numbers of paths."

The proposed scaling to 32 × 32 × 32 DiPaCo with 32,768 paths would require 32,768× more router evaluations than the 16 × 16 case, which is computationally infeasible.

The consequence. DiPaCo's quality depends significantly on routing quality. Generative routing (k-means) leaves a 0.7 PPL gap relative to discriminative routing even at P = 64 (Table 5). As the number of paths grows, the gap between generative and discriminative routing may widen — Figure 10 shows that discriminative routing gains are larger with more paths. If discriminative routing cannot scale, then the routing quality for very large DiPaCo models will be limited to generative methods (k-means, product k-means, or other unsupervised clustering), which do not directly optimize the language modeling objective. This creates a tension in DiPaCo's scaling story: the architecture supports arbitrarily many paths, but the best-known routing method for those paths requires computation that scales linearly with the number of paths, making it self-limiting.

Additionally, the discriminative router is a linear logistic regression on frozen features (the average hidden state from the first 32 tokens). As the number of paths grows, the classification problem becomes more challenging (more classes, fewer examples per class), and a linear classifier may not have sufficient capacity to discriminate between hundreds or thousands of path assignments. The bias correction mechanism (matching the target document-to-path distribution) addresses path starvation but does not improve the classifier's discriminative ability.

What evidence exists in the paper. Table 5 quantifies the discriminative vs. generative gap at P = 64. Figure 10 shows discriminative routing gains increase with path count from 8 to 64. Section 7.3 explicitly states the scaling limitation: "discriminative scaling by scoring all paths cannot scale to large numbers of paths." The paper notes that "we consider alternative generative routing approaches important for further study" but does not develop or evaluate any such alternatives for large path counts. The path sampling proposal (Section 2.6.2) — training only a subset of paths at each outer phase — is suggested for the training side but does not address the routing scaling problem.

Mitigation status. The paper acknowledges this limitation and suggests future work on "more sophisticated sharding" (Section 7), including "alternative generative routing approaches" (Section 7.3). However, the product k-means approach described in Section 7.3 is evaluated only briefly and shown to be worse than discriminative routing — it is an improvement over simple k-means but does not close the gap. The paper does not propose or evaluate any routing method that is both task-aware (like discriminative routing) and sublinear in the number of paths. A practitioner scaling DiPaCo beyond hundreds of paths would need to develop a new routing approach (perhaps approximate nearest-neighbor search over path embeddings, hierarchical routing where a coarse classifier first selects a subset of paths for scoring, or learned path embeddings that enable efficient similarity search) — none of which are explored in this work. This is a significant gap because routing quality is one of the two key ingredients (along with DiLoCo) that the paper identifies as critical to DiPaCo's effectiveness.

6.5 Communication and Hardware Assumptions Are Not Empirically Stress-Tested

The assumption or constraint. The paper's entire motivation rests on two assumptions (Section 2.1): "training compute (FLOPs) is relatively cheap" and "communication is relatively expensive." The experiments, however, were conducted in standard datacenter conditions where devices are co-located and communication is fast. Section 4.5's comparison against full synchronization was performed "in the setting where all devices are actually collocated and there is no constraint in terms of communication." The infrastructure section (Section 3) describes a design that could support geographically distributed training, but the reported experiments do not stress this capability. The paper provides no measurements of actual communication bandwidth, latency, packet loss, or worker heterogeneity in the experiments.

The consequence. The paper's central claim — that DiPaCo "facilitates training across poorly connected and heterogeneous workers, with a design that ensures robustness to worker failures and preemptions" (Abstract) — is supported by architectural arguments and infrastructure design, not by empirical stress-testing. Several specific questions are unanswered:

• How poorly connected can workers be? DiLoCo communicates every τ steps (150 in the experiments). The paper does not report the communication time as a fraction of total training time, nor does it vary the communication latency to find the threshold where DiLoCo degrades relative to full synchronization. If workers are separated by transcontinental latency (hundreds of milliseconds), the outer optimization phase — which requires gathering checkpoints from all workers, averaging, and redistributing — could become a significant fraction of total training time, eroding the wall-clock advantage.

• How heterogeneous can workers be? The backup pool design (Section 3.4) mentions using "multiple accelerator types using a low-tier priority," and the worker pool "can contain heterogeneous types of devices across different regions." But the paper does not report training with deliberately heterogeneous hardware — e.g., mixing GPU and TPU workers, or workers with different memory capacities and compute speeds. Heterogeneous workers would complete their inner optimization phases at different speeds, and the outer optimization must wait for the slowest worker using each module. The paper does not measure the straggler effect or propose mitigation strategies (e.g., adaptive τ based on worker speed).

• What failure rates can DiPaCo tolerate? The infrastructure design includes fault tolerance (task re-queueing, health monitoring), but the paper reports no experiments with induced failures — no preemption rates, no worker dropouts, no network partitions. The claim of "robustness to worker failures and preemptions" is untested. In a realistic poorly connected setting, workers might frequently go offline for minutes or hours; the paper does not show that DiPaCo continues to make progress or that the outer optimization remains coherent under such conditions.

• How much does the outer optimization cost in wall-clock time? Section 3.3 reports "average time per phase for outer update under 2 minutes" but does not report what fraction of total training time this represents, how it scales with the number of paths or the geographic distribution of workers, or how it compares to the per-step communication cost of synchronous training. The 2-minute figure is an operational observation, not a controlled measurement at different scales.

What evidence exists in the paper. Section 4.5 provides the only communication-stress experiment, but it is run in the fully collocated setting — DiLoCo matches full synchronization when communication is artificially restricted (to 1/150th the frequency) but the underlying hardware is still co-located and high-bandwidth. The infrastructure design (Figure 6, Section 3) is described in detail, and the backup pool (Section 3.4) and asynchronous checkpoint gathering (Section 3.3) are presented as scalability features. However, no latency measurements, bandwidth measurements, failure rate experiments, or heterogeneity experiments are reported. The paper does not demonstrate training across geographically distributed sites, despite language like "workers might use different hardware types … and might be placed in far away geographic areas" (Figure 4 caption).

Mitigation status. The paper does not attempt to empirically validate the robustness claims beyond the infrastructure design description. The authors present DiPaCo as "a first prototype towards a new paradigm of large-scale learning, one that is less synchronous and more modular" (Abstract), implying that the infrastructure validation is itself future work. The backup pool and elastic scaling design suggest awareness of real-world operational challenges, but without empirical stress-testing, the robustness claims remain aspirational. A practitioner considering DiPaCo for a genuinely poorly connected or heterogeneous deployment would need to conduct their own robustness evaluation, as the paper provides architectural guidance but no empirical guarantees about failure tolerance, communication thresholds, or straggler effects.

6.6 The Overfitting Problem for Independent Paths Is Diagnosed but Not Solved

The assumption or constraint. Section 4.3 and Table 2 demonstrate that Flat MoE (completely independent paths with no parameter sharing) overfits when the number of paths is large (256 paths, each trained on its own data shard). DiPaCo's shared-module design mitigates this by pooling data across paths — a module used by 16 paths is effectively trained on 16× more data. However, the paper does not solve the underlying problem: as the number of paths grows (the scaling trajectory the paper envisions with 32 × 32 × 32 DiPaCo at 32,768 paths), the per-shard data size shrinks linearly with the number of paths. Even with module sharing (where each module might be used by P / K_l paths), the effective data per module decreases as paths increase unless the total training data also scales proportionally.

The consequence. There is an inherent tension between DiPaCo's scalability and data availability. Adding more paths increases total model capacity, which should improve performance — and Figure 9 shows that it does, from 8 to 256 paths. But each additional path further subdivides the fixed C4 dataset into smaller shards. At some number of paths (unknown, because the paper does not measure scaling limits), the per-shard data will be too small for effective training even with shared modules, and overfitting will set in. The paper has no mechanism to determine how many paths are "too many" for a given dataset size, and the Table 2 results (Flat MoE overfits at 256 paths; DiPaCo with shared modules does not at 256 paths) only establish that sharing pushes the threshold further — not that the threshold has been eliminated.

Additionally, the top-2 overlapping shards technique (Section 2.4.4), which the 16 × 16 DiPaCo uses to combat overfitting, "increases the size of the shards" (Section 4) by assigning each document to two shards instead of one. But this also reduces the specialization of each path — a path now sees data that overlaps with another path's shard, which may dilute the benefit of routing. The paper does not explore how much overlap is optimal, or whether the optimal overlap ratio changes with the number of paths or the dataset size. At evaluation time, overlapping is not used because it would increase inference cost (requiring multiple paths to score each sequence). This creates an asymmetry: training benefits from overlapping but inference cannot use it, so the test-time performance may suffer if paths were overly regularized by the training-time overlap.

What evidence exists in the paper. Table 2 provides direct evidence of the overfitting problem for Flat MoE: 256 independent paths degrade to 14.1 PPL at 10K steps and continue overfitting at 64K steps (13.6 PPL even with overlapping shards and early stopping). The 16 × 16 DiPaCo with shared modules does not overfit at 256 paths, demonstrating that sharing is effective at this scale. However, no experiments push beyond 256 paths, so the overfitting threshold for DiPaCo is not measured. The paper does not vary the total dataset size (e.g., by subsampling C4) to characterize how the overfitting threshold depends on data volume. Figure 9 shows monotonic improvement from 8 to 256 paths, but this curve may eventually flatten or reverse — the paper provides no data about where that happens.

The paper also does not explore alternative overfitting mitigation strategies beyond overlapping shards and early stopping. Regularization techniques (weight decay, dropout, data augmentation), Bayesian approaches, or adaptive shard sizing (giving more data to paths with higher validation loss) are not considered.

Mitigation status. The paper identifies the overfitting problem and demonstrates that module sharing mitigates it at the 256-path scale, but does not characterize the scaling limits or propose additional mitigations for larger path counts. The path sampling proposal (Section 2.6.2) — training only a subset of paths at each outer phase — would exacerbate the overfitting problem by further reducing the effective data per path. The overlapping shards technique is a partial solution but is applied only at training time and its optimal configuration is not characterized. A practitioner scaling DiPaCo to thousands of paths would need to develop their own methods for determining when paths are overfitting, how much sharing is sufficient to prevent it, and whether training-time overlapping (which increases storage and reduces specialization) remains effective at larger scales. This is a fundamental limitation because it means DiPaCo's scalability — the primary motivation for the architecture — has an empirical ceiling determined by data availability that the paper does not quantify.

7. Implications and Future Directions

How This Work Changes the Landscape

DiPaCo introduces a methodological shift in how the field thinks about the relationship between model architecture and distributed training. The dominant paradigm — design the best architecture for a single device, then figure out how to distribute it — treats distribution as an after-the-fact engineering challenge. DiPaCo argues, and demonstrates empirically, that this order should be reversed: design the architecture around the constraints of the distributed training environment, and the distribution problem becomes dramatically simpler. This is not an incremental optimization of existing distributed training methods. It is a reframing of the problem statement itself — from "how do we make monolithic training work with less communication?" to "what architecture makes infrequent communication sufficient?"

The evidence for this reframing's validity comes primarily from Section 4.5, where DiPaCo trained with DiLoCo matches or exceeds the performance of fully synchronous training for 2 × 2 and 4 × 4 configurations (by 0.3 and 0.6 perplexity points, respectively), and falls short by only 0.1 perplexity for 8 × 8. This result is genuinely surprising under the standard paradigm, where more frequent communication is axiomatically better. Under DiPaCo's reframing, it makes sense: the architecture was designed so that each path is a self-contained function that benefits from independent local optimization, and the outer synchronization serves primarily to prevent shared modules from diverging too far — not to supply per-step gradient information. Adding more frequent communication provides diminishing returns because the architecture already encodes what needs to be shared (through module reuse structure) and what can remain independent (through path-specific modules).

The paper also resolves a specific tension in the modular ML literature. On one side, document-level expert models (Gross et al., 2017; Gururangan et al., 2023) demonstrate that routing entire documents to specialized experts works, but they hit an overfitting wall when the number of experts grows relative to the data. On the other side, token-level MoE models (Shazeer et al., 2017; Lepikhin et al., 2021) achieve excellent scaling by sharing parameters across all tokens, but they require full co-location and per-step synchronization. DiPaCo provides a principled middle ground: the hierarchical module structure with configurable per-level sharing ratios means that practitioners can dial the sharing-to-capacity ratio based on their specific constraints — more sharing when data per path is scarce, more independence when compute islands are abundant. Table 2 provides the diagnostic: Flat MoE (zero sharing) overfits at 256 paths while DiPaCo with partial sharing does not. Table 1 shows that at 64 paths, Flat MoE nearly matches DiPaCo, meaning sharing is unnecessary at that scale. This is actionable guidance, not just an observation: the required degree of parameter sharing is a function of the ratio of total parameters to total data.

The paper also makes the wall-clock vs. total FLOPs distinction explicit in a way that changes how modular training research should be evaluated. Most prior work on sparse models (Switch Transformer, GShard, token-level MoE) emphasizes training FLOP efficiency — matching dense model quality with fewer FLOPs. DiPaCo achieves the opposite: it spends substantially more total FLOPs (the paper explicitly acknowledges this in Section 6) but reduces wall-clock training time by 45% relative to a 1.3B dense model. This reframes the success criterion: if your constraint is total compute budget (energy, hardware cost, carbon), DiPaCo is not competitive. If your constraint is time-to-result (can you train the model before the conference deadline? can you iterate on architecture design in hours rather than days?), DiPaCo's distributed parallelism across independent compute islands becomes attractive. The paper's honest acknowledgment of this tradeoff — rather than trying to claim Pareto-optimality on all metrics — is itself a contribution: it forces the field to be explicit about which resource constraint matters most, rather than optimizing FLOP efficiency by default.

Research directions that become more attractive after this work:

• Communication-aware architecture design broadly — the idea that the architecture should be co-designed with the communication budget. This applies beyond language modeling to any domain where distributed training across loosely connected devices is desirable (federated learning, edge computing, multi-institutional collaborations with privacy constraints).
• Modular continual learning — DiPaCo's paths are natural units for incremental updates. A new domain can be accommodated by training new path-specific modules while keeping shared modules frozen, or by adding a new path that reuses existing modules plus new ones. The paper does not explore this, but the architecture makes it straightforward.
• Elastic training infrastructure — DiPaCo demonstrates that a task-queue-based training system with independent workers can absorb heterogeneous, preemptible hardware. This makes "spot market" cloud compute or volunteer-contributed resources (following PETALS/Borzunov et al., 2022) viable for training, not just inference.

Research directions that become less urgent:

• Token-level MoE for distributed training. If your goal is training across distant workers with limited communication, token-level routing (GShard, Switch Transformer) is the wrong starting point regardless of how much you optimize the communication algorithm — the architecture itself requires expert parameters to be available at each token step. DiPaCo shows that coarse routing is the architectural prerequisite for distributed training, making further optimization of token-level MoE for the low-communication setting somewhat moot.
• Purely optimization-based approaches to low-communication training (e.g., gradient compression, asynchronous SGD with staleness bounds). DiPaCo suggests that architectural changes (routing granularity, module sharing structure) provide larger and more robust gains than trying to make synchronous algorithms work under communication constraints. The 0.1 PPL gap between DiLoCo and full synchronization at 8 × 8 (Section 4.5) implies that the architecture-level solution is nearly sufficient, and further optimization gains would be marginal.

Follow-Up Research This Work Enables

Scaling laws of DiPaCo: path count, per-path size, and total parameters. The paper works at a single per-path scale (150M parameters) and shows monotonic improvement from 8 to 256 paths (Figure 9). The open question is whether this trend continues — and for how long — and how it interacts with per-path model size. A concrete experiment: fix total training FLOPs, then sweep over configurations of (number_of_paths, parameters_per_path, sharing_ratio) to find the compute-optimal modular configuration for a target dataset size. This is the modular analog of the Chinchilla scaling laws (Hoffmann et al., 2022), but with additional degrees of freedom (routing strategy, sharing structure, communication frequency). The paper's existing infrastructure — task queues, sharded outer optimization, elastic worker pools — makes such sweeps feasible because different configurations can run in parallel on heterogeneous hardware. A strong follow-up would produce plots analogous to Figure 9 but with total FLOPs matched across configurations, revealing whether there is a regime where DiPaCo is FLOP-competitive with dense training (which the current paper does not claim but also does not disprove).

Discriminative routing at scale via approximate nearest-neighbor search over path embeddings. The paper's discriminative router requires evaluating all paths on all held-out documents, which is O(P × N) and does not scale to P = 32,768 (Section 7.3). A concrete solution: train an embedding for each path during training (e.g., the average hidden state of documents assigned to that path, or a learned path embedding trained jointly), then use approximate nearest-neighbor search (FAISS, ScaNN) to assign new documents to the top-k nearest path embeddings for scoring. This reduces the discriminative routing cost from O(P) to O(log P) for the search plus O(k) for scoring only the top-k candidates. The key metric: at what value of k does routing quality match the full P-way scoring? The paper's bias correction mechanism (Section 7.2.1) would need to be adapted to handle the truncated candidate set. This would directly enable scaling DiPaCo to thousands of paths without sacrificing routing quality to generative k-means.

Characterizing the overfitting threshold for Flat MoE as a function of shard size. Table 2 shows Flat MoE overfits at 256 paths but not at 64 paths (Table 1). A systematic experiment: for a fixed dataset (C4), sweep the number of fully independent paths from 8 to 1,024, measure validation perplexity at convergence, and compute the per-shard token count. The expected result is a phase transition — at some critical shard size, validation perplexity stops improving with more paths and begins degrading. This threshold is a function of the model architecture (150M parameters) and the data distribution (C4's domain diversity), and characterizing it would provide practitioners with a simple rule: "with this model size, you need at least X tokens per shard to avoid overfitting." The paper's existing results suggest X is somewhere between total_C4_tokens / 64 (which doesn't overfit) and total_C4_tokens / 256 (which does, for Flat MoE), but a finer sweep would pin down the threshold and test whether it obeys a simple scaling relationship (e.g., tokens per shard should be at least C × parameter_count for some constant C).

KV-cache sharing across paths to enable low-latency frequent re-routing. Section 6 identifies the KV-cache recomputation cost as the primary inference bottleneck when re-routing every 64 tokens. If early DiPaCo modules (the first few transformer blocks) are shared across many paths, their KV-caches could be reused after a re-route — the new path would only need to recompute the KV-cache for path-specific later blocks. A concrete experiment: for a 16 × 16 DiPaCo, measure the fraction of KV-cache entries that are invariant across path switches (because the shared early modules produce the same keys/values regardless of which path they're part of), and measure inference latency with incremental KV-cache recomputation vs. full recomputation. The hypothesis is that if the first 6 of 12 transformer blocks are shared, a path switch requires recomputing only the last 6 blocks' KV-cache, reducing the recomputation cost by roughly 50%. If this holds, the "6× fewer parameters per forward pass" claim (Section 4.4) becomes close to an actual latency ratio as well, because the shared blocks' computation is amortized across path switches. This would make DiPaCo's inference advantage over dense models much stronger than the current paper demonstrates, since the paper only claims a parameter-count advantage without latency data.

Training DiPaCo on a dataset without natural domain structure to test routing necessity. The paper's routing mechanism — both generative and discriminative — relies on the existence of topical or stylistic clusters in C4 that paths can specialize to. A stress test: train the same 16 × 16 DiPaCo on (a) standard C4, (b) C4 with documents randomly shuffled (breaking topical coherence within documents but preserving token-level statistics), and (c) C4 with tokens randomly permuted (destroying all linguistic structure). If DiPaCo's advantage over a 150M dense baseline persists in condition (b), routing is doing more than exploiting domain clusters — it is enabling beneficial specialization even on artificially homogenized data. If the advantage disappears in (b), then DiPaCo's effectiveness depends on natural domain structure, and practitioners should not expect gains on datasets without such structure (e.g., uniformly formatted reasoning problems, synthetic data). This experiment would clarify the scope of DiPaCo's applicability and guide domain selection for modular architectures.

Path-level early stopping as a dynamic compute allocation mechanism. The paper uses path-specific early stopping (Section 2.7) to prevent overfitting on small shards, but frames it as a regularization technique. It could instead be viewed as a compute allocation mechanism: paths that converge faster (because they have easier or more homogeneous data shards) stop training earlier, freeing compute for paths that need more steps. A concrete experiment: during each outer phase, monitor each path's validation loss on its held-out shard. Allocate the inner steps τ adaptively — paths with plateauing validation loss stop early, while paths with decreasing validation loss continue for additional steps (up to some maximum). Compare total wall-clock time and final validation perplexity against the current fixed-τ approach. The hypothesis is that adaptive τ improves both metrics simultaneously: easier paths consume fewer FLOPs, and harder paths get more optimization, improving overall mixture quality. This connects DiPaCo to the broader literature on adaptive computation time and dynamic resource allocation in mixture models.

Practical Applications and Downstream Use Cases

Multi-institutional collaborative model training without data sharing. DiPaCo's architecture is well-suited to a scenario where multiple organizations each hold private text data (e.g., hospitals with clinical notes, law firms with case documents, universities with research papers) and want to jointly train a language model without pooling their data. Each institution trains one or more DiPaCo paths on their own data behind their firewall, using shared modules (whose parameters are synchronized via DiLoCo) to transfer general linguistic knowledge across institutions, while path-specific modules capture domain-specific vocabulary and style. The outer optimization — averaging parameter differences for shared modules only — requires communicating module parameters, not raw data, preserving data privacy. The paper's finding that DiPaCo matches dense model quality with 150M parameters per path means each institution only needs modest compute (a handful of GPUs) to participate. The 2-minute outer optimization phase and asynchronous checkpoint gathering (Section 3.3) make cross-institutional synchronization practical even over commodity internet connections. The discriminative routing could be adapted to use a small shared evaluation set (e.g., public domain text) to assign documents to the most appropriate institution-specific path without exposing private data.

Rapid prototyping and architecture search with elastic cloud resources. The backup pool design (Section 3.4) — spawning training workers on low-priority, preemptible accelerators — enables a specific workflow that is currently impractical with monolithic models: testing many architectural variants in parallel with minimal dedicated hardware commitment. A research team could maintain a small pool of dedicated workers (say, 16 GPUs) and use preemptible cloud instances to scale to 256 paths when testing a 16 × 16 DiPaCo configuration. If preemptions occur, the task queue reassigns the affected paths to remaining workers (Section 3.1); training slows but does not stop. The wall-clock time advantage (45% faster than a 1.3B dense model for equivalent quality, Section 4.1) means that each architectural variant can be tested in roughly half the time of training a comparable dense model from scratch. This enables a development cycle where multiple hypotheses about module sharing structure, routing strategy, and path count are tested simultaneously, with the best configuration promoted to full training on dedicated hardware.

Domain-adaptive deployment where only relevant paths are loaded at inference. A production language model serving diverse user queries — some about programming, some about medicine, some about entertainment — could be implemented as a DiPaCo where each path specializes in a different domain. At inference, the router (trained discriminatively on a small labeled dataset of domain exemplars) assigns each query to the appropriate path. The key benefit over a monolithic domain-adapted model: only the selected path's 150M parameters need to be loaded in GPU memory per query, rather than a 1B+ parameter model that covers all domains. If queries arrive with temporal locality (a user session about programming followed by another session about medicine), the system can keep recently used paths cached and switch between them at the cost of KV-cache recomputation (mitigated if early modules are shared, per the KV-cache sharing idea above). The paper's Table 3 result — routing every 64 tokens matches a 1B dense model — suggests that even within a single multi-domain query, the router can switch paths at paragraph boundaries (e.g., an explainer article that transitions from general background to technical details), using domain-specific capacity where needed without paying the memory cost of a monolithic model that must represent all domains simultaneously.

When to Prefer This Method

The paper articulates a specific set of assumptions under which DiPaCo is the preferred approach, stated explicitly in Section 2.1: training compute is relatively cheap, communication is relatively expensive, and no single compute island can host the full desired model. These conditions define the decision boundary against the two primary alternatives:

• Prefer DiPaCo over dense monolithic training when:

  • You have access to many independent compute islands (different GPU clusters, different data centers, volunteer-contributed hardware) that cannot be co-located into a single tightly connected cluster
  • The total model size you want exceeds what any single island can store in memory (the "cannot instantiate models as large as we would like" assumption)
  • Wall-clock training time matters more than total FLOP efficiency — you are willing to spend more total compute to get results faster
  • Worker preemption or heterogeneity is expected, and you need training to continue making progress despite individual worker failures

• Prefer dense monolithic training (or token-level MoE) over DiPaCo when:

  • You have a dedicated, co-located cluster with sufficient interconnect bandwidth — the assumptions that motivate DiPaCo don't hold, and the FLOP inefficiency (Section 6) is a pure cost with no offsetting benefit
  • Total energy, hardware cost, or carbon footprint is the primary constraint — DiPaCo is explicitly not optimized for FLOP efficiency and will consume substantially more total compute than a comparable-quality dense model
  • Inference latency is critical and you cannot amortize the cost of KV-cache recomputation or path switching — the paper provides no latency data, and the frequent test-time routing that enables DiPaCo's best results (Table 3) requires re-computing attention state for each re-routed chunk

• Prefer Flat MoE (fully independent experts, no shared modules) over compositional DiPaCo when:

  • The number of paths is small relative to the total data volume — Table 1 shows that at 64 paths, Flat MoE nearly matches 16 × 16 DiPaCo, and the simpler architecture (no outer optimization, no module sharing) is easier to implement without the Spanner-based checkpoint tracking infrastructure
  • You only need to deploy a subset of paths at inference time and have no need for shared representations — for example, when each path corresponds to a completely disjoint language or task where transfer is undesirable

• Prefer DiLoCo (dense model with DiLoCo optimization, no modularity) over DiPaCo when:

  • The desired model fits comfortably on each worker's memory and the goal is faster training through data parallelism with infrequent communication — at this point, DiPaCo's modularity adds complexity (routing infrastructure, path management) without the benefit of exceeding single-worker model capacity
  • Table 1 shows that DiLoCo on a 150M model reaches 13.8 PPL, which is better than the 150M dense baseline but worse than any modular DiPaCo variant — the modularity provides the additional capacity needed to absorb the extra FLOPs and tokens