[RFC] Deterministic training and versioned Megatron bitwise compatibility for production-common configurations

Summary And Motivation

We propose an opt-in deterministic training capability for DeepSpeed, with an
explicit, versioned numerical compatibility contract targeting pinned official
Megatron-LM training. The long-term objective is comprehensive bitwise
alignment across a maintainer-agreed matrix of mainstream production dense
and MoE configurations, not just repeatable loss values on small probes.

Repeated-run determinism is necessary but insufficient: two individually
deterministic implementations can differ because floating-point addition is
not associative, reduction geometry differs, or casts/fusions occur at different
boundaries. Matching Megatron does not imply that its arithmetic is more correct.
The reference is a reproducible numerical contract, not a universal quality ranking.

This would give users a practical way to validate framework migrations, diagnose
training regressions, and assess changes to kernels, optimizers, parallelism and
communication scheduling. We use OpGuard
to localize first differences and distinguish implementation bugs from valid
numerical-policy differences. OpGuard is a development/diagnostic tool, not a
required synchronized dependency in normal training.

This is a proposal for discussion under CONTRIBUTING.md, not a feature PR.
An exploratory local prototype exists; its limited validation below is not a
claim of production-wide support or upstream acceptance.

Proposed User Contract

There are two related but distinct guarantees:

1. Repeatability: independent runs under the same supported configuration,
   initialization, data order, RNG policy, topology and software/hardware stack
   produce identical tensor bytes at defined boundaries.
2. Megatron compatibility: a supported DeepSpeed configuration mapped to a
   pinned official Megatron numerical policy produces identical logical outputs,
   optimizer-consumed gradients, model/FP32 master weights and optimizer
   moments/counters at reviewed corresponding boundaries.

The final config surface is open for discussion. The prototype has a guarded
deterministic_mode flag and a separate research compatibility preset. We do
not propose that enabling repeatability silently selects every Megatron policy.
A versioned compatibility profile, separate from repeatability, appears useful.
Unsupported cells should fail clearly, or be explicitly labelled unverified;
they must not silently receive a bitwise-compatibility guarantee.

Initial scope is CUDA BF16 Transformer training with AdamW/FusedAdam and FP32
master weights, accumulation and communication. The eventual production matrix
should cover dense/MoE, supported ZeRO0--3 combinations, TP/SP and EP, GAS,
clipping, auxiliary losses, activation recomputation, common kernels, scheduler
and update-boundary checkpoint/resume. Additional PP/context parallelism,
offload, compilation and accelerators require separately agreed contracts and
acceptance, not assumptions inherited from the initial cells.

Bitwise guarantees are scoped to pinned versions and hardware/topology. They
are not cross-GPU-architecture, cross-NCCL-version or topology-changing trajectory
guarantees. ZeRO is a storage strategy, not a one-to-one Megatron config flag;
DeepSpeed Ulysses and Megatron TP-internal SP are not interchangeable semantics.

Exploratory Progress

The current evidence has two distinct reference tiers. These results must not
be combined into an untested Cartesian product of features.

Shared research helpers now implement canonical gradient mapping, matched
DP/EDP buffers, DP-before-SP-TP reduction and native dense/expert clipping norm
decomposition. Native ZeRO hook installation and partition writes remain in the
entry adapter. Moving these behind maintainable native interfaces is unfinished.

After helper extraction, 12 serial large-probe full/save/resume launches pass
direct complete old-tensor baseline comparisons and 96 state-manifest
comparisons. The indexed-data lifecycle adds 12 full/save/resume launches with
160 state-manifest/sample-prefix/sample-suffix/repeat comparisons. Manifests
retain every tensor's dtype, shape, byte count and SHA256 plus recursive scalars;
digest comparison is labelled separately from direct tensor-byte comparison.
These report counts are not counts of independent configurations.

Runtime checks cover optimizer state, counters, scheduler and CPU/CUDA/model-
parallel RNG; the indexed-data lifecycle also covers Python/NumPy RNG and the
actual consumer sample cursor. Prefetched dispatch counts are not treated as
consumed samples. Deterministic worker queues are reconstructed on resume;
arbitrary stochastic worker augmentation/queue-state recovery remains unverified.
Cross-framework mapping of every runtime field is also not complete.

Environment And Reproducibility Limits

Investigated baseline: DeepSpeed fc7bc301f; official Megatron core_v0.15.3
(309ffca); OpGuard 146d6bb523cf8364bbf2788cb319aef1a46f966b.
Runs used 2/4/8 NVIDIA L20 SM89 GPUs on a 16-GPU host, driver535.261.03,
PyTorch2.9.1/CUDA13.1, NCCL2.27.5, Ring collectives and a compatible SM89 Apex
build. Buckets, kernels, source provenance and commands are retained per run.
Grouped evidence uses official legacy grouped_gemm0.3.0 and candidate Torch
grouped execution, not TEGroupedMLP or a newer grouped backend.

Eight-rank ETP1 folding has actual native expert-data groups of size4.
It is not acceptance of the native ETP2/EDP2 configuration.

The prototype, reports and large tensor/trace artifacts are currently local or
in a private evidence archive; this issue does not present that archive as a
public reproducibility package. A portable public example and sanitized,
versioned commands/results are a next-stage deliverable. Large excluded traces
need separate artifact publication. Some compatibility shims and original
failed runs are recorded; success does not erase unsupported upstream arms.

Bugs Versus Policy Differences

Alignment investigation produced separate focused reports:

• [BUG] Low-precision optimizer initialization advances Adam and applies weight decay before training #8910: low-precision optimizer initialization advances Adam/apply weight decay.
• [BUG] ZeRO stage 1/2 marks optimizer-state fragments linked before later groups initialize #8911: ZeRO1/2 marks optimizer-state linkage complete before all groups initialize.
• [BUG] get_global_norm_of_tensors crashes on ranks with an empty norm contribution #8912: legitimate empty norm contribution crashes after replica deduplication.
• [Precision] BF16 AutoEP folding rounds FP32 TP reduction before FP32 gradient accumulation #8913: BF16 folding round-trip loses FP32 bits before FP32 accumulation; explicitly conditional on the precision contract.
• Per-parameter FusedAdam counters overlap with existing fix: track FusedAdam bias correction per parameter #8471; we will coordinate with that PR rather than duplicate it.

Different valid DP/TP order, buffer packing, norm decomposition, fusion/cast
placement and GAS reduction/accumulation order are compatibility policies,
not automatically bugs. Locally introduced overlap lifetime and auxiliary
folding-scale defects are recorded as draft/harness defects, not upstream
DeepSpeed defects. Genuine correctness fixes should be reviewed independently
of the compatibility flag and should not be hidden behind it.

Static-Agent Triage And OpGuard First-Difference Diagnosis

Our investigation uses a two-stage workflow:

1. Static-agent quick scan. An AI coding agent performs a bounded source
   scan around the earliest known mismatching interval: optimizer initialization,
   dtype/cast boundaries, accumulation, partition writes and collective
   ordering. If a simple source-level issue can be established by a small,
   discriminating reproducer, validate it directly. Static findings are
   hypotheses until supported by execution evidence; do not keep guessing
   cleanup or synchronization changes when this scan cannot explain the bytes.
2. OpGuard first-difference localization. When static triage cannot establish
   the cause, capture reviewed corresponding boundaries and find the earliest
   observed differing input/output. For GAS, include local backward, FP32
   accumulation, reducer inputs and partition writes, not just collectives.
   Replay captured equal inputs under one changed policy at a time, then rerun
   the original workload without diagnostic synchronization to validate the fix.

The role of the static agent is inexpensive triage, not a correctness oracle.
OpGuard also is not sufficient on its own: event correspondence must be checked,
exact bytes compared, and controlled interventions used to establish causality.
An earliest observed difference is only relative to the capture coverage.

Historical Evidence Of Diagnostic Effectiveness

The following cases provide concrete evidence of localization and causal
validation, rather than merely counting passing tests:

These cases show that the workflow can separate forward/backward correctness,
collective policy, clipping geometry and partition lifecycle, and can distinguish
valid numerical differences from local bugs. They support a qualitative claim
of effective diagnosis with explicit causal controls. We have not collected a
systematic static-only versus OpGuard time-to-root-cause comparison, so we do
not claim a measured debugging speedup or universal automatic diagnosis.

There is also an important negative example: synchronous OpGuard collective
captures suppressed a draft overlap storage-lifetime race. Non-copying
dispatch/event observations were needed to show allocator reuse before reduction
completion; adding only record_stream fixed uninstrumented runs without an
extra stream wait. Per-call captures inside Megatron's coalescing manager also
initially observed outputs before execution; corrected start/finish-boundary
captures preserved the original first difference. Thus capture placement and
perturbation checks are part of the workflow, not details to ignore.

The reports behind these cases are retained locally/private as
full-moe/z2-opguard-investigation/, production/eightcard-actual-layout-replay/,
acceptance/grouped-norm-replay.json, acceptance/gas-opguard/report.json
and overlap-rootcause/FINDINGS.md. They are not public artifact links yet.
The public replay milestone will include sanitized per-case traces, exact
commands, original failures and controlled replay results. A prospective
evaluation can additionally record static-triage budget, unresolved hypotheses,
capture coverage, first-difference location, time to causal reproduction and
uninstrumented acceptance, without assuming that every case requires OpGuard.

Rough Implementation Design

• Keep default DeepSpeed behavior unchanged. Isolate opt-in logic in shared
  modules; use small engine/optimizer/ZeRO wiring changes, not pervasive hot-path
  branches or a second training framework embedded in DeepSpeed.
• Define a versioned numerical profile describing cast/fusion boundaries,
  routing/activation/weighting, optimizer equations, accumulation order,
  clipping ownership and norm packing, DP/EDP/TP scaling and reduction geometry.
• Use explicit process-group/mesh metadata, canonical parameter correspondence
  and buffer descriptors. Reuse AutoTP/AutoEP and checkpoint work rather than
  creating another mesh or conversion stack.
• Integrate compatible layouts with native ZeRO bucket/partition mechanisms.
  Complete replicated FP32 buffers in the current research adapter are a
  diagnostic bridge, not an acceptable claim of scalable ZeRO reduce-scatter.
• Save reference-profile identity, topology/configuration manifest and numerical
  state with checkpoints. Define ownership between framework state and entry
  data/RNG state; reject incompatible restores. Start at update boundaries.
• Preserve overlap when dependence and storage lifetime are proven. An overlap
  option or stream alone is not concurrency evidence; avoid diagnostic
  synchronization in final correctness runs and benchmarks.

Next Stages And Acceptance Gates

1. Native integration and public replay. Publish portable examples and
   sanitized evidence; agree the profile/config API and parameter/buffer mapping.
   Integrate shared reducer/norm/aux policies with native AutoTP/AutoEP/ZeRO,
   removing research monkeypatches without changing the pinned oracle. Revalidate
   original cells; defaults and unsupported cases remain explicitly separate.
2. One combined production lifecycle cell. Combine BF16/FP32 master,
   TP2 EP2, ZeRO2/3, GAS2/4, active general clipping, grouped experts, router aux,
   full recompute, nonconstant scheduler and actual multiworker indexed data.
   Verify independent runs and fresh-process resume across epochs, sample IDs,
   model/master/moments/update counts, RNG and consumed cursor. Specify stochastic
   worker/prefetch semantics before claiming random augmentation recovery.
3. Topology and kernel expansion. Establish actual TP2 EP2 EDP2 semantics
   separately from current ETP1 folding; expand TP/EP/EDP sizes and multi-node
   groups, uneven/idle experts and bucket/oversized-parameter regimes. Validate
   GAS+overlap independently. Add production grouped/attention backends,
   PP/context parallelism and offload in maintainer-prioritized stages.
4. Production-scale long runs, quality and performance. At least 1,000
   updates per representative dense/MoE cell, independent repeats and midpoint
   new-process restore; streamed exact comparisons at agreed boundaries with no
   OpGuard diagnostic synchronization. Include larger realistic model/sequence
   sizes and a real-data training/validation task. Follow CONTRIBUTING.md full-
   training quality requirements where arithmetic changes affect convergence.
   Report serial unchanged-DeepSpeed / compatible-DeepSpeed / official-Megatron
   benchmarks, twice each, fixed workload and stack, slowest-rank mean/median/p95,
   tokens/s, peak memory and variability. Separate matched policies from tuned
   defaults, physical overlap evidence from flags, and failed arms from passes.

Release acceptance should be a published finite production matrix with exact
versions, exclusions, coverage and performance costs, not "all configs work."
The aspiration is comprehensive mainstream coverage of that agreed matrix.
No blanket claim should precede combination, long-run and scalable-path evidence.

Current timings do not establish a production performance win. Small-model
three-arm oracle comparisons can be slower than Megatron; a two-repeat official
1,000-update timing exists, but is neither a thousand-step byte acceptance nor
a three-arm production benchmark. Some historical timing batches have provenance
or shared-host limitations and are not used to claim certified rankings.

Alternatives And Coordination

Seed settings / torch.use_deterministic_algorithms alone provide neither
cross-framework arithmetic alignment nor complete checkpoint data lifecycle.
Tolerance-only comparisons remain useful for quality/performance exploration,
but cannot substitute for this bitwise contract. Globally serializing every
collective may simplify debugging but is not the proposed production endpoint.

Searches of issues and open PRs for deterministic/bitwise/Megatron did not find
an existing complete proposal with this contract. There is related work in
the Q4 roadmap #8672, AutoEP/AutoTP folding #8670, checkpointing #8848 and
optimizer fixes #8471. This proposal adds a numerical contract, diagnostic
workflow and acceptance matrix; it should coordinate with those owners rather
than duplicate their implementations.

Maintainer feedback requested: initial production models/configs to prioritize;
whether a separate versioned compatibility profile is the preferred API;
native bucket/mesh integration boundaries; and CI hardware/maintenance ownership.
Contributor contact: @ziming-zh. This prototype, experiments and proposal were
prepared with AI assistance; upstream code contributions require human review
and normal project contribution requirements.