Fix FlashOptim FSDP2 checkpoint serialization

[kashif]

Fixes two small FSDP2/Accelerate interop issues in FlashOptim when used through HF Trainer / TRL / Accelerate. Since v0.1.4 landed on main (which added the frozen-param continue in state_dict()), the pre-first-step case is already handled upstream — this PR now only carries the remaining gaps.

Summary

• Uneven-shard fallback during full-state-dict export. _wrap_state_as_dtensor previously raised ValueError on unevenly sharded parameters. In FSDP2 FULL_STATE_DICT export (accelerate/utils/fsdp_utils.py), the caller only needs a gathered state dict — not DCP-ready DTensors — so the raise aborted valid exports. Now the function silently returns, leaving tensors unwrapped (still a correct gathered export).
• Clearer error for CPU-offloaded shards. FlashOptim's quantized state and fused CUDA kernels require CUDA-resident shards at step time. Previously, a CPU-offloaded shard (e.g. fsdp_offload_params: true) crashed with an opaque kernel error during step(). Now _ensure_state_initialized raises a targeted ValueError pointing at fsdp_offload_params: false.
• Regression test for Accelerate prepare(). Pins that FlashOptimizer.state_dict() on a never-stepped optimizer returns {"state": {}, "param_groups": [...]} matching vanilla AdamW. Accelerate's AcceleratedOptimizer.__init__ calls optimizer.state_dict() during accelerator.prepare() (to move state to the correct device) — this test guards against regressing that contract.

Why

Single-GPU training worked, but FSDP2 checkpointing through Trainer/TRL/ Accelerate exposed two serialization assumptions that don't hold in distributed preparation: (a) unevenly sharded parameters can make DTensor wrapping impossible even when a valid full-state export is still possible, and (b) CPU-offloaded shards are not a meaningful FlashOptim mode and previously failed with a low-level quantization error instead of a clear incompatibility message.

Caveats

• Uneven-shard handling now favors successful full-state-compatible export over strict DTensor wrapping for that case.
• CPU-offloaded parameter shards remain unsupported; this PR only improves the error message.

[josejg]

Thanks for the contribution @kashif

can you elaborate what

• make FlashOptimizer.state_dict() behave like a standard PyTorch optimizer before the first step, since accelerate may inspect optimizer state during prepare()

means?

[kashif]

Thanks for the contribution @kashif

can you elaborate what

• make FlashOptimizer.state_dict() behave like a standard PyTorch optimizer before the first step, since accelerate may inspect optimizer state during prepare()

means?

after rebasing onto v0.1.4, this fix doesnt apply anymore. v0.1.4's

if param_number not in opt_state:
    continue

in state_dict() already handles the pre-first-step case (same condition as frozen params: no state allocated → skip).

The original motivation for calling it out was that accelerate.AcceleratedOptimizer.__init__ calls optimizer.state_dict() during accelerator.prepare() to move state to the target device when device_placement=True — if FlashOptim had raised or returned garbage on empty state, Trainer + Accelerate would break before step 1.

So now, I've dropped the redundant code and kept only a regression test (test_state_dict_before_first_step_matches_torch_empty_state) that pins the contract, so future refactors can't regress it. I've also updated the PR description to reflect the narrower scope.