Missing model.safetensors.index.json in final hf checkpoint

[legarwin]

Bug description

When using torchtitan with huggingface checkpoint for in and output, there seems to be a incosistent wrong behaviour for the save hf model.

If you provide a hf model with missing model.safetensors.index.json you get a fair warning

WARNING - model.safetensors.index.json not found at hf_assets_path: ./tests/assets/tokenizer/model.safetensors.index.json. Defaulting to saving a single safetensors file if checkpoint is saved in HF format

and the saved final hf checkpoint behaves as expected. It saves a single hf safetensor (no matter if there were split safetensor files before) + index.json.

However, if the model.safetensors.index.json is not missing when loading the initial hf model, the saved final hf checkpoint seem to behave wrong. It keeps the safetensor splits for the output (based on the input model), but the model.safetensors.index.json is completely missing.

See the output here for a debug run with a single safetensor for the input and output model:

I think the issue can be found in this if condition:
The first one skips the step of save the model.safetensors.index.json and also the call of method consolidate_safetensors_files_on_every_rank is not saving it.

Can you have a look into that? Thank you! :)

Versions

1. Tested on latest pytorch nightly: 2.10.0.dev20251210+cu128
2. Use default debug model toml torchtitan/models/llama3/train_configs/debug_model.toml with checkpoint enabled and an adapted one where i pointed to a local downloaded model with tokenizer and model.safetensors.index.json, see

[model]
hf_assets_path = "<local-path>/llama3_debugmodel"

[checkpoint]
enable = true
folder = "checkpoint"
interval = 10
export_dtype = "float32"
async_mode = "disabled"  # ["disabled", "async", "async_with_pinned_mem"]
last_save_model_only = true
last_save_in_hf = true
initial_load_path = "<local-path>/llama3_debugmodel"
initial_load_in_hf = true

3. Tested issue on torchtitan run_train.sh with no modifications

[wwwjn]

Thank you for your report!

It keeps the safetensor splits for the output (based on the input model), but the model.safetensors.index.json is completely missing.

The behavior of "keeps the safetensor splits" is expected, and we intentionally want to keep the same splits.
But I think we do need to make a copy of model.safetensors.index.json and put it in path of saving checkpoint, so next time when loading from the saved HF checkpoints it functions as expected. Do you want to create a PR to fix this behavior?

[legarwin]

@wwwjn thanks for your feedback :). I totally agree with it.
When i will find time for it next week, i will try to prepare a PR.

A bit off-topic: Do you think it would be beneficial to also dump the tokenizer to the hf output checkpoint?
Then it would better align with the general hf format and maybe someone changes the tokenizer while training, e.g. using unused token.

Qwen3 HF example:

[tianyu-l]

You could use https://github.com/pytorch/torchtitan/blob/main/scripts/download_hf_assets.py to download anything that is available on HF.

But I think we do need to make a copy of model.safetensors.index.json and put it in path of saving checkpoint

This should be the job of DCP checkpoint saving, rather than torchtitan job??

[legarwin]

@tianyu-l you are saying it should go to pytorch? I agree. I just noticed the issue while using torchtitan.
Do you think i should move the issue or do you think it is fine to keep it here and i just open a PR on pytorch side?

Probably that's the place where the fix should go to: https://github.com/pytorch/pytorch/blob/main/torch/distributed/checkpoint/_consolidate_hf_safetensors.py

[tianyu-l]

cc @ankitageorge
People would like to see index file also propagated in the save folder. Do you think DCP should write it as the final source of truth?

[ankitageorge]

cc @ankitageorge People would like to see index file also propagated in the save folder. Do you think DCP should write it as the final source of truth?

ya that makes sense, I'll work on updating this in DCP

[legarwin]

@ankitageorge so are you taking care for the change in pytorch or shall i open a PR with a fix for it?