[RFC] Add optional CPU staging for LocalSGD averaging

[Skuld7451]

RFC 0003: Optional CPU staging for LocalSGD averaged parameters

Summary

Add an opt-in LocalSGD synchronization path that all-reduces one parameter at a
time and stages each averaged result on CPU until quorum commit. This lowers
peak accelerator memory at synchronization time in exchange for host memory,
device-host transfers, and serialized collective waits.

This proposal is algorithm-level TorchFT functionality. It contains no
backend-specific API.

Dependency Boundary

The implementation uses only the existing Manager.allreduce() and
Manager.should_commit() contracts plus Tensor/DTensor copy operations. It
does not depend on generalized stream utilities, process-group accelerator
discovery, or FSDP-safe quorum recovery. It can be adopted directly on the
common TorchFT baseline with no source merge-order requirement.

Running LocalSGD on a particular backend still requires that integration to
provide a usable process group for its tensors. That runtime requirement is
unchanged and does not make the CPU-staging implementation depend on a specific
process-group registration change.

Problem

The existing LocalSGD _average() path retains averaged copies for all model
parameters on the accelerator until quorum commit. Large models can exceed
available memory during synchronization even when the ordinary forward/backward
step fits. TorchTitan cannot solve this cleanly because LocalSGD owns averaging,
commit, and DTensor reconstruction.

Proposal

API change

The public LocalSGD constructor changes from:

LocalSGD(manager, model, optimizer, sync_every)

to:

def __init__(
    self,
    manager: Manager,
    model: nn.Module,
    optimizer: optim.Optimizer,
    sync_every: int,
    offload_averaged_parameters_to_cpu: bool = False,
) -> None:
    self._offload_averaged_parameters_to_cpu = (
        offload_averaged_parameters_to_cpu
    )

At torchft/local_sgd.py:88-115, the new
offload_averaged_parameters_to_cpu argument is appended after sync_every and
defaults to False. Existing positional and keyword calls remain valid and
continue to use the original synchronization path. Callers opt in by passing
offload_averaged_parameters_to_cpu=True; no return type, optimizer hook, or
checkpoint interface changes.

The generated API reference obtains this signature and parameter description
from the LocalSGD docstring. docs/source/local_sgd.rst additionally documents
the invocation and its memory/throughput trade-offs.

Sequential averaging and staging

At torchft/local_sgd.py:170-187, the opt-in path clones one local parameter,
waits for its TorchFT all-reduce, transfers it to CPU, and releases the
accelerator temporary before moving to the next parameter.

for parameter in parameters:
    averaged_parameter = extract_local_tensor(parameter)
    self._manager.allreduce(averaged_parameter).wait()
    averaged_parameters.append(averaged_parameter.detach().cpu())
    del averaged_parameter

if self._manager.should_commit():
    for parameter, averaged_parameter in zip(parameters, averaged_parameters):
        _copy_local_tensor_to_parameter(parameter, averaged_parameter)

If quorum does not commit, CPU-staged values are discarded and local model
parameters remain unchanged.

Tensor and DTensor writeback

At torchft/local_sgd.py:45-69, move a staged tensor back to the parameter's
device. For DTensor parameters, reconstruct a DTensor with the original mesh,
placements, shape, and stride before copying.

if isinstance(parameter, DTensor):
    parameter.data.copy_(
        DTensor.from_local(
            local_tensor,
            parameter.device_mesh,
            parameter.placements,
            shape=parameter.shape,
            stride=parameter.stride(),
        )
    )
else:
    parameter.data.copy_(local_tensor)

Semantics and Trade-offs

• Quorum start, should_commit, and rollback semantics are unchanged.
• Peak accelerator overhead changes from approximately all local averaged
  parameter copies to one local parameter copy plus transfer/runtime workspace.
• Host memory must hold all local averaged parameters until commit.
• Immediate .wait() serializes parameter collectives and may reduce throughput.
• CPU-to-accelerator writeback is synchronous (non_blocking=False) to avoid
  lifetime and ordering ambiguity at commit.
• This option offloads averaged model parameters only. It does not offload model
  weights, gradients, or optimizer state.

Theoretical memory reduction

Let s_i be the byte size of parameter i's local tensor on the current rank:

S = sum(s_i)       # all local parameter copies
M = max(s_i)       # largest single local parameter copy

The original path creates every copy and launches every all-reduce before it
waits, so its additional active accelerator memory during synchronization is
approximately:

Delta_device_original ~= S + W_parallel

W_parallel represents concurrent collective and runtime workspace. The CPU
staging path creates one copy, waits for its all-reduce, moves it to CPU, and
releases it before advancing. Its additional active memory is approximately:

Delta_device_offload ~= M + W_single

Ignoring allocator granularity and workspace differences, the ideal saving is
S - M, or a fraction 1 - M / S. Models with many similarly sized parameter
tensors approach a larger percentage; a single dominant tensor limits the
reduction to M.

This describes active memory in the synchronization interval. It does not
guarantee a matching reduction in memory_reserved() or in the whole-step
global peak. A caching allocator can retain blocks, and a higher forward,
backward, or optimizer peak can hide the LocalSGD reduction.

Theoretical costs

• Each synchronization point adds approximately S of resident CPU staging
  memory until quorum commit.
• A successful commit adds about S D2H plus S H2D traffic, or 2S total
  host-link traffic. A rejected commit skips H2D writeback but still stages D2H.
• All-reduce payload volume is unchanged, but the original path can launch
  several collectives before waiting. Immediate per-parameter .wait() calls
  reduce communication overlap.
• For one synchronization every K=sync_every training steps, average added
  step latency is roughly (T_offload_sync - T_original_sync) / K.
  sync_every=1 therefore exposes the worst throughput cost, while a larger
  interval amortizes it.

Compatibility

• Default False follows the existing _average() path byte-for-byte.
• Existing constructor calls require no changes.
• The option is accelerator-neutral and has no vendor-extension dependency.
• Plain Tensor and DTensor parameters are supported.
• Callers such as TorchTitan may expose a recipe/config flag, but the algorithm
  belongs in TorchFT because TorchFT owns LocalSGD commit semantics.

Non-goals

• Optimizer-state swapping or virtual optimizer support.
• Asynchronous/pipelined host transfers.
• Bucketized LocalSGD averaging.
• DiLoCo offload changes.

Validation

The complete LocalSGD unit-test module passes after adding direct coverage:

python -m unittest -v torchft.local_sgd_test
Ran 11 tests
OK

The three new tests cover quorum commit, quorum rejection without model
mutation, plain Tensor writeback through the commit test, and DTensor
reconstruction with the original mesh, placements, shape, and stride. The
existing default LocalSGD and DiLoCo tests remain green.

Wider downstream validation has also run two-replica FSDP LocalSGD training
with sync_steps=1 and completed three steps.

Downstream switch comparison

On 2026-08-18, a strict A/B run and a synchronization-interval probe were
added. The cases differed only in local_sgd_offload_to_cpu:

Item	Configuration
Hardware	One host, four accelerator devices, 64 GiB per device
Software	PyTorch 2.14.0.dev20260701 with a compatible third-party accelerator stack
Topology	2 replicas x 2 devices; FSDP2 within each replica and a custom process group across replicas
Model	DeepSeek-V4 debugmodel, 4 layers, 80,762,013 parameters
Training	Eager, sequence length 512, local batch 1, global batch 4, Full AC, AdamW
TorchFT	LocalSGD, sync_steps=1, init_sync=False
Reproducibility	Seed 42 with deterministic algorithms enabled
Revisions	TorchFT 593faa8, including the TF-3 change now published as aa7a29b; matched downstream integration revisions

Both cases used the same eager model path and the same matched downstream
integration revisions. This keeps unrelated model-kernel and framework API
differences outside the A/B comparison and does not alter the LocalSGD path.

Functional and end-to-end results used six steps per case, discarded step 1 as
warmup, and averaged ten throughput samples from steps 2-6 on rank 0 of both
replicas:

Metric	Original path	CPU staging	Difference
Training completion	6/6 steps	6/6 steps	Both passed
Mean throughput	919.0 tokens/s	789.2 tokens/s	-14.1%
Logged global peak memory_reserved	About 2.93 GiB	About 2.93 GiB	No visible reduction
Loss / gradient norm	Baseline sequence	Identical at every step	No difference

The synchronization probe synchronized the device and reset peak statistics
immediately before _perform_sync(), then measured entry active memory and
max_memory_allocated() for that interval. A second four-step run discarded
the first synchronization and used six stable samples from steps 2-4 across
both replicas:

Metric	Original path	CPU staging	Difference
Additional active HBM during sync	154.09 MiB	63.13 MiB	-90.97 MiB (-59.0%)
Mean _perform_sync() time	176.39 ms	235.16 ms	+33.3%
Active HBM after sync	Returned to entry value	Returned to entry value	No leak observed
Reserved HBM during sync	Unchanged	Unchanged	Allocator retained blocks

The measured 154.09 MiB -> 63.13 MiB follows the theoretical S -> M
shape. For this model, each synchronization also requires about 154.09 MiB of
host staging and a committed synchronization transfers about 308.19 MiB over
the host link in both directions. An unchanged global reserved peak does not
invalidate the feature: forward/backward and allocator caching dominate this
small four-layer run. The option can affect OOM viability when synchronization
copies are the global peak in a larger model, but a production recipe still
needs measurement with its real sharding, synchronization interval,
inter-node bandwidth, and CPU-memory budget.

Alternatives

TorchTitan could reimplement LocalSGD synchronization around TorchFT, but that
would duplicate quorum and rollback logic and make other TorchFT users unable to
use the memory-saving path. Bucketization could improve throughput but still
requires an explicit memory policy; it is a possible follow-up rather than a
replacement for the opt-in sequential path.

Rollout

Keep the option disabled by default. Framework integrations should expose it
only where synchronization-time memory is the limiting factor and should
measure host memory and sync latency before enabling it by default for a recipe.