perf: keep scalar optimizer step counters on CPU in state offload

[samsja]

Summary

optim_cpu_offload (on by default) moves optimizer state between CPU and GPU around each step. _move_states also shipped the 0-dim step counters to CUDA — and torch's AdamW reads each one back with .item() per parameter per step, turning every read into a device sync (587 syncs/step on Qwen3-30B-A3B, visible as aten::item ≈ 1.6 s in profiler traces at seq 8K on 8×H200).

This keeps scalar step counters on CPU — their canonical torch.optim location — so the optimizer reads them without touching the device.

Measured effect is modest (8K step median 3.18 → 3.13 s) because the syncs were mostly waiting on the state transfers rather than adding to them — the dominant cost of state offload remains the ~2×30 GiB/rank PCIe streaming itself. The change removes the per-parameter serialization points and keeps step counters in the standard representation for checkpoints.

Losses unchanged (verified on paired 5-step runs).

🤖 Generated with Claude Code

[samsja]

Heads-up: #3234 has since removed the state-offload mode entirely (offload is now full-offload-only and disabled by default), which deletes the _move_states code this PR patches. This fix is still correct for main as long as the current default state-offload path exists there — merge order decides whether it lands or becomes moot.

🤖 Generated with Claude Code