Resuming Training on a Different GPU Count: What to Validate

A checkpoint that loads is not necessarily a checkpoint that resumes the same training run. Changing the GPU count can alter the data-parallel world size, shard ownership, samples per optimizer update, collective topology, random-number consumption, and the number of batches remaining in an epoch. The first task is therefore to declare which properties must survive the move. Model parameters and optimizer history usually must survive exactly in value, while throughput and rank-local layouts are expected to change.

Write a migration manifest before launching: source and destination world sizes, parallelism dimensions, framework and checkpoint format, model commit, optimizer type, scheduler position, precision mode, global samples or valid tokens processed, and the intended global batch. Separate data-parallel ranks from tensor or pipeline ranks. Eight GPUs arranged as tensor parallel 2 by data parallel 4 do not process eight independent data batches. That distinction determines the batch equation and the state layout you need to restore.

A full, unsharded state dictionary is conceptually portable but can demand substantial host memory and load time. A rank-local checkpoint may encode assumptions about the world size or partitioning used when it was written. PyTorch Distributed Checkpoint supports load-time resharding, and its state-dictionary APIs are intended to produce model and optimizer state that can be resharded across different trainer counts or parallelisms. That capability still depends on using the matching APIs and a checkpoint whose metadata and files are complete.

Test portability with a disposable destination job before the real restart. Allocate the destination model, load every model and optimizer key, and fail on missing or unexpected entries. Then compute stable summaries such as parameter count, tensor shapes, dtypes, and per-tensor checksums before and after reconstruction. A checksum mismatch may be expected if a format converts dtype, but it must be explained. Never infer portability from the presence of one metadata file while a rank shard is missing.

Adam-like optimizers carry moment tensors associated with parameters, while schedulers carry counters and sometimes additional state. PyTorch’s general checkpoint guidance explicitly saves the optimizer state dictionary because buffers and hyperparameters change during training. When parameters are resharded, optimizer tensors must follow the same logical parameters, not the old rank numbers. Reinitializing the optimizer turns a resume into a warm start and can produce a transient that looks like a precision or data problem.

After loading, inspect optimizer parameter groups, learning rate, weight decay, step counters, and state tensor shapes. Pick several named parameters from early, middle, and late layers and confirm that their moment statistics are finite and nonzero where expected. Run one update with logging before committing the job. The acceptance rule should be strict: all expected optimizer entries load, the scheduler reports the saved update index, and no state tensor changes until an optimizer step actually occurs.

Global batch equals microbatch per data-parallel replica multiplied by the data-parallel degree multiplied by gradient-accumulation steps. If the world size changes and the other factors remain fixed, the global batch changes too. That changes optimizer-update frequency, sample-based evaluation intervals, and the meaning of a step-based learning-rate schedule. Decide whether the invariant is global samples per update, valid tokens per update, or simply the original schedule position. Do not let the launcher choose by accident.

Consider an illustrative restart from 16 data-parallel GPUs with microbatch 2 and accumulation 8. The global batch is 16 x 2 x 8 = 256 samples. On 10 data-parallel GPUs, accumulation 8 would yield 10 x 2 x 8 = 160. A fixed microbatch of 2 on every rank cannot preserve 256 exactly because 256 / (10 x 2) = 12.8 accumulation passes. Options include a controlled batch change, a deliberately variable-sized microbatch with correct sample or token normalization and compatible collective control flow, or a destination data-parallel degree that divides the target. These numbers are illustrative and make no performance claim.

Rank-local sampler positions are fragile when the number of ranks changes. Prefer a global progress coordinate such as committed sample IDs, global valid-token count, or a deterministic shard cursor. The restart should map that coordinate onto the new rank set. Epoch and batch index alone can be ambiguous when padding, dropped tails, variable-length packing, or gradient accumulation is involved. Save the accumulation phase too, or checkpoint only after a complete optimizer update.

PyTorch DistributedSampler builds a deterministic ordering from a common seed plus epoch, pads or drops indices to divide work across replicas, then slices indices by rank. A new replica count changes that slicing and can change padding. For exact continuation, materialize or reconstruct the global order and resume after the last committed global unit. For an acceptable approximate continuation, document the possible replay or omission bound and verify it from emitted sample identifiers.

Bitwise equality across different topologies is often the wrong acceptance test. Collective reduction trees can change floating-point addition order, and different shard layouts can change kernel selection or random-number consumption. PyTorch notes that deterministic-algorithm mode alone does not guarantee complete reproducibility. The useful contract is usually functional: the first resumed loss, gradient norm, update magnitude, and short-window trajectory remain within declared tolerances and all values remain finite.

Use a two-level check. First, run an evaluation-only forward pass on fixed examples immediately before and after migration and compare logits or losses at a tolerance appropriate to the saved dtype. Second, run perhaps 20 optimizer updates with identical global inputs and compare scalar traces. A measurable decision rule could reject the migration if any nonfinite value appears, if the first update index differs, or if relative loss and update-norm deviations exceed a predeclared 1 percent. The threshold is illustrative and should reflect normal run-to-run variation.

Use a diagnostic sequence that limits wasted GPU time. First inspect checkpoint completeness without GPUs. Next load on the destination topology and run evaluation only. Then execute one optimizer update while capturing parameter delta, learning rate, loss scale, skipped-step status, gradient norm, and sample IDs. Finally run a short window long enough to cross one logging boundary and, if practical, one evaluation boundary. Compare against a control resume on the original topology.

Promote when state coverage is complete, global-batch behavior matches the declared choice, processed-sample continuity meets its bound, and the short-window metrics remain inside tolerance. Throughput should be measured separately because new communication geometry may make scaling disappointing even when correctness is sound. Record the validated parallelism layout and batch invariant, not merely a total GPU count. A short test cannot prove final convergence, so a material batch change still requires a representative quality run.

PyTorch 2.7 Distributed Checkpoint documentation

PyTorch Saving and Loading Models tutorial

PyTorch 2.7 deterministic algorithms documentation

FP4 Inference on Blackwell Cost Guide 2026: 8x Claims vs Production Reality

A/B Testing for AI Models: Production Inference Experimentation Framework