fix: restore local optimizer state without distributed broadcast

[Liuchenbing-2026]

PR type

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

PR information

Single-process NativeFSDPStrategy training can save optimizer state without an initialized distributed process group. Loading it nevertheless enables broadcast_from_rank0, which makes PyTorch query the default group and raise ValueError: Default process group has not been initialized.

Enable rank-zero broadcast for loading only when a process group is initialized. Saving behavior remains unchanged. Regression tests exercise actual AdamW checkpoint I/O and subsequent optimization both without a process group and with a real one-rank CPU/Gloo group. They compare restored learning rate, step and moment tensors, then require exact equality of the next parameter update against uninterrupted training. The existing placement test and broadcast-option control remain covered.

AI assistance was used for implementation and testing.

Experiment results

Compared the same tests on main 9953ea1ebbbdbe59c3831702157d3c147ca0ee26 and the fixed implementation:

Source	Result
Main	1 failed, 3 passed
Fixed	4 passed, no skips

Only the no-process-group checkpoint round trip fails on main. Both local and initialized one-rank round trips pass after the fix, including moment restoration and subsequent weights at rtol=atol=0.

PYTHONPATH=src PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 OMP_NUM_THREADS=2 python -m pytest \
  tests/transformers/test_native_fsdp_strategy.py -q --confcutdir=tests/transformers
pre-commit run --all-files

All pre-commit hooks passed; the test file also passed manual isort/yapf/flake8 checks because repository hooks exclude tests. Environment: Linux/aarch64, Python 3.12.13, Torch 2.10.0, CPU only. A one-rank Gloo run is not multi-rank FSDP sharding validation; this rerun does not claim the latter.

[Liuchenbing-2026]

@Yunnglin could you please review this focused single-process optimizer-resume fix? The same four regressions reproduce the missing-default-process-group failure on current main and pass on this head. Actual checkpoint I/O restores AdamW moments and the next update exactly, both with no process group and with a real one-rank Gloo group. The final published files were verified byte-for-byte against the rerun snapshot; this does not claim multi-rank sharding validation. Thanks!