NoLoCo: No-all-reduce Low Communication Training Method for Large Models

In plain words

Training a large language model usually means many machines that keep their copies of the model identical by exchanging all of their numbers at regular points. That exchange slows down as more machines join and as the links between them get worse, so training tends to need expensive, tightly linked clusters.

This paper removes the global exchange. Each machine trains on its own for a while. At the sync point it only blends its numbers with one randomly picked partner, and the partners are picked again at every sync. Over many rounds the copies drift toward each other without ever being fully aligned. Copies that hold the same layer are also mixed in a second way: during training, each sample is sent through a randomly chosen chain of machines, so a layer sees input that came through many different copies of the rest of the model. The authors change the momentum update used at sync time so that the copies stay close together. On language-model pre-training, the method reaches a better final quality than a widely used low-communication method in most runs, and it is faster in wall-clock time when the network is slow.

TL;DR

• Removes the global parameter synchronisation of DiLoCo-style training. Each outer step synchronises only inside a random pair of replicas, re-drawn every outer step §3.2 sparse_optimizer_c.py.
• Inner steps route each sample through a random one-to-one pairing of replicas between neighbouring pipeline stages §3.1 dipaco_pipeline.py.
• The outer update is a modified Nesterov momentum with an extra local-averaging term. A convergence result is given for a quadratic model under stated assumptions §3.2 [thm. 1] app. D.
• Author-reported: validation perplexity below DiLoCo in all 8 Reddit runs and 6 of 8 C4 runs, by up to about 4% §4.2 tab. 1. Author-reported wall-clock speed-up over DiLoCo of about 4x at 100 Mb/s §4.3 tab. 2.
• The official code is MIT-licensed. The paper's own equation for the outer update differs in sign from the code, and the local-averaging coefficient is a separate symbol in the paper but tied to the outer learning rate in the code. Read Paper vs code before implementing.

Problem

Data-parallel training keeps replicas identical with a global exchange, often an all-reduce, which can dominate step time on slow or high-latency links §1 §2. DiLoCo lowers how often replicas exchange, but each exchange is still global, so its cost and the waiting it causes grow with the number of workers §2 §4.3. The authors evaluate up to 1024 accelerators in a latency simulation §4.3 app. G. Their convergence result assumes a quadratic objective with a positive-definite curvature matrix, a Gaussian distribution for the shift of the optimum, and SGD with a constant inner learning rate §3.2.

Method

NoLoCo is data parallelism (DP) combined with pipeline parallelism (PP) §3. It changes two things: during inner steps, activations take a randomly chosen pipeline path, and during outer steps, replicas synchronise only inside small groups §3.

Notation: N is the number of model instances, n the size of a synchronisation group (n = 2 in the experiments), phi_{t,i} the slow weights of instance i, theta_{t+1,i} its fast weights after m inner steps, delta its outer momentum, alpha the Nesterov momentum, beta the outer learning rate and gamma the local-averaging coefficient §3.2. The paper's m is the number of local (inner) steps between outer steps. Section 4.1 sets 50 for NoLoCo and H = 100 for DiLoCo §4.1.

Inner step: random pipeline routing

1. Each pipeline stage is replicated across DP ranks. Before each sample is sent to the next stage, a permutation of the DP ranks is drawn, so that each replica sends to exactly one replica of the next stage §3.1 alg. 1 dipaco_pipeline.py.
2. The receiving replica uses the inverse of the same permutation to find its sender dipaco_pipeline.py. The paper's pseudocode gives a random neighbour for each direction instead, using a RandomRank function that is not defined app. A alg. 1.
3. Gradients return along the same path as the forward activations §3.1.

# Random pipeline routing for one sample  [§3.1] [alg. 1]
# Source: dipaco_pipeline.py#L41-L75
# Assumes P pipeline stages and D data-parallel replicas. Every rank derives
# the same permutation from a shared seed, so sender and receiver agree.

def route_sample(dp_rank, stage, sample_seed, num_stages, num_replicas):
    # Seed is shared by all ranks and differs per sample
    # [code: dipaco_pipeline.py#L155-L158]
    seed = sample_seed * num_stages + stage
    # One-to-one map from sender rank to receiver rank
    # [code: dipaco_pipeline.py#L71-L75]
    permutation = random_permutation(num_replicas, seed)
    next_rank = permutation[dp_rank]
    # The receiver at stage + 1 inverts the same map
    # [code: dipaco_pipeline.py#L41-L47]
    prev_rank = permutation.index(dp_rank)
    return next_rank, prev_rank

Outer step: pairwise synchronisation with modified momentum

4. Every outer step, all ranks draw the same random permutation of the DP ranks. The permutation is seeded by the optimizer step counter, and it is cut into groups of n §3.2 sparse_optimizer_c.py.
5. Each member of a group sends its slow weights and its outer gradient to the group root. The root averages both and sends the averages back. Each member then applies the update below sparse_optimizer_c.py.
6. Outer gradient of each instance eq. 1:

\Delta_{t,i} = \theta_{t+1,i} - \phi_{t,i} eq. 1

7. Momentum update, as written in the paper eq. 2. The sign of the first sum is discussed in Paper vs code:

\delta_{t,i} = \alpha\,\delta_{t-1,i} - \frac{\beta}{n}\sum_{j}\Delta_{t,j} - \gamma\left(\phi_{t,i} - \frac{1}{n}\sum_{j}\phi_{t,j}\right) eq. 2

8. Slow weights move by the momentum eq. 3, and the fast weights restart from the new slow weights §3.2 sparse_optimizer_c.py:

\phi_{t+1,i} = \phi_{t,i} + \delta_{t,i} eq. 3

The update the official code actually runs is the following. It is the version to implement. The code uses beta in place of gamma, and it adds the outer-gradient term with a plus sign sparse_optimizer_c.py.

# Outer step for one replica i  [§3.2] [eq. 1] [eq. 3]
# Code form: there is no separate gamma; the pull term uses beta = 0.7
# [code: sparse_optimizer_c.py#L61]
# Input:  theta (fast weights), phi (slow weights), delta (momentum),
#         group (list of n = 2 replicas), alpha = 0.5, beta = 0.7
# Output: new phi and delta; the replica restarts from phi

def outer_step(theta, phi, delta, group, alpha, beta):
    # Outer gradient of this replica  [eq. 1]
    outer_grad = theta - phi
    # Average slow weights and outer gradients over the group  [§3.2]
    mean_phi = mean(phi of each member of group)
    mean_grad = mean(outer_grad of each member of group)
    # Momentum with the outer-gradient term and the pull toward the group mean
    # Gradient-term sign follows the code, not eq. 2
    # [code: sparse_optimizer_c.py#L57-L63]
    delta = alpha * delta + beta * mean_grad + beta * (mean_phi - phi)
    # Move slow weights by the momentum  [eq. 3]
    phi = phi + delta
    # Restart the inner phase from the new slow weights
    # [code: sparse_optimizer_c.py#L64-L66]
    theta = phi
    return phi, delta, theta

# Group draw for one outer step  [§3.2]
# Source: sparse_optimizer_c.py#L104-L109
# Every rank computes the same groups, so no extra communication is needed.

def make_groups(num_replicas, step, n):
    # numpy default_rng seeded with step + 1
    # [code: sparse_optimizer_c.py#L6-L8]
    order = random_permutation(num_replicas, seed = step + 1)
    num_groups = num_replicas // n
    # Group g takes positions g, g + num_groups, g + 2 * num_groups, ...
    # [code: sparse_optimizer_c.py#L19-L25]
    return [order[g::num_groups] for g in range(num_groups)]

Other points from the text:

• With the whole network in one group, eq. 2 reduces to the DiLoCo outer momentum, and the last term vanishes §3.2.
• The paper suggests that slow weights could be sent at the end of the previous outer step, so that communication overlaps with the next inner steps §3.2. The official code does not do this (see Paper vs code).
• Inner optimiser: Adam, gradient clipping at norm 1.0, cosine learning-rate schedule with linear warm-up §4.1 app. B.

Convergence result

Theorem 1 §3.2 states that, for a quadratic loss with a Gaussian shift of the optimum, a constant inner learning rate omega, and infinitely many outer steps, the expected slow weights go to zero and their variance is proportional to omega squared. The proof is in appendix D (outer and inner recursions) and appendix E (the expectation is the same for all replicas). The variance part relies on a conjecture stated in appendix D.2, and it neglects covariance between non-consecutive fast weights app. D.1 app. D.2. Appendix D.3 gives a stability range for gamma: with alpha and n as above, it requires sqrt(n / (2(n-1))) * alpha < gamma < sqrt((n / (2(n-1))) * (2 + alpha^2)) app. D.3. Derived from that bound with n = 2: 0.5 < gamma < 1.5 for alpha = 0.5.

Hyperparameters

Values for the main experiments §4.1. Table 3 gives the architectures and is reproduced in full below.

Table 3, model parameters tab. 3:

Latency simulation app. G: inner-step latency is LogNormal with mu = 1 and sigma^2 = 0.5. Per-message time has the expectation t_c = exp(mu + sigma^2 / 2) fig. 5 app. G.

Results

Author-reported. Setup: decentralised DP with PP, Llama-style models of 125M, 1.3B and 6.8B parameters, Reddit and C4 pre-training, FSDP and DiLoCo as baselines §4.1. The paper gives no number of random seeds §4.2.

Final validation perplexity tab. 1. Bold marks the lowest value among the low-communication methods (DiLoCo and NoLoCo). Columns: total accelerators, data-parallel world size (DP), pipeline stages (PP).

Derived from the table: NoLoCo has the lower perplexity than DiLoCo in 8 of 8 Reddit runs and 6 of 8 C4 runs. The largest DiLoCo-vs-FSDP gap on C4 is in the small model at DP 16 (39.1 against 29.1, about 34%). The text reports up to 30% for C4 small at DP 16, which matches NoLoCo's 37.7 §4.2 tab. 1.

Perplexity curves, training-loss curves and weight spread are shown in figures fig. 2 fig. 3 app. F. The authors report that NoLoCo's gap to DiLoCo is a few percent early and late in training, and that the late-stage gap depends on dataset and model size §4.2 fig. 3. The authors also report that the spread of replica weights grows early and shrinks as the learning rate decays, with a Pearson correlation of 0.91 to 0.97 between that spread and the learning rate §4.2. They take this as support for controlling divergence through the learning-rate schedule §4.2 [thm. 1].

Wall-clock inner-step time on 32 workers, medium model, pure DP, 100 inner steps per outer step, two outer steps tab. 2 §4.3. Throughput is H / T_outer §4.3.

Derived from the table: NoLoCo is about 4.05x faster than DiLoCo and about 351x faster than FSDP at 100 Mb/s; about 2.34x and 196x at 1 Gb/s; and nearly equal to DiLoCo at 100 Gb/s (1.64 s against 1.70 s) §4.3 tab. 2. The authors combine this with per-step convergence to claim about 4x faster time-to-target on consumer networks §4.3. The Table 2 time column is labelled "per inner step", but the throughput values match H / time only if the time is per outer step of 100 inner steps (see Paper vs code).

Tree all-reduce versus pairwise averaging, simulated over world sizes from 4 to 1024: the expected ratio grows with world size and with message-time variance §4.3 fig. 5. The authors state that DiLoCo falls behind by about 20% at 1024 accelerators when only the blocking wait for the slowest worker is counted §4.3.

Hyperparameter sweeps for alpha, beta, gamma and H are shown only as figures on the 1.3B C4 model §5 fig. 6; the numbers are not in the text. Local group size, on the 125M model with 0.5M training tokens: n in {2, 4, 8} gives perplexity that improves weakly with n §5 fig. 7. Random routing, against fixed routing with the outer step switched off, lowers the spread of weights by about 15% for the small model and about 10% for the medium model §5 fig. 8.

Table 4, batch-size check on the medium model at 64 accelerators, four PP stages tab. 4. The 1M-token column matches the medium 64-accelerator row of Table 1 tab. 1:

Comparisons with other low-communication methods are in appendix C app. C fig. 9 fig. 10 fig. 11. All use the medium model on C4 with DP 8 and PP 2 app. C. The paper reports NoLoCo ahead of Epidemic Learning fig. 9 and PALSGD fig. 10. For SparseLoCo, throughput is estimated by scaling DiLoCo's communication time by a compression ratio of 1.66%, taken from another paper's table, not measured app. C fig. 11.

Paper vs code

The official repository is gensyn-ai/noloco, pinned at a1b4a425bdc4050a356cf9f4bae7c383419703ab. The paper links it in §1 and in appendix B. The repository has an MIT licence with copyright held by Gensyn. The paper links it, and the LaTeX class is Gensyn's (gensyn.cls), so it is accepted as official.

1. Sign of the outer-gradient term. Eq. 2 subtracts beta / n times the sum of outer gradients, with the outer gradient defined as theta minus phi eq. 1 eq. 2. The code adds beta times the group mean of the outer gradient sparse_optimizer_c.py. Taken literally, eq. 2 moves the slow weights away from the fast weights. The appendix derivations use the same printed sign app. D.1 app. E. Implement the code's sign.
2. Gamma is not a separate parameter in code. The paper treats gamma as its own hyperparameter §3.2 with a stability range app. D.3, and sweeps it in figure 6c fig. 6. The code has no gamma. The local-averaging term is multiplied by the outer learning rate beta = 0.7 sparse_optimizer_c.py. The value used in the paper is not stated, so the gamma sweep cannot be reproduced from this code.
3. Routing granularity. The paper says routing is re-randomised at each step §3.1. The code draws a new permutation for each sample: the seed is a counter that advances by the batch size on each forward pass, and each row of the batch is a separate micro-batch pipeline_llama.py pipeline_llama.py dipaco_pipeline.py dipaco_pipeline.py.
4. Routing pseudocode. Algorithm 1 draws a random rank for the previous and the next replica, with RandomRank undefined app. A alg. 1. The code uses one permutation per sample, so each replica sends to exactly one replica and receives from exactly one dipaco_pipeline.py dipaco_pipeline.py.
5. Config names. Configs named noloco_* without mean use PipelineDiLoCoOptimizer, an all-reduce DiLoCo baseline with H = 100 and alpha = 0.3 noloco_1b.yaml pipeline_diloco_optimizer.py. The NoLoCo runs are the *_mean_* configs, which use SparseOptimizerC noloco_mean_1b.yaml. The mean reading is inferred from the optimizer class; the repository does not say so.
6. Context length. Table 3 gives 2048 tab. 3. The text and the configs use 1024 §4.1 c4_base.yaml. This card uses 1024, but the paper does not resolve the conflict.
7. Vocabulary and tokeniser. The paper says a SentencePiece tokeniser with 128,000 tokens §4.1. The model vocabulary is 128,002 and the tokeniser is meta-llama/Meta-Llama-3-8B pipeline_llama.py c4_base.yaml.
8. Optimiser state. The code averages only weights. Adam moments stay on each replica sparse_optimizer_c.py. The paper does not say what happens to them. The state-averaging code in the DiLoCo optimizer is commented out pipeline_diloco_optimizer.py.
9. Communication overlap. The paper says slow weights can be sent at the end of the previous outer step to overlap with compute §3.2. The code uses blocking CPU send and receive of full parameter tensors, two per parameter, inside the outer step sparse_optimizer_c.py. No overlap is implemented.
10. Initial weights. The slow weights start as a copy of the initial model sparse_optimizer_c.py. The repository does not seed or broadcast the weights (no manual_seed and no broadcast in src/noloco). The paper does not say how replicas are initialised. Identical starting weights are a precondition the code does not enforce.
11. Table 2 label. The caption says "time per inner step" tab. 2. The throughput column equals 100 / time, which means the time is per outer step of 100 inner steps (for example 100 / 61.80 = 1.618) tab. 2 §4.3.
12. Latency speed-up. The text says the NoLoCo/DiLoCo step-time ratio "approaches" log2(32) = 5 at 100 Mb/s §4.3. Table 2 gives 61.80 / 15.26 = 4.05 tab. 2.
13. Latency simulation length. Appendix G says 250 outer iterations app. G. Figure 5 says 500 outer steps fig. 5. The simulation code is not in the repository.
14. Sweep naming. Section 5 names alpha as the local averaging parameter, eta as the outer learning rate and beta as the outer momentum §5. Figure 6 and eq. 2 use alpha for momentum, beta for the outer learning rate and gamma for averaging fig. 6 eq. 2.
15. Default values in the code. The builders default to outer momentum 0.9 sparse_optimizer_c.py and pipeline_diloco_optimizer.py. The configs set 0.5 for NoLoCo and 0.3 for DiLoCo, as the paper states.

Checked and consistent: group size 2, 50 inner steps for NoLoCo, 100 for DiLoCo, beta 0.7, alpha 0.5 and 0.3, learning rates per size, 1000 warm-up steps, clipping at 1.0, 25,000 steps, and the pairwise group draw §4.1 tab. 3 sparse_optimizer_c.py.

Pitfalls

• Use the right config. Reproduce NoLoCo with the *_mean_* configs. The noloco_* files without mean run the DiLoCo baseline.
• Pick gamma explicitly. The paper gives no value for gamma in the text. The code uses beta = 0.7. Keep the stability range in mind: with n = 2 and alpha = 0.5 the printed bound is 0.5 < gamma < 1.5 app. D.3.
• Fix the sign. Do not implement eq. 2 as printed. Use the code's sign sparse_optimizer_c.py.
• Routing is per sample. A new permutation per sample is what the code does. A per-step permutation, as the text describes, is a different and untested variant §3.1 dipaco_pipeline.py.
• Communication is not smaller per sync. Each replica sends its slow weights and its outer gradient to the group root, which sends two averaged tensors back. That is two full-size tensors per parameter in each direction per outer step sparse_optimizer_c.py. The saving comes from removing the global barrier, not from fewer bytes. The authors say the per-byte volume is about the same as DiLoCo app. C.
• Speed-ups are hardware-specific. Wall-clock numbers come from H100 nodes with a throttled NIC app. B and per-step timings tab. 2. The 4x figure is for consumer-grade bandwidth in the authors' setting, not a general result.
• The convergence result is narrow. It is for a quadratic toy problem with constant learning rate. The variance argument uses a conjecture app. D.2. Do not read it as a guarantee for transformer training.
• Small gains need repeats. The perplexity gains are up to about 4%. The paper does not state the number of seeds or repeats per setting. The authors attribute the gain to regularisation from weight noise, but they do not test this §4.2. All training stays within one epoch.
• Replica spread follows the learning rate. The paper argues that the learning-rate schedule controls how far replicas drift §4.2 [thm. 1]. The variance result ties the spread to the inner learning rate, so a schedule that stays high keeps replicas further apart [thm. 1]. This is an inference from the theorem; the paper runs no constant-rate comparison.
• Initial weights must match. The code assumes identical starting weights and does not enforce it (Paper vs code, item 10).
• Adam state is local. Moments diverge across replicas and are not synced sparse_optimizer_c.py. Results may depend on this.
• Sweeps are figure-only. The alpha, beta, gamma and H sweeps and the group-size and routing ablations are only in figures §5 fig. 6 fig. 7 fig. 8.

Lineage

• Builds on: DiLoCo, the all-reduce outer optimiser used as the primary baseline §1 §4.1; the Lookahead optimiser, whose slow and fast weights are reused §3.2; Epidemic Learning and gossip averaging, which motivate pairwise mixing §2; pipeline-routing work SWARM and DiPaCo, which the paper cites for random and learned routing §2.
• Compared against: FSDP, DiLoCo §4.1; Epidemic Learning, SparseLoCo, PALSGD, Async LocSGD, GWTF §4.1 app. C.
• Superseded by: not stated.