Fixed weight tying causes corrupted safetensors file for lm_head.weight in HF checkpoint (#3764)

## Problem
When training with qwen3 1.7B, all configs set
`enable_weight_tying=True` in torchtitan/models/common/decoder.py, so in
`to_hf`, the converter skips `lm_head.weight` because the tied weight is
already covered by `model.embed_tokens.weight`.
However, `fqn_to_index_mapping` is loaded verbatim from
`model.safetensors.index.json`， which maps `lm_head.weight` to the 2nd
safetensors file. When `to_hf` returns, the mapping references a key
absent from the converted state dict, causing the creation of an output
file for `lm_head.weight` with empty contents: `_FqnData` (`shape=[]`,
`dtype=""`, `data_offsets=[0,0]`), and when this checkpoint is loaded,
it will get rejected:
   ```
SafetensorError: Error while deserializing header: invalid JSON in
header:
   unknown variant ``, expected one of `F32`, `F16`, `BF16`, `F8_E4M3`,
   `F8_E5M2`, `I64`, `I32`, `I16`, `I8`, `U8`, `BOOL` at line 1 column 1
   ```

## Root cause
In torchtitan/components/checkpoint.py, `HuggingFaceStorageWriter` is
created with the original mapping, and passed in to
torch.distributed.save().

## Reproduction

I use commit #b70844d5, but should be reproducible on the latest
version.
Training:
```bash
bash run_train.sh--module qwen3     --config qwen3_1_7b     --checkpoint.last_save_model_only=True     --checkpoint.last_save_in_hf=True     --checkpoint.export_dtype=float16     --training.steps=3     --dataloader.dataset=c4_test     --checkpoint.enable=True
```

The safetensors produced by this training run should be corrupted.
Observe that there is a 78bytes `model-00002-of-00002.safetensors` in
outputs/checkpoint/step-3.

Verify the break by loading from the folder:
```bash
python3 -c "
from safetensors import safe_open
safe_open('outputs/checkpoint/step-3/model-00002-of-00002.safetensors', framework='pt') "

```


## Proposed fix

See changes: In torchtitan/models/qwen3/state_dict_adapter.py, at the
end of to_hf(), before return hf_state_dict, filter
fqn_to_index_mapping. This prevents the empty entry from being saved.

I’m offering this as a temporary fix for Qwen3 (Qwen3.5 and Llama also
seem susceptible to this problem), and would love to hear from the
maintainers on a more general, robust solution.

---------

Co-authored-by: Mystri <hanboyou@huawei.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

torchtitan/distributed/utils.py

@@ -546,9 +546,7 @@ def set_pg_timeouts(
         for mesh in parallel_dims.get_all_one_dimensional_meshes().values()
     ] + [None]
     for group in groups:
-        torch.distributed.distributed_c10d._set_pg_timeout(
-            timeout, group
-        )  # pyrefly: ignore[deprecated]
+        torch.distributed.set_timeout(timeout, group)
 
 
 @torch.no_grad()

torchtitan/models/llama3/state_dict_adapter.py

@@ -101,6 +101,8 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                     self.model_config.enable_weight_tying  # pyrefly: ignore [missing-attribute]
                     and key == "lm_head.weight"
                 ):
+                    if self.fqn_to_index_mapping:
+                        self.fqn_to_index_mapping.pop("lm_head.weight", None)
                     continue
                 new_key = to_hf_map[key]
 

torchtitan/models/qwen3/state_dict_adapter.py

@@ -120,6 +120,8 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                     continue
                 # pyrefly: ignore [missing-attribute]
                 if self.model_config.enable_weight_tying and key == "lm_head.weight":
+                    if self.fqn_to_index_mapping:
+                        self.fqn_to_index_mapping.pop("lm_head.weight", None)
                     continue
                 new_key = to_hf_map[key]
                 hf_state_dict[new_key] = value

torchtitan/models/qwen3_5/state_dict_adapter.py

@@ -186,6 +186,8 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]:
                 if tt_key == "lm_head.weight" and getattr(
                     self.model_config, "enable_weight_tying", False
                 ):
+                    if self.fqn_to_index_mapping:
+                        self.fqn_to_index_mapping.pop("lm_head.weight", None)
                     continue
                 hf_value = value
                 # Linear weight (out, C*T*H*W) → Conv3d weight (out, C, T, H, W)