Fix NaN for full-parameter SFT with DeepSpeed ZeRO-1/2 on torch<2.9 (AdamW int32 flat-partition overflow)

[kiscad]

Problem

On torch<2.9, torch._foreach_* kernels index tensors with signed int32 elements and write out of bounds when a flat tensor exceeds INT32_MAX (2,147,483,647) elements. DeepSpeed ZeRO-1/2 flattens each optimizer param group into one FP32 flat partition per DP rank, so full-parameter training of large models crosses this boundary and the AdamW step corrupts memory — loss turns NaN right after the first optimizer step.

Reproduced with Qwen3.5-9B full SFT (--optim adamw_torch --deepspeed zero1, bf16, 4×A100):

• an independent torch._foreach_add_ probe on a 2,238,382,272-element tensor writes 65,344 elements past the end (guard corrupted, 3/3 repeats); AdamW(foreach=False) and a chunked reference are clean;
• training loss is 0.24 at step 1 and NaN from step 2; the same command on 8 GPUs (partition below the boundary) stays finite.

Fix

In SwiftMixin.create_optimizer, after the HF Trainer creates the optimizer and before DeepSpeed receives it: on torch<2.9 + ZeRO stage 1/2 + torch.optim.AdamW, set foreach=False on the optimizer defaults and every param group so the update takes the single-tensor path. ZeRO-3 (per-parameter partitioning, no oversized flat), torch≥2.9, and non-AdamW optimizers are untouched.

The fallback costs ~2% on a 2 GiB flat (ZeRO feeds the base optimizer a few large flat tensors, where foreach has little advantage), and it does not change the optimizer's group layout or checkpoint format.

Verification

Run	Loss
4×A100, fix disabled	0.24 → NaN (step 2+)
4×A100, fix enabled	0.24 → 0.15 → 0.11 → 1.20 → 0.32, no NaN
8×A100, fix enabled (no-op)	all finite, no NaN

[kiscad]

Hi @tastelikefeet , when you have a chance, could you please take a look at this PR? Thanks!

[tastelikefeet]

Thanks for your code contribution! I will read the PR asap, thanks again!

[tastelikefeet]

Hi, how about make a patch of:

optimizer.defaults['foreach'] = False
    for group in optimizer.param_groups:
        group['foreach'] = False

[kiscad]

Hi, how about make a patch of:

optimizer.defaults['foreach'] = False
    for group in optimizer.param_groups:
        group['foreach'] = False

Thanks for the suggestion — agreed, this is much simpler and avoids changing the optimizer’s group layout or checkpoint format. Done.

Re-verified end-to-end on 4× A100s: the loss stays finite with the fix, but goes NaN from step 2 without it.