[Levius-Fubuki]

Fixes #8910.

FP16/BF16 ZeRO-0 wrappers currently call optimizer.step() with zero gradients during construction, advancing Adam counters and applying AdamW decay before training. Remove the synthetic updates from both wrappers and allocate lazy state on the first real update. Move torch Adagrad's eagerly created state onto the corresponding FP32 masters, preserving the supplied optimizer and group settings. Existing checkpoint counters are loaded unchanged.

Device validation: 2 x NVIDIA GeForce RTX 3090 24 GB available, new cases use one GPU; driver 580.105.08, torch 2.12.1+cu130, CUDA 13.0, NCCL 2.29.7, Python 3.12.3. Baseline bf222a2 fails the new counter/initialization regression. With the fix:

python -m pytest --forked tests/unit/runtime/half_precision/test_fp16.py::TestOptimizerInitialization --torch_ver=2.12.1+cu130 --cuda_ver=13.0 -q
9 passed
python -m pytest --forked tests/unit/runtime/half_precision/test_dynamic_loss_scale.py tests/unit/runtime/half_precision/test_fp16.py::TestAdamwFP16Basic tests/unit/runtime/half_precision/test_fp16.py::TestLambFP16 --torch_ver=2.12.1+cu130 --cuda_ver=13.0 -q
9 passed

New cases cover FP16/BF16 Adam/AdamW/FusedAdam with decay 0/0.1, plain FP32 update references, actual engine checkpoints before the first update and after 20 updates, continued training, two-group eager Adagrad, and both FusedLamb gradient APIs (counter/restore checks for LAMB). Tests explicitly disable clipping to isolate optimizer updates. Changed-file pre-commit, syntax and diff checks pass; independent spec/code reviews found no remaining issue.

A direct-wrapper CUDA measurement with 1,048,576 FP16 parameters, torch AdamW(lr=0.01, weight_decay=0.1), 3 repeats and 40 updates reproduces master 0.499500006 / step 1 after baseline construction versus 0.5 / step 0 with the fix. Median construction peak allocated memory is 24 MiB versus 8 MiB; overall training peak is 36 MiB for both. Allocation shifts to the first actual step. The short quadratic-loss trajectory falls 0.25→0.002882 baseline and 0.25→0.002622 fixed; these are different update semantics, not a convergence-quality comparison. Timing is noisy at this scale (recorded constructor/first/steady medians in the experiment artifacts); no throughput improvement or production-convergence claim is made.

Duplicate check found no open PR covering #8910. #8673 modifies mixed-precision FusedAdam updates but retains this constructor path. Dedicated BF16 and ZeRO 1/2/3 initialization were inspected and already avoid this synthetic step.

Unvalidated: Apex legacy Adam, other accelerators, offload, pipeline/tensor/MoE parallelism, cross-stage/world-size checkpoint conversion, generic prepopulated custom optimizer state. Adagrad frozen/zero-element state entries are a preexisting separate corner outside these trainable-parameter cases.

AI assistance (OpenAI Codex) was used for investigation, implementation and device tests. The human submitter confirmed manual diff review and test-result verification and authorized submission. Device tests above were executed on the requested server with Codex; no additional human test execution is asserted.

Review follow-up (1a372ef): generalized Adagrad state promotion beyond sum and added an engine regression for additional eager tensor state, unchanged counters/metadata, and a plain-torch update reference. On Apple M3/macOS ARM64, Python 3.12.13, torch 2.14.1, CPU/Gloo: TestAdagradLowPrecisionState reports 1 passed, 1 skipped (FP16 unsupported). The BF16 regression fails before this follow-up and passes after it. Changed-file pre-commit, syntax and diff checks pass. The GPU results above were obtained on the original PR revision 7d2a5573; CUDA tests were not rerun for this follow-up because the server is offline.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	🔄 Running since 2026-10-27T01:32:15.907457Z	7d2a557	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[sfc-gh-truwase]

Treating sum specially makes this fragile to future Adagrad changes. Consider how to generalize upcasting all appropriate tensors in state.

[Levius-Fubuki]

Addressed in 1a372ef.

The migration now iterates over the state entries and promotes tensors whose dtype matches the original low-precision parameter to the FP32 master's dtype/device, instead of naming sum. Higher-precision step counters retain their dtype/device, and boolean tensors and non-tensor metadata are preserved.

Added an engine regression covering an additional eager accumulator, counter/metadata preservation, and a parameter update against plain torch Adagrad. The BF16 case fails before this change because the additional accumulator remains BF16, and passes after it.

Local validation on Apple M3, macOS ARM64, Python 3.12.13, torch 2.14.1, CPU/Gloo:

DS_ACCELERATOR=cpu CUDA_VISIBLE_DEVICES=0 LOCAL_SIZE=1 GLOO_SOCKET_IFNAME=lo0 DS_UNITTEST_TIMEOUT=90 python -m pytest --forked tests/unit/v1/half_precision/test_bf16.py::TestAdagradLowPrecisionState --torch_ver=2.14.1 -q

Result: 1 passed, 1 skipped (FP16 is unsupported by this CPU accelerator). Changed-file pre-commit, syntax checks and git diff --check pass. CUDA validation was not rerun for this follow-up because the GPU server is offline; the earlier GPU results apply to the previous PR revision.