FusedAdam step counter desynchronizes across DP ranks with empty param_groups

[redmoe-moutain]

Describe the bug

In FusedAdam, the step counter (group["step"]) for a parameter group becomes desynchronized across Data Parallel (DP) ranks if that group is empty on some ranks. The optimizer logic skips the step increment for any group where len(group["params"]) == 0.

This creates a critical issue when resuming from a checkpoint. A rank that holds parameters for a specific group (e.g., RMSNorm params in a no_weight_decay group) may load its optimizer state from another DP rank where that group was empty. This causes it to load an incorrect, stale step value. An incorrect step leads to a wrong bias_correction term (1 - beta1 ** step), causing training instability or divergence upon resumption.

The image below illustrates this problem in a PP=2, EP=4, DP=8 setup after 2640 steps. The step counter for the no weight decay group is 1 on several ranks (orange) because they contain no such parameters, while it is correctly 2641 on others.

Within the same data parallel group (pp=0, ep=2), the rank dp=6 (global rank 6) has the correct step 2641. However, rank dp=2 (global rank 2) has an incorrect step of null. When resuming, rank 6 might load the state from rank 2, inheriting the wrong step and breaking the bias_correction calculation.

Suggested Fix:
The step increment logic should be moved before the check for an empty parameter list to ensure it always executes.

# In FusedAdam's step() method
for group in self.param_groups:
    # Increment step for all groups on all ranks first
    if "step" in group:
        group["step"] += 1
    else:
        group["step"] = 1
    
    # Then, skip the kernel launch for empty groups
    if len(group["params"]) == 0:
        continue
    
    # ... rest of the logic