Return frozen parameters as fp32 from get_fp32_state_dict_from_zero_checkpoint

[vineethsaivs]

get_fp32_state_dict_from_zero_checkpoint documents itself as returning "a single fp32 consolidated state_dict", and its own usage example feeds the result straight to model.load_state_dict. For a bf16 or fp16 run it returns a mixed-dtype dict instead: trainable parameters and buffers are fp32, every frozen parameter is not.

Reproduced against synthesized ZeRO-2 and ZeRO-3 checkpoints (world size 2, one trainable parameter, two frozen ones, one buffer):

stage 2: {'buf': float32, 'f_bf16': bfloat16, 'f_fp16': float16, 'trainable': float32}
stage 3: {'buf': float32, 'f_bf16': bfloat16, 'f_fp16': float16, 'trainable': float32}

Root cause

Trainable parameters come from the optimizer's fp32 master weights, so they are fp32 by construction. Buffers are upcast explicitly in parse_model_states:

buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}

A frozen parameter has no fp32 master copy: DeepSpeedEngine._get_param_fragment_func saves param.detach().cpu() (or param.ds_tensor.detach().cpu() under ZeRO-3), so the fragment carries the model's dtype. _zero2_merge_frozen_params and _zero3_merge_frozen_params then place that tensor in the state dict unchanged, and the eager return path calls to_torch_tensor(state_dict) with no dtype, so nothing casts it.

This is the usual shape for LoRA-style training, where the base model is frozen and only the adapter is trained, so most of the returned dict is the un-upcast half.

Fix

Pass dtype=torch.float32 in the eager branch, which is the boundary that carries the fp32 contract:

    if lazy_mode:
        return state_dict
    else:
        return to_torch_tensor(state_dict, dtype=torch.float32)

Lazy mode is deliberately left dtype-preserving. convert_zero_checkpoint_to_state_dict consumes it and applies its own --dtype, so upcasting there would both undo that option and force a full fp32 copy of the frozen weights during conversion, which is exactly the memory that lazy mode exists to avoid. Tensor.to returns the same object when the dtype already matches, so the trainable parameters and buffers are not copied.

Test

tests/unit/checkpoint/test_convert_checkpoint.py::test_fp32_state_dict_upcasts_frozen_params, a CPU test in the same monkeypatched style as the existing test_checkpoint_file_output_dtype. It asserts the eager path upcasts a bf16 frozen parameter and, in the same test, that lazy mode still hands it back as bf16 so the split above does not regress silently.

1 failed before the change, 5 passed after; the 4 existing CPU tests in the file pass on both sides.

The wider check was a round-trip harness over synthesized ZeRO-2 and ZeRO-3 checkpoints (world sizes 1 to 4, odd numels, several parameter groups, frozen parameters, buffers, tied parameters, sharded and unsharded output, eager and lazy) comparing the reconstruction against the weights the checkpoint was built from. Values and shapes already round-trip in every case; dtype was the only thing that did not match the documented contract.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 13805fc9f7

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

This is a non-merge commit, but its commit message has no Signed-off-by: line, so it violates the repository's commit requirements and may fail DCO/signoff validation before it can be merged; please amend/recreate the commit with --signoff using the configured git name and email.

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

Useful? React with 👍 / 👎.

[ebarkhordar]

I ran this against 13805fc in a clean container (torch 2.14.0+cpu). The fp32 contract argument holds for the export entry points, but this function has a third caller that is not an export path: deepspeed/runtime/engine.py:4561, inside DeepSpeedEngine._load_checkpoint.

if self.zero_optimization_partition_weights() and not load_optimizer_states and not self.has_moe_layers:
    checkpoint['module'] = get_fp32_state_dict_from_zero_checkpoint(load_dir)

One positional argument, so no lazy_mode and exclude_frozen_parameters stays False. A ZeRO-3 load with load_optimizer_states=False now materializes every frozen parameter as fp32 on every rank. Same call shape, one frozen bf16 base plus a small fp32 adapter:

master 80f19f30d   base.weight bfloat16    8.00 MiB   total  8.01 MiB
   PR  13805fc9f   base.weight float32    16.00 MiB   total 16.01 MiB

Those tensors are cast again when they land in the model, so this is extra CPU memory rather than a wrong result, and it costs most on frozen-base runs, which are exactly the ones carrying the most frozen bytes.

The docstring does promise fp32, so the engine may be the caller in the wrong here. The smaller edit is to upcast where fp32 is actually promised, which after this change is load_state_dict_from_zero_checkpoint at :787; convert_zero_checkpoint_to_fp32_state_dict already gets there by passing dtype down the lazy path, so it needs nothing. That leaves the engine load reading what the checkpoint held.

I only exercised the synthesized state dict this function is handed, not a real ZeRO-3 checkpoint, and I did not measure the engine path end to end.

[vineethsaivs]

You are right, fixed in 7360bb4.

get_fp32_state_dict_from_zero_checkpoint now takes dtype=torch.float32, and engine.py:4561 passes dtype=None, so the engine load reads what the checkpoint held, exactly as on master. The fp32 contract stays on the documented public path.

I kept the upcast as the default rather than moving it to load_state_dict_from_zero_checkpoint, because the docstring example calls this function directly and feeds the result to model.load_state_dict.

Test extended: test_fp32_state_dict_upcasts_frozen_params now also asserts dtype=None keeps the frozen bf16 and the trainable fp32. pytest tests/unit/checkpoint/test_convert_checkpoint.py: 5 passed, 1 skipped.

[ebarkhordar]

Read 7360bb4: engine.py:4563 passes dtype=None, and to_torch_tensor returns the contiguous tensor untouched when dtype is falsy, so the engine load keeps the checkpoint dtype exactly as on master. Keeping the fp32 default here rather than moving it makes sense to me, the docstring example calls this function directly. I read the diff and did not rerun the memory measurement.