ZeRO-1/2: add bounded quantized weight all-gather

[adenzhou1350]

Summary

Extend the existing zero_quantized_weights opt-in to ZeRO stages 1 and 2. After each optimizer step, weight partitions are quantized in bounded chunks, gathered, and dequantized one source rank at a time into the persistent flat buffer.

Large weights use 64-element symmetric INT8 groups. Parameters with at most 65,536 elements use a bounded original-dtype side channel; they are masked before quantization and restored after gathering. The bucket limit applies across all data-parallel ranks, with a minimum of one quantization group per rank. Unsupported FP32 and ZenFlow combinations are rejected during optimizer initialization.

The patch also removes an unused cuda_profiler_api.h include from gemm_test.h; CUDA 13 no longer ships that header and the quantizer does not use the profiler API. The default ZeRO-1/2 path is unchanged.

Validation (f2c781d)

• Changed-file pre-commit checks passed.
• Full focused test file: 28 passed, 0 failed/skipped — 22 CPU cases (including simulated DP64 buffer bounds) and 6 real two-GPU RTX 5090 cases. The GPU cases cover ZeRO-1/2 FP16/BF16 training and public gather reconstruction across chunk/partition boundaries. The quantizer was freshly built from this revision.
• Qwen3-4B-Base, sequence 2048, BF16, ZeRO-2 DP4 on four PCIe RTX 5090 GPUs: four fresh-process arms in baseline/candidate/candidate/baseline order, each with 1 warmup + 19 measured steps.
• Within-arm loss/gradient trajectories and sampled parameter hashes reproduce exactly across orders. Cross-arm maximum loss difference: 0.004051; 1,194 parameter samples: relative L2 0.000460, max absolute difference 0.0078125, zero atol=0.02, rtol=0.02 violations.
• Whole training-step speedup: 1.08515x; within-run 95% bootstrap CI [1.08418, 1.08590]. Baseline medians: 1.19319 / 1.19339 s; candidate: 1.09943 / 1.09989 s.
• Peak allocated/reserved memory: 28.88 / 30.25 GB in both arms, within the unchanged 90% / 95% gates.

The DP4 run used PyTorch 2.11.0+cu130, CUDA 13.0, NCCL CTA8, a 26M reduce bucket and a 50M all-gather bucket. Both arms used the same exact source; only zero_quantized_weights differed.

Scope and risks

This is lossy, opt-in synchronization. The short BF16 Qwen run is not a convergence study; its within-run interval does not establish cross-day or cross-machine robustness. FP16 has focused DP2 correctness coverage, not whole-model performance qualification. MoE, offload, other models/accelerators, NVLink/RDMA, and long-horizon convergence are not established. Small-parameter-heavy workloads may not benefit from the extra side-channel collectives.

Prepared and tested with Codex assistance.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bb2383dadf

ℹ️ 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]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c0debf7df7

ℹ️ 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".

[adenzhou1350]

Addressed all four findings in f2c781d: early FP32/ZenFlow rejection, aggregate all-ranks buffer limits, and public reconstruction tests instead of private-layout assertions.

Fresh-head validation: 28/28 focused tests pass, including six real DP2 cases; the same-condition, order-balanced DP4 Qwen run passes correctness/memory gates at 1.085x speedup. The PR body now has the updated results and their scope.