[adenzhou1350]

Current-head maintenance — October 30

Merged master in 9f7da5cbfe4eb5aa7b65324bd1add3fce25682bd. Resolved the import conflict against master 8d4f29c75f4a: retain CUDAQuantizer while adopting upstream ZeRORuntimeException and get_norm_dtype. The rest of master merged normally; the original quantized all-gather test file is byte-identical to the previous head.

• All seven PR-changed files pass the repository's complete pre-commit run --files ... hooks on isolated Linux / Python 3.12, including clang-format, YAPF and flake8. No source changes were made by those hooks.
• All five authored Python files parse successfully; the tested archived Git tree matches this pushed merge commit.
• Native Windows Python 3.14 could not run the old flake8/pyflakes combination (ast.Str was removed) and lacked clang-format; the same required hooks were therefore run on Linux/Python 3.12 without changing hook versions or bypassing hooks.
• This is an import-conflict maintenance refresh. No new GPU/distributed training, convergence or performance run was performed at this head. The earlier device measurements below remain historical.
• Signed-off normal fork push and clean worktree verified. Hosted CI for the new head awaits its normal upstream authorization.

---

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.

Current-main maintenance — October 1

Merged current master (3509874) and kept both the quantized helpers and the new optimized-parameter import. CPU-focused tests: 22 pass, with 6 distributed GPU cases not rerun. All configured changed-file pre-commit checks pass on Linux/Python 3.12. Earlier GPU correctness/performance results remain historical; this refresh does not claim a fresh GPU qualification.

[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.

[adenzhou1350]

Merged current master (3509874) and kept both the quantized helpers and the new optimized-parameter import. CPU-focused tests: 22 pass, with 6 distributed GPU cases not rerun. All configured changed-file pre-commit checks pass on Linux/Python 3.12. Earlier GPU correctness/performance results remain historical; this refresh does not claim a fresh GPU qualification.

[adenzhou1350]

Merged current master in 65a03d91 and resolved the import conflict, retaining both sets of helpers. Pre-commit passes; the focused suite passes 22 CPU cases and all 6 two-GPU RTX 5090 cases (ZeRO-1/2, FP16/BF16 training and reconstruction) on PyTorch 2.11.0/CUDA 13.0. This is a correctness refresh; the earlier performance measurements were not rerun.

[adenzhou1350]

Refreshed in 9f7da5cbfe4e. Conflict resolved by retaining CUDAQuantizer and adopting upstream ZeRORuntimeException/get_norm_dtype. The complete pre-commit checks pass on all seven PR files in isolated Linux/Python 3.12; syntax checks pass and the original quantized test is byte-identical. No new current-head GPU/distributed or performance qualification is claimed; earlier device results remain historical.

This refresh changes no measured performance claim; current hosted CI is separate from earlier-head evidence.