Fix subgroup optimizer metadata inconsistency

[st-bang97]

Description

This PR addresses Issue #7819.

When using ZeRO Stage-3 + CPU offloaded optimizer (CPUAdam) with 2+ subgroups, ds_adam_step can be invoked multiple times within a single global optimizer step. Before this fix, the internal bias-correction-related state (e.g., _betta2_t → _bias_correction2) could become inconsistent across subgroup invocations within the same step, leading to subgroup-wise optimizer state divergence.

---

Solution

Make IncrementStep() step-consistent under repeated calls in the same global step.

Change

Update IncrementStep() to only advance or recompute state when step != _step, preventing subgroup-to-subgroup drift inside a single step.

inline void IncrementStep(size_t step, float beta1, float beta2)
{
    if (beta1 != _betta1 || beta2 != _betta2) {
        _step = step;
        _betta1 = beta1;
        _betta2 = beta2;
        _betta1_t = std::pow(_betta1, step);
        _betta2_t = std::pow(_betta2, step);
    } else {
        if (step != _step) {
            _step++;
            if (_step != step) {
                _betta1_t = std::pow(_betta1, step);
                _betta2_t = std::pow(_betta2, step);
                _step = step;
            } else {
                _betta1_t *= _betta1;
                _betta2_t *= _betta2;
            }
        }
    }
}

1. Reproduction (Before Fix)

2. Verification (After Fix)

[tohtana]

Hi @st-bang97,
Thank you for the contribution!
Sorry, I just noticed you've already addressed #7819. Please disregard my comment.

Overall this PR looks good to me. Here are some suggestions:

• Apply the same fix to the header for XPU
• Add some tests to avoid addition (you can reuse some from Fix CPUAdam same-step subgroup drift in ZeRO-3 (#7819) #7859 if you find them useful)
• Fix format