Fix zero_to_fp32 --safe_serialization on models with tied weights

[vineethsaivs]

zero_to_fp32 --safe_serialization aborts on any checkpoint whose model ties weights, which is the usual embedding / lm_head pair, so most decoder-only models are affected:

RuntimeError:
    Some tensors share memory, this will lead to duplicate memory on disk and
    potential differences when loading them again: [{'lm_head.weight', 'model.embed_tokens.weight'}].

Root cause

_get_fp32_state_dict_from_zero{2,3}_checkpoint recovers a tied parameter by pointing it at the tensor it is tied to:

for pair in zero_model_states[0].shared_params:
    if pair[1] in state_dict:
        state_dict[pair[0]] = state_dict[pair[1]]

to_torch_tensor then deliberately keeps that aliasing so the pair costs one tensor rather than two:

if tensor_id in converted_tensors:  # shared tensors
    shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
    torch_state_dict[name] = shared_tensor

That is exactly what safetensors.torch.save_file refuses. The pickle path is fine because torch.save stores shared storage once and restores the aliasing on load.

Note that ordinary (untied) parameters are not affected even though they are all views into one flat vector: safetensors only rejects tensors that overlap, and sibling narrow() slices are disjoint. Only the tied pair, which is the same region twice, trips it.

Fix

to_torch_tensor grows a share_tensors flag, and convert_zero_checkpoint_to_state_dict passes share_tensors=not safe_serialization at both call sites, so the tied entry gets its own copy on the safetensors path and nowhere else. Passing it to the return_empty_tensor=True call as well means split_torch_state_dict_into_shards sees two distinct storages and sizes the shards for what is actually written.

Cloning rather than dropping the duplicate key keeps the documented contract that the output loads with load_state_dict(): a tied module still lists both names in its state_dict(), so a strict load needs both. The pickle path is untouched and still shares the tensor, which test_output_dtype_conversion_preserves_shared_tensors and test_checkpoint_file_output_dtype already assert.

Test

tests/unit/checkpoint/test_convert_checkpoint.py::test_safe_serialization_with_tied_weights, parametrized over max_shard_size None and "1GB" so both the single-file and the sharded writer are covered. It is a CPU test in the same monkeypatched style as the existing test_checkpoint_file_output_dtype.

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

Separately I exercised the change against synthesized ZeRO-2 and ZeRO-3 checkpoints on disk (world sizes 1 to 4, odd numels, several parameter groups, frozen parameters, buffers, tied parameters, fp16/bf16 output, sharded and unsharded, safetensors and pickle), reconstructing the state dict and comparing it to the weights the checkpoint was built from. Everything round-trips, and the only failures before the change were the tied-weight-plus-safetensors combinations.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 235e5e2578

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

Avoid copying placeholder tensors during sharding

When safe_serialization=True with the default max_shard_size, the sharding pass calls this in return_empty_tensor=True mode to size shards without materializing weights, but a tied GB-scale weight now hits this branch and clone() copies/touches the entire placeholder tensor. That can allocate/commit another full embedding during planning and OOM before shard saving even when the tied keys would be split; use a fresh empty tensor for the return_empty_tensor case instead of cloning.

Useful? React with 👍 / 👎.

[vineethsaivs]

Good catch, fixed in 3de47db. The sizing pass only needs shapes, so cloning one placeholder into another was pure cost, and it lands exactly where it hurts: a tied embedding, planned before any shard is written.

        if tensor_id in converted_tensors:  # shared tensors
            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
            if share_tensors:
                torch_state_dict[name] = shared_tensor
            elif return_empty_tensor:
                torch_state_dict[name] = torch.empty(shared_tensor.shape, dtype=shared_tensor.dtype)
            else:
                torch_state_dict[name] = shared_tensor.clone()

The separate storage still has to exist, or split_torch_state_dict_into_shards would dedupe the tied keys onto one pointer and undercount the file it is about to plan, so this allocates rather than sharing.

New test_sizing_pass_does_not_clone_a_tied_placeholder spies on Tensor.clone and asserts the planning call makes none, while the saving call still makes exactly one and the two written tensors are equal with different data_ptr(). Against the previous head it fails with AssertionError: the sizing pass cloned [torch.Size([16])].

$ python -m pytest unit/checkpoint/test_convert_checkpoint.py -k "not TestConvert"
before: 1 failed, 6 passed, 1 skipped
after:  7 passed, 1 skipped