[ys2025-AI]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Summary

AccelerateStrategy.save_optimizer_checkpoint / load_optimizer_checkpoint can crash with KeyError: 0 inside torch.distributed.checkpoint.state_dict when an optimizer checkpoint is saved after a no_grad evaluation of an FSDP2 model whose parameters are not all individually wrapped by fully_shard (VL models with a visual tower, tied-weight tail groups, etc.). This PR reshards any FSDP2 param group left unsharded before the state-dict conversion.

Root cause

Three behaviors combine:

1. accelerate shards the FSDP2 root param group with reshard_after_forward=None, which PyTorch's heuristic resolves to False (accelerate/utils/fsdp_utils.py:883-888); tied-weight tail groups are also created with reshard_after_forward=False (accelerate/utils/fsdp_utils.py:879).
2. torch FSDP2 swaps module params to full nn.Parameter on unshard (to_unsharded) and back to sharded DTensor on reshard (to_sharded). With reshard_after_forward=False, reshard() early-returns in the FORWARD state, so unsharding is only undone by the post-backward reshard. A no_grad evaluation never runs backward, so params stay materialized as full nn.Parameter, while the optimizer still holds the sharded DTensor it was switched to by fsdp2_switch_optimizer_parameters.
3. get/set_optimizer_state_dict maps optimizer params back to model params by object identity. The stale full-nn.Parameter objects are not in model.named_parameters(), so _get_optim_state_dict raises:

# torch/distributed/checkpoint/state_dict.py:941
group[_PARAMS] = [fqn_pid_mapping[pid] for pid in group[_PARAMS]]
# KeyError: 0

Training itself is unaffected (gradients flow through the unsharded params and are reduce-scattered into the sharded DTensors the optimizer holds); only the identity-based checkpoint conversion breaks.

Repro conditions: accelerate FSDP2 + a no_grad eval pass + params outside individually-wrapped groups. On Qwen3.6-27B (qwen3_5, VL) with 8-rank HSDP LoRA training this crashes at the first eval-then-save on every rank, with 220/1212 optimizer params missing from model.named_parameters().

Fix

Add AccelerateStrategy._reshard_unsharded_fsdp_params(model) — walk the model tree and call the public FSDPModule.reshard() API ("freeing the unsharded parameters if they are allocated"; a safe no-op for already-sharded groups) — and invoke it in the FSDP2 branches of both:

• save_optimizer_checkpoint, before get_optimizer_state_dict
• load_optimizer_checkpoint, before set_optimizer_state_dict (the load path has the symmetric identity-mapping problem when resuming right after an eval)

Cookbook-level workarounds are unnecessary after this.

Impact

• Fixes checkpoint save/load after eval for accelerate-FSDP2 runs (dense + VL + any model with non-individually-wrapped params)
• No behavior change for groups already sharded; cost is O(#FSDP groups) cheap no-op calls at checkpoint frequency

Experiment results

Environment: 8 × Ascend910 (64GB), CANN 9.0.0, torch 2.10.0 / torch_npu 2.10.0, accelerate 1.15.0, transformers 5.5.4.
Task: cookbook/transformers/fsdp2.sh LoRA SFT (FSDP2, fsdp_size=2 × dp_size=4, Muon optimizer, GSM8K_zh dataset).

Before (this PR) — training healthy, crash at first eval-then-save on every rank

[OSD-DIAG] after-first-step:  optimizer params missing from named_parameters: 0/1212
[OSD-DIAG] after-eval:        optimizer params missing from named_parameters: 220/1212  (all visual.* lora)
  opt: DTensor  shape=(1, 1152)    <- sharded, held by optimizer
  mdl: Parameter shape=(8, 1152)   <- full unsharded Parameter, left by the no-grad eval
torch/distributed/checkpoint/state_dict.py:941: KeyError: 0

After (this PR) — full run passes end to end, 0 exceptions

sh fsdp2.sh --log-interval 5   # 1000 GSM8K_zh samples, batch 8, 125 steps

step   5  train loss 0.93870
step  80  train loss 0.51889 | eval loss 0.47690 -> checkpoint-80 saved
step 120  train loss 0.48911 | eval loss 0.45929 -> checkpoint-120 saved
step 125  train loss 0.53532 -> last-checkpoint saved

checkpoints: checkpoint-40 / checkpoint-80 / checkpoint-120 / last-checkpoint
             (each: LoRA adapter + optimizer.pt + scheduler.pt)

Quick regression (8 steps, eval every 2): 4 evals + 4 checkpoint saves all pass; no behavioral change to the training loop.