Add opt-in TorchAO HiFloat8 Linear and AutoEP training

[hazelduan]

Selected Linear and AutoEP GroupedExperts train with HiFloat8 through one public TorchAO quantize_ call. AO now owns both configuration types and both registered handlers; DeepSpeed selects modules and manages training lifecycle without defining a grouped quantization handler or choosing a GEMM callable.

RFC #8777 documents the architecture, quick start and historical compute/end-to-end comparisons.

Implementation:

• HiFloat8GroupedConfig and the existing Linear recipe derive from AO's shared HiFloat8Config. Grouped uses the inherited 15/15/224 defaults and from_dict validation. Legacy HiFloat8LinearConfig native-range defaults remain compatible. No nested policy wrapper is added.
• Engine parses the same numerical JSON into the applicable immutable recipes, checks native/lifecycle capabilities before any conversion, builds one ModuleFqnToConfig map and calls quantize_ after AutoEP/device preparation and before optimizer creation.
• AO's grouped handler accepts a module exposing the existing grouped_mm slot and binds the unchanged HiFloat8 grouped functional/autograd implementation. It imports no DeepSpeed container. The former deepspeed/runtime/torchao.py adapter is removed.
• The accelerator lazily resolves public torchao.hifloat8 configuration/probe APIs. Dense-only selection does not require grouped exports; selected experts require HiFloat8GroupedConfig before conversion. There is no DeepSpeed reference to hifloat8_grouped_mm for execution.
• Native BF16/FP16 grouped GEMM remains in deepspeed/ops/npu, independent of AO. DeepSpeed retains routing, padding, activation, communication, parameter/state identity and caller-facing dtype.
• Qwen3.5 text-wrapper/checkpoint fixes and NPU permutation generation remain in this PR as separate functional commits.

Merge this fragment into the existing training configuration and use the normal deepspeed.initialize(..., config=ds_config) call:

{
  "bf16": {"enabled": true},
  "zero_optimization": {"stage": 2},
  "hifloat8": {
    "enabled": true,
    "backend": "torchao",
    "module_name_patterns": ["*.experts", "*.shared_experts.*_proj"],
    "min_numel": 65536
  }
}

Install the paired sources with torchao>=0.18,<0.19, matching Torch/TorchNPU and CANN, and automatic activation enabled. No explicit import torchao_npu is needed. AutoEP topology, optimizer/batch settings and trainable-parameter selection remain separate. See the tutorial.

October 19 validation on Ascend950PR, original Conda Python3.11.15, Torch/TorchNPU2.11.0, TorchAO0.18.0, CANN9.1.0 and driver25.7.rc1:

• AO HiFloat8/namespace suites: 71 passed, 1 empty-Dense native-contract skip. DeepSpeed focused suites: 342 passed, 32 CUDA-only skips, including CPU reference cases.
• A fresh process imports torch and public torchao APIs, converts a Linear and a plain grouped module together, and runs native forward/backward. There is no explicit torchao_npu import and no DeepSpeed module loaded. The grouped configuration derives from HiFloat8Config.
• Actual AutoEP EP2/ZeRO-2: paired 30-step runs plus a reverse-order 100-step pair in each precision, BF16 and HiFloat8 separately. Inputs, losses, gradient norms and all 19 final tensors per rank are exact against the previous paired DS/AO sources. Two layers/four experts, hidden64/intermediate128, batch1/sequence16, seed42 and rank-dependent data seed123+rank, AdamW1e-3, same CPU affinity. This is a small actual AutoEP model, not a full Qwen convergence run.
• NPU allocation peaks are identical before/after: 17,627,648 bytes initialization; 1,035,863,040 bytes training. Initial 30-step P50 changes were BF16 +14.57%, HiFloat8 +1.05%; the reverse-order 100-step changes were BF16 −0.50%, HiFloat8 +0.45%. Ordinary BF16 execution code is unchanged. The samples show timing sensitivity and do not establish a stable speedup or regression.
• DP2/ZeRO-2/GAS2 checkpoint/resume retains exact subsequent losses, parameters and optimizer state. Typed Linear/grouped recipes have equal numerical fields; grouped modules reuse the same grouped recipe. Experts are replicated, so this is not AutoEP checkpoint proof.
• Local CPU focused tests: 40 passed, 82 NPU-dependent skips. The first three reconstructed DeepSpeed commits were also tested independently: 199/6, 6/24 and 10/37 passed/skipped. Changed-file checks cover the full feature diff. Final trees match the validated source, with only a separately checked documentation wording correction after the initial native suites.

The previous published source pair was DS 6b91ff09 / AO 9c6dfb4. Reproduction: /data/dxq/hif8_a5_20260914/experiments/ao_grouped_handler_20261019_01 (run_all.sh, postchecks.sh, source manifests, exact configs/rank tables, per-rank results and comparisons). Original packages and host permissions were preserved; device access uses a task-local non-root container identity and cache paths.

The existing feature histories are consolidated by responsibility, preserving the final trees. No new PR was created.

DeepSpeed:

• 98caec79 — fix(autoep): preserve Qwen3.5 routing and rank-local checkpoints
• 9a87125d — feat(npu): expose native grouped GEMM through accelerator capabilities
• b49b298f — perf(autoep): generate large permutation indices on NPU
• 95dd780c — feat(hifloat8): integrate AO-owned quantize_ training handlers

AO:

• 6b034fd — fix: preserve NPU auto-activation during TorchAO discovery
• ea917fd — feat(ops): add differentiable native NPU grouped GEMM
• a388ef3 — feat(hifloat8): register Linear and grouped training workflows
• e0d59bd — bench(hifloat8): separate timing, memory and kernel attribution

Current sources: DeepSpeed 95dd780c, AO e0d59bd (MR31), Swift 3e9994baf. Swift changes are documentation only. Draft boundaries remain BF16 parameters/ZeRO0–2 on NPU; reject ZeRO3, pipeline, AutoTP/custom MPU, explicit MoE and fused gate/up overrides. Full-parameter convergence, LoRA, AutoEP checkpoint recovery, quantized checkpoint ingestion, offload/compile, EP8, multi-node and CUDA low precision remain outside the established support.

[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	✅ Completed 2026-10-09T10:18:37.931586Z	2623227	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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2623227bb9

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

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

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Add the required Signed-off-by trailer

This is a one-parent, non-merge commit, and git log --format='%B' 7e5c4fd3085d9076a170d2540f2f3e3de5d7e8f9 confirms that its body has no Signed-off-by trailer, so it does not meet the repository's DCO requirement; sign off the commit before merging.

AGENTS.md reference: AGENTS.md:L8-L9

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Add the required license header to the tutorial

This newly added tutorial starts directly with Jekyll front matter and contains neither the required SPDX line nor the DeepSpeed Team line; add the two comments inside the YAML front matter so the page remains valid while satisfying the new-file license requirement.

AGENTS.md reference: AGENTS.md:L13-L17

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Restrict FP16 outputs in BF16-only HiFloat8 training

When hifloat8.config.compute_dtype is set to the documented "float16" value, this accepts the policy even though HiFloat8 initialization requires a BF16 model. The tutorial states that GEMM results use the selected working dtype, and _run_experts_npu returns that result unchanged, so selected Dense/GroupedExperts modules emit FP16 into BF16 residual paths; FP16 plus BF16 promotes to FP32, after which the next ordinary BF16 linear fails on mismatched matrix dtypes. Reject FP16 here or cast converted module outputs back to their input dtype.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Guard the Qwen integration against older Transformers

In the inspected cpu-torch-latest manual PyPI path, Transformers defaults to 4.51.3 and the workflow runs the full unit suite, but that supported release predates all four Qwen3_5Moe* symbols imported here. Because the import is unconditional, this test errors rather than skipping the feature whose preset is itself version-gated; use pytest.importorskip plus the existing skip_unless_transformers_has pattern before constructing these models.

Useful? React with 👍 / 👎.