Affinity-Aware Sharding for Delayed Tensor Parallelism

Eloi de Reynal

cs.LGSep 12, 2026

Affinity-Aware Sharding for Delayed Tensor Parallelism

Authors: Eloi de Reynal

Organizations: Kog Team

Abstract

Delayed Tensor Parallelism (DTP) removes the blocking all-reduce of tensor-parallel Transformer inference. Every device adds its own partial output to its residual stream (and broadcasts it) immediately, but only gathers (receives) the other devices' partials $δ$ modules later. A TP to DTP change therefore amounts to a real architecture change, and dense Transformer models need to be retrained or distilled after adaptation. We show that DTP breaks the permutation symmetry of neurons inside FFNs and of KV heads inside attention modules, and that this symmetry breakage makes the sharding itself a modelling decision. We show that maximising the affinity between the KV heads and the FFN neurons co-located on a device, by permuting the dense model before sharding, speeds up the distillation or retraining process. The affinity is measured with a first-order approximation of the damage that losing a head's contribution does to each neuron's output, and the co-located affinity is maximised with a coordinate-ascent optimiser that alternates an exact balanced assignment of neurons with an exhaustive search over the KV head partitions. The whole procedure takes under two minutes on one GPU for Qwen3-0.6B and Danube3-500M. On these models, at $δ=1$ , the affinity-optimised layouts reach any distillation target in about half to two thirds of the steps needed by the naive contiguous layouts, over the whole 10k-step range we tested, and every optimised seed beats every contiguous seed and all but one of the sixteen random layouts. We also show that the co-located affinity score at initialisation predicts the KL to the base model after training, across seventeen layouts ranging from anti-optimised to optimised (Pearson $-0.81$ and $-0.89$ ).

Figures & tables

Panel of Figure 1 — Figure 1: What device $l$ sees under DTP at $\delta=1$ (equation 4 ). Solid: available when the module runs. Dashed red: not yet available. The two edges the layout can control are heads $(i)\to$ neurons $(i)$ (the up edge) and neurons $(i)\to$ KV groups $(i+1)$ (the down edge).

Model	Layout	$n$	KL@500	KL@1000	ppl@1000	step matching contig.@1000
Qwen3-0.6B	optimised ( fo )	3	$0.361\pm 0.009$	$\mathbf{0.303\pm 0.003}$	$\mathbf{17.01}$	347 ( $2.9\times$ )
	contiguous	3	$0.500\pm 0.055$	$0.415\pm 0.033$	19.20	1000
	random	8	$0.449\pm 0.044$	$0.374\pm 0.015$	18.31	—
Danube3-500M	optimised ( fo )	3	$0.763\pm 0.056$	$\mathbf{0.667\pm 0.043}$	$\mathbf{16.05}$	565 ( $1.8\times$ )
	contiguous	3	$0.832\pm 0.010$	$0.738\pm 0.009$	17.45	1000
	random	8	$0.844\pm 0.023$	$0.737\pm 0.018$	17.27	—

Table 1: KL to the dense model (nats per token, mean

\pm

sample sd over seeds) and perplexity on the 16 WikiText-2 blocks after 1000 distillation steps, and the step at which the mean optimised curve first reaches the contiguous arm’s final KL. Random is eight layouts with one seed each.

Figure 2: Qwen3-0.6B,

L=4

\delta=1

: KL to the dense model (left) and perplexity (right) versus distillation step. Bold: mean of three training seeds; faint: individual seeds; grey: the eight random layouts. The Danube3-500M version is Figure 6 in the appendix.

Figure 3: KL to the dense model against the co-located fo score

J

of the layout (equation 6 ), untrained (left), after 500 (middle) and 1000 steps (right); Qwen3-0.6B top, Danube3-500M bottom. One point per layout, seeds averaged where there are several (ringed), individual seeds faint; dashed line: least-squares fit over the seventeen layouts. Untrained, the score says nothing; after training it explains most of the variance.

Figure 4: Qwen3-0.6B, 10,000 distillation steps (164M tokens), one seed, evaluation every 250 steps. Top: KL and perplexity on a log step axis. Bottom left: the KL gap contiguous minus optimised with its size relative to the contiguous KL. Bottom right: the compute multiplier, the contiguous step

s

divided by the interpolated step at which the optimised run first reached the same KL (grey per evaluation, black rolling median of five). Danube3-500M: Figure 7 .

step	optimised KL	contiguous KL	gap	gap / contiguous	opt. ppl / contig. ppl	multiplier
500	0.392	0.470	0.079	16.7%	18.59 / 20.32	$\geq 2.0$
1000	0.333	0.397	0.064	16.1%	17.46 / 18.75	2.10
2000	0.281	0.325	0.043	13.4%	16.50 / 17.38	1.81
4000	0.237	0.283	0.046	16.3%	15.81 / 16.66	2.03
6000	0.207	0.247	0.040	16.3%	15.25 / 15.94	1.59
8000	0.186	0.218	0.032	14.9%	14.94 / 15.48	1.64

Table 2: The two 10k-step runs on Qwen3-0.6B (one seed). The multiplier is the contiguous step

s

over the optimised step reaching the same KL.

		Qwen3-0.6B		Danube3-500M
Layout	What moves	$J$	KL@1000	$J$	KL@1000
anti-optimised	everything, $J$ minimised	8.42	0.515	4.91	0.831
contiguous (3 seeds)	nothing	13.99	0.415	7.46	0.738
random (8 layouts)	everything, random	$\approx 13.9$	0.374	$\approx 7.7$	0.737
4 layers	evenly spaced layers, incl. layer 0	15.27	0.425	8.46	0.678
8 layers		16.46	0.397	9.67	0.677

Table 3: Ablation layouts (single seed unless noted), ordered by Qwen3-0.6B score. Score

J

is the co-located fo affinity (equation 6 ); its random expectation is 13.75 on Qwen3-0.6B (55 edges) and 7.75 on Danube3-500M (31 edges). KL after 1000 steps. Differences below 0.01 are within evaluation noise.

Figure 5: Qwen3-0.6B ablations. Left: KL at step 1000 for every layout, ordered by score (in the tick labels); shaded bands are mean

\pm

sd of the three optimised (blue) and three contiguous (orange) seeds. Middle: KL curves of the component ablations against the two seed means. Right: KL at 1000 steps against the number of layers optimised (evenly spaced, always including layer 0).

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Danube3-500M seed curves, as Figure 2 .

Figure 7: Danube3-500M 10k-step runs, as Figure 4 . The relative gap falls from 16% to 4.6% but the optimised run is ahead at all 40 evaluations; the multiplier is 1.6 to 2.4

\times

at the tabulated steps and its rolling median of five stays between 1.45 and 2.7

\times

Figure 8: Danube3-500M ablations, as Figure 5 . Neurons-only is the best 1000-step layout; heads-only equals contiguous; four evenly spaced layers give 85% of the gain.

Figure 9: Qwen3-0.6B: co-located fraction of the up edge (heads to same-layer neurons) per layer for each optimised layout, random $=0.25$ . Layer 2, which hosts the massive-activation circuit of § 5.6 , is fully separable: the optimiser co-locates 100% of its affinity, the anti-optimised layout 0%, and the head sweep alone already reaches 0.74 with contiguous neurons. Eight layers reach 0.6 or more and none is below 0.41.

Figure 10: Danube3-500M: as Figure 9 . Layers 1 and 15 are the only ones above 0.6; layer 1 is the only layer where the head sweep adds to the neuron assignment (0.78 with contiguous heads, 0.92 with optimised heads).

Figure 11: Left: KL at step 0 against KL at step 1000 for every 1000-step run (Qwen3-0.6B top, Danube3-500M bottom). Right: full KL trajectories on a log scale for the contiguous and optimised seed-1 runs and the best and worst random layouts untrained. On Danube3-500M the contiguous layout is the best untrained layout of the seventeen and the optimised one is twelfth; by step 100 the order has flipped and it never flips back.

Tensor Parallelism

Tensor parallelism splits neural-network computations across GPUs to fit larger models and reduce latency through concurrent execution.

See all 20 papers

Apr 29, 2026

Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

We present tensor and sequence parallelism (TSP), a parallel execution strategy that folds tensor parallelism and sequence parallelism onto a single device axis. In conventional multi-dimensional parallelism layouts, tensor parallelism (TP) shards model weights while sequence parallelism (SP) shards tokens, reducing per-device parameter or activation memory, respectively. Traditionally, each scheme is assigned its own mesh dimension. TSP instead assigns each rank both a weight shard and a sequence shard, reducing both parameter and activation memory along the same device axis. We implement this design with two runtime schedules. For attention, ranks iterate over broadcast parameter shards and reconstruct context through a sequence-wise key/value exchange. For gated MLPs, weight shards circulate in a ring while partial outputs accumulate locally. By sharding both weights and activations across the same devices, TSP trades additional communication volume for reduced memory overhead. We provide a theoretical communication and memory analysis, describe our implementation of TSP attention and gated MLP blocks, and benchmark TSP against TP, SP, and TP+SP. These results position TSP as a hardware-aware alternative for long-context and memory-constrained model training, and as a viable axis of parallelism in concert with existing parallelism schemes such as pipeline and expert parallelism for dense and mixture-of-expert models.

Vasu Shyam, Anna Golubeva, Quentin Anthony

Zyphra

Transformer Inference Transformer Tensor Parallelism Memory-Efficient Inference

Ask PDF

Aug 24, 2026

ShardMeter: Sharded and Geo-Distributed Training Without the Guesswork

Training large-scale AI models often outgrows a single data center, demanding sharded, multi-cluster, and decentralized training. However, the huge space of resource allocations makes exhaustive benchmarking and manual tuning impractical, while performance depends on tightly coupled factors like model size, GPU memory, batch size, bandwidth, and sharding strategy. We introduce ShardMeter, a lightweight analytical performance model that predicts the end-to-end runtime of transformer-based workloads across arbitrary sharded, distributed, and even decentralized training. Given a model's characteristics and a target hardware topology, ShardMeter estimates per-GPU and per-island throughput, training cost, total wall-clock time, and identifies performance bottlenecks. Our analysis reveals diminishing-return regimes as island size increases, quantifies transitions between compute- and communication-bound scaling, evaluates hyperparameter trade-offs, and models cost-throughput for large-scale decentralized training. ShardMeter exposes these insights to quickly explore the configuration space, choose near-optimal deployment plans, and avoid costly trial and error.

Tim Beringer, Patrick Diem, Felix Wolf +1

Technical University of Darmstadt · PanocularAI

Distributed Training Decentralized Learning

Ask PDF

Oct 12, 2026

BARTER: Compact Dynamic Sparse Training with Standard Dense Kernels

Dynamic sparse training (DST) aims to reduce training costs while maintaining accuracy. Existing DST methods mask out parameters but keep matrix dimensions unchanged. Training speedups therefore depend on specialized support, such as sparse kernels. We introduce BARTER, a structured DST framework that reduces training cost using standard dense kernels by keeping training tensors physically compact while dynamically reallocating capacity across layers. Unlike existing works that require computing gradients for all parameters, BARTER instead uses statistics collected only from active units during ordinary backpropagation to prune and clone attention heads and feed-forward channels under separate fixed global budgets. At each topology update, it resizes the affected parameter tensors and remaps their optimizer states, so subsequent training steps operate directly on the compact model. Training and fine-tuning experiments with Transformer-based vision models on two image classification benchmarks at densities from

0.1

0.7

show competitive Top-1 accuracy relative to full-tensor DST baselines. On an NVIDIA L40, BARTER achieves speedups of up to

6.29\times

in measured training segments and reduces peak allocated GPU memory by up to

81.3\%

relative to these implementations. An iPhone 13 Pro Max training benchmark additionally shows a

2.78\times

speedup over the full-tensor DST baseline.

Yueyuan Sui, Payal Mohapatra, Qi Zhu +1

Northwestern University, USA · Arm Inc., USA

Dynamic Sparse Training Neural Network Pruning Memory-Efficient Training Efficient ViTs Attention Head Pruning

Ask PDF

Abstract

Tensor Parallelism

Explore similar work

Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

ShardMeter: Sharded and Geo-Distributed Training Without the Guesswork

BARTER: Compact Dynamic Sparse Training with Standard Dense Kernels