Keep ZenFlow Adam step counters in float32

[vineethsaivs]

ZenFlow's Adam counter stops at 256 for BF16 and 2048 for FP16, freezing bias correction while training continues.

Cause: Five initialization paths store the counter in the parameter dtype. Old checkpoints retain that dtype when loaded.
Fix: Initialize counters as float32 and promote restored FP16/BF16 counters in both selective optimizers, preserving existing float64 counters and the saved value.
Test: DS_ACCELERATOR=cpu python -m pytest tests/unit/runtime/zenflow/test_zf.py -k "counter or test_num_selected_columns or test_split_affinity": new training, checkpoint restoration, and resumed updates past both limits. All 21 selected ZenFlow tests pass on Apple M2 Pro CPU, torch 2.11.0; pre-commit passes. Distributed offload was not exercised. Previously lost step counts cannot be reconstructed.

[ebarkhordar]

I read this at 4e23970 and went looking for where a bf16 step tensor comes from inside DeepSpeed, to see which half of the patch each case needs.

The creation half lands where it should: zenflow_stage_1_and_2.py:81 builds the selective optimizer over self.bit16_groups, so param.dtype is bf16 or fp16 for every real ZenFlow run, not just a hand-built one.

The __setstate__ half is what I could not place. Which path hands these classes a checkpoint whose selective step is bf16? ZenFlowZeroOptimizer defines no state_dict and no load_state_dict of its own, so it inherits DeepSpeedZeroOptimizer.state_dict, which serializes self.optimizer. selective_optimizer is a separate instance attached beside it (engine_stage3.py:120 does the same for stage 3), and selective_optimizer.state_dict( and selective_optimizer.load_state_dict( occur zero times in the package.

If that reading is right, the selective step, exp_avg and exp_avg_sq never enter a DeepSpeed checkpoint, so __setstate__ can only fire for someone driving ZenFlowSelectiveAdamW through torch APIs directly, and a ZenFlow run resumed from a checkpoint restarts its selective moments at zero whatever the dtype was. That would make the second half a different bug from the first one rather than the resume story for it.

I did not run this. The above is from reading the class bodies at that sha and searching the package for the two calls, so it is a question and not a claim about behaviour.