fix: metric deferral + optimizer optimization

[Dogacel]

Speedup is near 5-10%, because saved computation is already mostly overlapped. However those optimizations ensure we have more room to overlap in the future.

1. Defer metrics to be reported one step ahead -> ensures main loop is not blocked by .item() calls produced by metrics.
2. Replace python for in optimizer with torch._foreach_copy_ -> %50 faster optimizer steps
3. LRU cache EAGLE-3 block mask (~%5 speedup)

Final profiler result:

Before:

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b9f890d801

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Avoid forcing fused AdamW for CPU-offloaded training

When fsdp_cpu_offload is enabled, apply_fsdp2(..., cpu_offload=True) installs CPUOffloadPolicy, so this optimizer's fp32 master parameters can be CPU tensors. Forcing fused=True makes AdamW require fused-supported accelerator tensors and raises on the first optimizer step instead of falling back to the single/foreach implementation, regressing CPU-offloaded training from working to failing; let PyTorch choose the backend or gate this on all master params being on a fused-supported device.

Useful? React with 👍 / 👎.

[Dogacel]

Avoid forcing fused AdamW for CPU-offloaded training

When fsdp_cpu_offload is enabled, apply_fsdp2(..., cpu_offload=True) installs CPUOffloadPolicy, so this optimizer's fp32 master parameters can be CPU tensors. Forcing fused=True makes AdamW require fused-supported accelerator tensors and raises on the first optimizer step instead of falling back to the single/foreach implementation, regressing CPU-offloaded training from working to failing; let PyTorch choose the backend or gate this on all master params being on a fused-supported device.

Useful? React with 👍 / 👎.

Should I disable if fsdp_cpu_offload then?