Safetensors

[gabe-l-hart]

Description

Closes #1249

This PR implements support for downloading and converting model checkpoints from huggingface which use the safetensors format rather than the pth binary (pickle) format.

Changes

• Allow the tensor map file to be found under different *.index.json names
• Allow loading model files with safetensors.torch.load when needed
• Allow downloading safetensor files if they exist without pth files or the model explicitly prefers them

Testing

• I have tested that the download and load for llama3.1 are un-changed with these changes (no safetensors are downloaded)
• I have verified (on my WIP branch for Granite Code) that for models with only safetensors, they are not ignored and can be cleanly converted

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1255

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9db974d with merge base d8c0aaf ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[facebook-github-bot]

Hi @gabe-l-hart!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[gabe-l-hart]

NOTE: The CLA is in process since I'm contributing through the IBM corporate CLA.

[byjlw]

Thanks @gabe-l-hart
I appreciate you contributing. Code looks good, please rebase when you're ready and i'll approve and merge.
Also since it looks like you're already adding another model to the list and have tested it feel free to bring in the definition so others can benefit by downloading via the cli

[gabe-l-hart]

Thanks @byjlw! I'll get the rebase done shortly. As you noted, this is part of my work to add Granite Code support. I have a single branch pointer for all that work, but I have also tried to organize the commits so that they can be merged as bite-sized features rather than a big overhaul. Right now, the addition of the configs for Granite Code are at the end of the chain since the model won't work at all until the other features are present (it also currently only works in the python runtime). I'm totally open to whatever convention the torchchat team prefers in terms of contribution chunking, so just let me know what you prefer!

[byjlw]

chunks are definitely better. Would love to learn more about the overall goals for the Granite code support

[gabe-l-hart]

Would love to learn more about the overall goals for the Granite code support

Good point, I didn't ever spell this out! Let me open a top-level issue about that support that I can use to track it all.

[gabe-l-hart]

Top-level Granite Code support issue: #1262

[byjlw]

Actually though, when i tested it by changing the the model.json to use safetensors for 11b base download errored out. It looks like this code works as long as it's not using definitions that use the torchtune format. Do you mind testing this case and resolving the issue?

"meta-llama/Llama-3.2-11B-Vision": {
        "aliases": ["llama3.2-11B-base", "Llama-3.2-11B-Vision-base"],
        "distribution_channel": "HuggingFaceSnapshot",
        "distribution_path": "meta-llama/Llama-3.2-11B-Vision",
        "prefer_safetensors": true
    },

Fetching 19 files: 100%|█████████████████████████████████████████████████████████████████████| 19/19 [1:06:30<00:00, 210.04s/it]
Converting meta-llama/Llama-3.2-11B-Vision to torchtune format...
Traceback (most recent call last):100%|████████████████████████████████████████████████████| 4.99G/4.99G [30:11<00:00, 3.52MB/s]
  File "/Users/byjlw/Documents/source/working/torchchat/torchchat.py", line 85, in <module>| 4.92G/4.92G [12:17<00:00, 3.34MB/s]
    check_args(args, "generate")
  File "/Users/byjlw/Documents/source/working/torchchat/torchchat/cli/cli.py", line 52, in check_args
    download_and_convert(args.model, args.model_directory, args.hf_token)
  File "/Users/byjlw/Documents/source/working/torchchat/torchchat/cli/download.py", line 123, in download_and_convert
    _download_hf_snapshot(model_config, temp_dir, hf_token)
  File "/Users/byjlw/Documents/source/working/torchchat/torchchat/cli/download.py", line 82, in _download_hf_snapshot
    convert_hf_checkpoint_to_tune( model_dir=artifact_dir, model_name=model_config.name)
  File "/Users/byjlw/Documents/source/working/torchchat/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/byjlw/Documents/source/working/torchchat/torchchat/cli/convert_hf_checkpoint.py", line 184, in convert_hf_checkpoint_to_tune
    raise RuntimeError(f"Could not find {consolidated_pth}")
RuntimeError: Could not find /Users/byjlw/.torchchat/model-cache/downloads/meta-llama/Llama-3.2-11B-Vision/original/consolidated.pth

[gabe-l-hart]

Ah, good catch! I guess N == 2 is a small sample size. I'll dig into your error and see how far I can get.

[byjlw]

@ebsmothers can also help

[gabe-l-hart]

Great. I'll report progress or blockers as they come up. @ebsmothers let me know if you dig in and get anywhere!

[ebsmothers]

@byjlw so I'm not super familiar with how torchchat is handling checkpoint conversion but if you're switching from the .pth format to .safetensors format you will no longer be able to just do the simple move that's happening here. The safetensors format for Llama 3.2 11B Vision distributes the model weights across multiple files (see here) so they will need to be loaded and merged into a single state dict as we do in torchtune here.

[gabe-l-hart]

That makes sense. A similar approach is being taken in convert_hf_checkpoint to load the sharded weights and convert them into the .pth format needed by torchchat. If I'm understanding the logic in convert_hf_checkpoint_to_tune correctly, it looks like it's really just a simpler version of the logic in convert_hf_checkpoint that doesn't require the tensor renaming or permuting. If that's the case, I think the fix should be to hoist out the shard-loading logic into a helper and then only do the post-processing in convert_hf_checkpoint before resaving as model.pth.

I finally have the safetensors downloaded locally, so I'll see how far I can get with this approach.

[gabe-l-hart]

I have the mechanics of this working for the safetensors weights with Llama-3.2-11B-Vision-Instruct, but it appears that there is a naming difference between the tensors in original/consolidated.pth and the safetensors weights. This is likely similar to the name remapping needed in the standard conversion logic.

Given that, and given the fact that these models do already have checkpoints that match the target naming scheme, I think it might make sense to leave the PR as-is for now and not switch prefer_safetensors to true for these models. I could also see an argument for using safetensors all-the-way-through to avoid the known pickle vulnerabilities with .pth, but this PR doesn't address that anyway since the models are converted to .pth during the conversion process.

[leseb]

Why is this an error? Thanks!

[gabe-l-hart]

Good catch. See my response over on your PR: #1346 (comment)