{
  "schemaVersion": 1,
  "id": "atlas-lesson-distributed-gradient-weighting",
  "title": "Uneven batches and the global gradient",
  "category": "learning",
  "format": "Atlas lesson",
  "description": "Preserve the intended training objective when ranks contribute different numbers of valid examples or tokens.",
  "canonicalUrl": "https://semiaifoundry.com/learn/atlas/distributed-gradient-weighting/",
  "publishedDate": null,
  "revisionDate": "2026-10-07",
  "author": null,
  "topic": "Executable software",
  "series": null,
  "license": null,
  "rightsUrl": "https://semiaifoundry.com/ai-atlas/licenses.html",
  "version": "1.3.0",
  "status": "Published educational lesson",
  "evidenceScope": "The PyTorch source is pinned to v2.8.0 and supports the default averaging convention, not a claim about every backend or custom hook. The squared-error data, gradient weighting derivation and SGD step are authored examples. Join semantics after input exhaustion are outside this fixed-participant update.",
  "contentCoverage": "published-lesson",
  "artifactUrl": "https://semiaifoundry.com/ai-atlas/lesson-data/distributed-gradient-weighting-31b4b96883c8d7f0.json",
  "interactiveUrl": "https://semiaifoundry.com/ai-atlas/#distributed-gradient-weighting",
  "related": [],
  "supersedes": [],
  "supersededBy": [],
  "contentHash": "fa5db0229ca8551d9fa4a1a842352e5782e61de99b1855aa1f914540ea86451a",
  "contentLength": 8644,
  "jsonUrl": "https://semiaifoundry.com/agents/publications/atlas-lesson-distributed-gradient-weighting.json",
  "markdownUrl": "https://semiaifoundry.com/agents/publications/atlas-lesson-distributed-gradient-weighting.md",
  "sections": [
    {
      "id": "section-1",
      "title": "Uneven batches and the global gradient",
      "start": 36,
      "end": 520
    },
    {
      "id": "section-2",
      "title": "The mechanism",
      "start": 520,
      "end": 2594
    },
    {
      "id": "section-3",
      "title": "Science and assumptions",
      "start": 2594,
      "end": 2622
    },
    {
      "id": "section-4",
      "title": "Global item-mean gradient",
      "start": 2622,
      "end": 3396
    },
    {
      "id": "section-5",
      "title": "Compensate a mean-across-ranks reducer",
      "start": 3396,
      "end": 4260
    },
    {
      "id": "section-6",
      "title": "Engineering example",
      "start": 4260,
      "end": 5012
    },
    {
      "id": "section-7",
      "title": "Distinctions and tradeoffs",
      "start": 5012,
      "end": 5550
    },
    {
      "id": "section-8",
      "title": "Specifications and interfaces",
      "start": 5550,
      "end": 5902
    },
    {
      "id": "section-9",
      "title": "Interfaces",
      "start": 5902,
      "end": 6100
    },
    {
      "id": "section-10",
      "title": "Check your understanding",
      "start": 6100,
      "end": 6129
    },
    {
      "id": "section-11",
      "title": "Two ranks have counts 2 and 6 and mean gradients 1 and 3. What is the global item-mean gradient?",
      "start": 6129,
      "end": 6359
    },
    {
      "id": "section-12",
      "title": "A mean-across-two-ranks reducer receives local means. Which scale factors preserve the eight-item objective?",
      "start": 6359,
      "end": 6626
    },
    {
      "id": "section-13",
      "title": "Does equal synchronized state prove the chosen sample weights were correct?",
      "start": 6626,
      "end": 6850
    },
    {
      "id": "section-14",
      "title": "Evidence and scope",
      "start": 6850,
      "end": 7390
    },
    {
      "id": "section-15",
      "title": "Sources",
      "start": 7390,
      "end": 7892
    },
    {
      "id": "section-16",
      "title": "Continue exploring",
      "start": 7892,
      "end": 7915
    },
    {
      "id": "section-17",
      "title": "Foundations",
      "start": 7915,
      "end": 8293
    },
    {
      "id": "section-18",
      "title": "Related lessons",
      "start": 8293,
      "end": 8644
    }
  ],
  "citations": [
    {
      "title": "Explore in Atlas",
      "url": "https://semiaifoundry.com/ai-atlas/#distributed-gradient-weighting"
    },
    {
      "title": "PyTorch v2.8.0 — DistributedDataParallel implementation and contract",
      "url": "https://github.com/pytorch/pytorch/blob/v2.8.0/torch/nn/parallel/distributed.py"
    },
    {
      "title": "NVIDIA NCCL — Collective Operations",
      "url": "https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html"
    },
    {
      "title": "Published lesson data",
      "url": "https://semiaifoundry.com/ai-atlas/lesson-data/distributed-gradient-weighting-31b4b96883c8d7f0.json"
    },
    {
      "title": "Atlas releases and licensing",
      "url": "https://semiaifoundry.com/ai-atlas/licenses.html"
    },
    {
      "title": "Loss, gradients & optimizer",
      "url": "https://semiaifoundry.com/learn/atlas/loss-optimizer/"
    },
    {
      "title": "Training data pipeline",
      "url": "https://semiaifoundry.com/learn/atlas/training-data/"
    },
    {
      "title": "Data, tensor, pipeline, and state parallelism",
      "url": "https://semiaifoundry.com/learn/atlas/parallelism/"
    },
    {
      "title": "Collective communication library",
      "url": "https://semiaifoundry.com/learn/atlas/nccl/"
    },
    {
      "title": "Training state and activation rematerialization",
      "url": "https://semiaifoundry.com/learn/atlas/training-memory-state/"
    },
    {
      "title": "Evaluation & the AGI frontier",
      "url": "https://semiaifoundry.com/learn/atlas/evaluation/"
    },
    {
      "title": "Floating-point reductions and compiler permissions",
      "url": "https://semiaifoundry.com/learn/atlas/floating-point-reduction-order/"
    }
  ],
  "recordHash": "bebd393548a344aa4a027459440f57ebd985e56e9a38a2e06463feb8eaf2fd99",
  "contentFormat": "text/markdown",
  "content": "Executable software · Atlas v1.3.0\n\n# Uneven batches and the global gradient\n\nPreserve the intended training objective when ranks contribute different numbers of valid examples or tokens.\n\n[Explore in Atlas](https://semiaifoundry.com/ai-atlas/#distributed-gradient-weighting)Logical values and distributed ownership; no physical scale or vendor timing\n\nLearning objectiveDerive and test a globally normalized gradient from unequal local batch means, then distinguish objective equivalence from replica synchronization.\n\n## The mechanism\n\nA synchronized gradient can be consistent across replicas and still represent the wrong objective. Suppose the intended loss is the mean over all valid training items in one optimizer update. A rank that averages two items and a rank that averages six have already divided by different denominators. Averaging their two gradient means assigns the smaller batch more influence per item. The global item mean instead weights each local mean by its valid-item count. PyTorch’s pinned DistributedDataParallel contract documents averaging gradients between workers, so the loss normalization and collective normalization must be considered together. The count can mean examples, unmasked target tokens or a sum of explicit sample weights; choosing it is part of the objective, not a hardware detail. This lesson derives one scalar update from a concrete squared-error loss. It assumes the same current parameter at both ranks and no cross-sample operation such as batch-dependent normalization. A rank exhausting its entire input stream is a separate synchronization problem: matching collective participation does not by itself fix per-item weighting.\n\n1.  Define which items count in the loss and record their global count N for this optimizer update. Exclude padding only when the objective says to exclude it.\n2.  Compute local loss sums or local means and retain the valid count associated with each reduction.\n3.  For a rank-mean reducer over P active participants, scale a local mean gradient by P nᵣ/N before the averaging collective. Equivalently reduce local sums and divide once by N.\n4.  During gradient accumulation, use the count over the entire update window rather than silently averaging unequal microbatch means.\n5.  Apply the same synchronized gradient and optimizer state transition at every replica, then compare the result with the unpartitioned reference objective.\n6.  Handle empty local contributions and exhausted ranks with explicit participation rules; never divide an empty local mean by zero or simply omit a required collective.\n\n## Science and assumptions\n\n### Global item-mean gradient\n\n$$\nN=\\sum_{r=0}^{P-1}n_r,\\qquad g=\\frac{1}{N}\\sum_{r=0}^{P-1}n_r\\bar g_r\n$$\n\nSum local valid-item counts to N. Weight each local mean gradient by its count and divide their sum by N.\n\nSymbols and units\n\nN\n\nTotal valid equally weighted items in the update\n\ncount\n\nnr\n\nValid items at rank r\n\ncount\n\nP\n\nParticipating rank count\n\ncount\n\nḡr\n\nMean local gradient at shared parameter state\n\nloss/parameter\n\ng\n\nGlobal item-mean gradient\n\nloss/parameter\n\nEach local ḡᵣ is the mean of gradients for nᵣ valid equally weighted items. Multiply by nᵣ to recover the local sum.\n\nAssumptions N>0, same parameters and differentiable separable per-item loss. No cross-sample dependence; sample weights replace counts when the objective uses unequal weights.\n\n### Compensate a mean-across-ranks reducer\n\n$$\ng=\\frac{1}{P}\\sum_{r=0}^{P-1}\\left(\\frac{P n_r}{N}\\bar g_r\\right),\\qquad w_{\\mathrm{new}}=w-\\eta g\n$$\n\nBefore a mean-across-ranks reduction, scale local mean gradients by P times $n_{r}$ over N. Scalar SGD subtracts learning rate times global gradient.\n\nSymbols and units\n\nP\n\nFixed participants in the mean reducer\n\ncount\n\nnr, N\n\nLocal and global valid-item counts\n\ncount\n\nḡr, g\n\nLocal and global mean gradients\n\nloss/parameter\n\nw, wnew\n\nParameter before and after the update\n\nparameter\n\nη\n\nLearning rate for the declared objective scaling\n\nparameter²/loss\n\nThe factor P nᵣ/N ensures the subsequent division by P leaves the intended global item mean.\n\nAssumptions P fixed participating ranks for this update; no custom reducer scaling. The scalar SGD update has no momentum, clipping, weight decay or loss scaling.\n\n## Engineering example\n\nTake ℓ(w;y)=½(w−y)² at w=10. Rank 0 has targets \\[9,9\\], producing gradients \\[1,1\\], count 2 and mean 1. Rank 1 has \\[7,7,7,7,7,7\\], producing six gradients of 3, count 6 and mean 3. The unpartitioned objective has gradient (2×1+6×3)/8=2.5. Averaging rank means gives (1+3)/2=2 instead. With learning rate 0.1, correct SGD gives wnew=9.75; the unweighted rank mean gives 9.8. For a reducer that divides by P=2, scale the local means by 2×2/8=0.5 and 2×6/8=1.5. It then averages 0.5 and 4.5 to obtain 2.5. If the ranks each held four valid items, equal weighting would be appropriate. Token-level training needs the valid target-token count when that is the declared loss denominator; sequence counts need not be proportional.\n\n## Distinctions and tradeoffs\n\nA mean across ranks is a mean across training items only when local denominators and weights make it so.\n\nKeeping exact counts and a reference objective adds bookkeeping and sometimes a small count reduction. It prevents an apparent scaling improvement from silently changing sample influence.\n\nThe number of completed updates is useful only relative to the training objective those updates optimize. Ragged data, masking and uneven accumulation can change that objective through an unnoticed denominator.\n\n## Specifications and interfaces\n\nNormalization unit\n\nValid examples or tokens\n\nChoose the intended loss denominator explicitly.\n\nReplica consistency\n\nEqual synchronized gradients\n\nDoes not establish agreement with the desired global objective.\n\nRuntime scope\n\nDefault mean reduction\n\nCustom hooks, joined ranks or summed losses can change the scale.\n\n### Interfaces\n\nPer-rank loss contributions Input\n\nLocal sums or means, valid counts, shared parameter state\n\nGlobal optimizer update Output\n\nOne correctly normalized gradient applied consistently\n\n## Check your understanding\n\n### Two ranks have counts 2 and 6 and mean gradients 1 and 3. What is the global item-mean gradient?\n\n1.  2\n2.  2.5\n3.  4\n\nAnswer and reasoning\n\n2.5\n\nThe weighted numerator is 2×1+6×3=20 across eight items, so the result is 2.5.\n\n### A mean-across-two-ranks reducer receives local means. Which scale factors preserve the eight-item objective?\n\n1.  1 and 1\n2.  2 and 6\n3.  0.5 and 1.5\n\nAnswer and reasoning\n\n0.5 and 1.5\n\nScaling by P nᵣ/N gives 0.5 and 1.5; the reducer then averages 0.5 and 4.5.\n\n### Does equal synchronized state prove the chosen sample weights were correct?\n\n1.  Yes\n2.  No\n3.  Only if the network is fast\n\nAnswer and reasoning\n\nNo\n\nAll replicas can agree on the same incorrectly normalized gradient.\n\n## Evidence and scope\n\nThe PyTorch source is pinned to v2.8.0 and supports the default averaging convention, not a claim about every backend or custom hook. The squared-error data, gradient weighting derivation and SGD step are authored examples. Join semantics after input exhaustion are outside this fixed-participant update.\n\nThe ideal ring all-reduce lab varies communication payload, ranks, startup and bandwidth. It does not execute these matrix shards or alter the optimizer; the exact computation is worked through in this lesson.\n\n## Sources\n\n-   [PyTorch v2.8.0 — DistributedDataParallel implementation and contract](https://github.com/pytorch/pytorch/blob/v2.8.0/torch/nn/parallel/distributed.py)\n-   [NVIDIA NCCL — Collective Operations](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html)\n\n[Published lesson data](https://semiaifoundry.com/ai-atlas/lesson-data/distributed-gradient-weighting-31b4b96883c8d7f0.json) · [Atlas releases and licensing](https://semiaifoundry.com/ai-atlas/licenses.html)\n\n## Continue exploring\n\n### Foundations\n\n-   [Loss, gradients & optimizer](https://semiaifoundry.com/learn/atlas/loss-optimizer/)\n-   [Training data pipeline](https://semiaifoundry.com/learn/atlas/training-data/)\n-   [Data, tensor, pipeline, and state parallelism](https://semiaifoundry.com/learn/atlas/parallelism/)\n-   [Collective communication library](https://semiaifoundry.com/learn/atlas/nccl/)\n\n### Related lessons\n\n-   [Training state and activation rematerialization](https://semiaifoundry.com/learn/atlas/training-memory-state/)\n-   [Evaluation & the AGI frontier](https://semiaifoundry.com/learn/atlas/evaluation/)\n-   [Floating-point reductions and compiler permissions](https://semiaifoundry.com/learn/atlas/floating-point-reduction-order/)"
}