Preserve initialized optimizer state in low-level ZeRO

[sdjasj]

What changed

• migrate optimizer state that was eagerly initialized on working parameters to the corresponding local ZeRO master shards
• partition parameter-shaped state with the same padding and rank layout as the master parameter
• convert floating-point accumulators to the FP32 master dtype while preserving scalar state such as CPU step counters
• remove the stale working-parameter state entries
• add a distributed Adagrad regression using FP16 parameters and initial_scale=8.332635365271916

Why

Adagrad creates step and sum state in its constructor. Low-level ZeRO later replaced each working parameter in the optimizer groups with a master shard but left that preinitialized state keyed by the old parameter. The first optimizer step therefore found an empty state for the master shard and raised KeyError: 'sum'.

Impact

Eagerly initialized optimizers such as Adagrad retain correctly sharded state when wrapped by LowLevelZeroPlugin.

Fixes #6396.

Validation

• python -m pytest tests/test_zero/test_low_level/test_zero1_2.py::test_adagrad_initialized_state -q (2 GPUs; 1 passed)
• compileall and git diff --check: passed
• full file: the new regression passed; an existing BF16 DDP comparison failed before the optimizer step on one element (0.0078125 > 0.004)