Keep the parameters a model names for fp32 in fp32 under ZeRO-3

[pengdurice]

Keep the parameters a model names for fp32 in fp32 under ZeRO-3

Follows: #8699 (merged), which kept the listed buffers in fp32 and left the listed parameters to this PR.

Summary

Hugging Face transformers models name the tensors that must stay in fp32 (_keep_in_fp32_modules_strict, and
_keep_in_fp32_modules for fp16). #8699 added data_types.keep_in_fp32_modules (default "auto": the model's own
lists) and kept the listed buffers in fp32 under ZeRO-3. Listed parameters were still trained in bf16, because
ZeRO-3 flattens each sub-group into one buffer of one dtype. Examples of listed parameters: DeepSeek-V4's
hyper-connection weights, attention sinks and norms; gpt-oss's attention sinks.

This PR lets ZeRO-3 (without parameter offload) hold partitions of two dtypes, so listed parameters are stored,
gathered, updated and checkpointed in fp32 next to the bf16 or fp16 ones.

Duplicate check

Searched open and recent pull requests on 2026-10-15 ("keep_in_fp32", "fp32 modules", "mixed dtype"): only #8699
(merged, buffers). #8686 (reconstruction of shared-submodule parameters in ZeRO checkpoints) is unrelated.

The change

File	What
deepspeed/runtime/keep_in_fp32.py	keep_tensors_in_fp32(module, pattern, include_params=True) converts matching parameters too; a ZeRO-3 parameter is converted in its partition (ds_tensor), so later gathers return fp32 and a later checkpoint load writes exact fp32 values
deepspeed/runtime/zero/partition_parameters.py	after a transformers model's own __init__ (when its fp32 list is complete), its already-partitioned parameters are converted in place, before any checkpoint is loaded; with parameter offload only buffers are kept
deepspeed/runtime/zero/stage3.py	the training dtype is the first parameter that is not fp32; parameters kept in fp32 form sub-groups of their own (a sub-group is one flat buffer of one dtype), get their own flat gradient buffer (fp32 accumulation) and reduce their gradients in fp32; gradient buckets are sized per dtype
deepspeed/runtime/engine.py	ZeRO-1/2 and no ZeRO: the engine's mixed-precision cast skips nothing new (those optimizers hold one dtype), keeps buffers, and warns that listed parameters are trained in the training dtype; an explicit parameter pattern below ZeRO-3 is refused
docs/_pages/config-json.md	data_types.keep_in_fp32_modules now covers parameters (ZeRO-3) and buffers
tests/unit/v1/half_precision/test_keep_in_fp32_modules.py	tests below

Measured

DeepSeek-V4-Flash, first 4 layers, ZeRO-3, bf16, 1 node: 54 listed parameters and 1 buffer are
fp32 and exact after loading (bf16 before); the run without a list is bitwise equal to master over 20 steps; with the
list the mean loss differs from master by -0.0063 (lower on 15 of 20 steps). Throughput and memory were not
measured separately; the listed parameters are norms, sinks and small mixing weights, a small share of each model.

Tests

Test	Checks
TestZeroInitKeepsParamsFp32::test_loaded_values_survive	under zero.Init, listed parameters keep exact fp32 checkpoint values after loading
TestZero3TrainsKeptParamsInFp32::test_update_below_bf16_resolution_is_kept (zero.Init on/off, contiguous gradients on/off, optimizer offload on/off, gradient accumulation 1/2)	an update smaller than bf16's resolution survives in the fp32 parameter
TestZero3TrainsKeptParamsInFp32::test_all_reduce_fetch	the fp32 and bf16 partitions are gathered with one all-reduce per dtype
TestZero3TrainsKeptParamsInFp32::test_checkpoint_round_trip	save and load keep the fp32 parameters exact
TestZero3TrainsKeptParamsInFp32::test_param_offload_keeps_only_buffers	with parameter offload only buffers are kept (one flat buffer of one dtype)
TestZero3Fp16KeptParams::test_fp16_trains	fp16 with a loss scaler trains with fp32 listed parameters
TestStagesBelowThreeKeptParams	ZeRO-0/1/2 cast listed parameters to the training dtype and keep buffers; an explicit parameter pattern is refused

Hardware: 1 node of 8 H200 (driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1.

pytest unit/v1/half_precision/test_keep_in_fp32_modules.py     # master 1c862f7df + this change: 47 passed
                                                               # master 5f266e26d + this change: 47 passed, exit 0

pre-commit run --files passes on every changed file. Regression runs at the time of the first version: identical pass counts with and without the change on test_zero, test_offload_states and
tests/unit/runtime/zero.

Compatibility

Cell	Validated
ZeRO-3 bf16 and fp16, zero.Init on and off, optimizer offload on and off, gradient accumulation	yes (tests)
ZeRO-3 with parameter offload	buffers only, by design (tested)
ZeRO-0/1/2	listed parameters in the training dtype with a warning (tested)
Universal checkpoint conversion of fp32 listed parameters	not run
Accelerators other than CUDA	not run

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).