FSDP2 optimizer state is not DCP-shardable (no dcp.save/load, no optimizer-state resharding)

[thad0ctor]

Summary

Gefen's FSDP2 optimizer state cannot be checkpointed through torch.distributed.checkpoint (DCP) save()/load(), and cannot be resharded across a different GPU count / topology. It resumes only through the in-memory full-state gather path, and only on the same topology.

This is a pre-existing design limitation (documented in COMPATIBILITY.md → Optimizer checkpoint scope), filed here to track a proper fix. It is independent of the CPU-offload work in #80 — it was re-surfaced while adding a resume test there.

What doesn't work

dcp.save/dcp.load of the optimizer state fails. Gefen packs its FSDP2 optimizer state as a variable-size rank-local payload (_gefen_rank_local_payload_*) because, under fully_shard, each rank learns its own codebook and block geometry, so the per-rank state genuinely differs in size. DCP's fixed-shape load planner rejects it:

ValueError: Size mismatch between saved torch.Size([15484]) and current:
torch.Size([2861]) for optimizer.state.0.weight._gefen_rank_local_payload_0

(15484 = a warmed optimizer's payload; 2861 = a fresh template's — the planner requires them to match.)

Consequences:

• No sharded / parallel optimizer-state disk checkpoint (each rank writing its own shard in parallel).
• No optimizer-state resharding — save on N GPUs, resume on M is not supported; resume requires the same GPU count and sharding layout.

What does work (supported path)

• Model weights via DCP — unaffected; weights are ordinary DTensors, so dcp.save/load of the model shards works normally.
• Optimizer state via full-state gather — get_state_dict(full_state_dict=True, cpu_offload=True) → torch.save → torch.load → set_state_dict(broadcast_from_rank0=True). Same-topology only. This is what tests/test_gefen_fsdp2_checkpoint.py and the PR CPU-offloaded training compatability: CPU-resident stepping + DeepSpeed ZeRO CPU-offload #80 resume test use.

Why

Inherent to the quantized-codebook design under FSDP2: per-rank heterogeneous codebooks/geometry ⇒ variable-size per-rank state ⇒ incompatible with DCP's uniform, fixed-shape, shard-addressable planner. The rank-local blob is the current workaround that preserves correctness at the cost of DCP-shardability.

Impact

• Fine-tuning / same-topology resume: fine — the full-state path covers it.
• Large-scale distributed pretraining with elastic / resharded checkpoints: blocked.
• Checkpoint-time memory: the full-state gather concentrates the entire optimizer state on rank-0 CPU (~1 byte/param, e.g. ~70 GB for a 70B model).

Possible direction

Give the FSDP2 optimizer state a DCP-compatible representation — fixed-shape per-rank tensors plus a resharding-aware planner / Stateful implementation — so dcp.save/load and N→M resharding work. There is an existing unmerged branch, feat/dtensor-portable-resharding ("Add sharded portable DCP resharding", "Prove checkpoint resume from two ranks to four"), that appears to target exactly this; it should be validated and landed (or superseded).

Acceptance criteria

• dcp.save/dcp.load round-trips the FSDP2 optimizer state without a size-mismatch error.
• Optimizer-state resume across a changed GPU count (e.g. 2 → 4) is exact.
• Checkpoint-time memory does not require gathering the full optimizer state onto a single rank.
• Regression test exercising a real dcp.save/dcp.load disk round trip (currently impossible, hence the in-memory path in tests/test_fsdp2_cpu_offload.py / test_gefen_fsdp2_checkpoint.py).

[thad0ctor]

Resolved by #83 (merged as df3a615), which adds GefenDCPState — a standalone torch.distributed.checkpoint adapter for plain Gefen's FSDP2 optimizer state.

Against this issue's acceptance criteria:

• dcp.save/dcp.load round-trips the FSDP2 optimizer state without the variable-size rank-local-payload size mismatch — save emits dense fp32 Shard(0) DTensors (dequantized against the real learned per-rank codebook) and load re-blocks them back into Gefen's compact ~1 byte/param form. Covered by a real disk round-trip test.
• N→M resharding (2→4 and 4→2, incl. a real 4-GPU NCCL reshard) — verified against a native reference run at the target topology.
• No single-rank gather at checkpoint time — DCP handles sharded I/O directly, replacing the full-state rank-0 gather.
• Regression tests for a real dcp.save/dcp.load disk round trip — added.

One honest deviation from the wording here: the resharded resume is a correct continuation within ~256-level quantization noise, not bit-exact, even same-topology. That is inherent to routing momentum through a dense reshard and re-blocking against a freshly learned per-shard codebook; it is documented, and same-topology users are steered to the native get_state_dict/set_state_dict path, which remains bit-exact.

Scope is plain Gefen only — Muon and Hybrid fail closed at construction. Reshardable DCP for those is tracked in #84 (GefenMuon) and #85 (GefenMuonHybrid). PR #68 (feat/dtensor-portable-resharding), which targeted this same gap, is superseded by #83.