[Core] Integrate fastsafetensors loader for loading model weights

[manish-sethi]

This PR integrates fastsafetensor library for loading model weights. Fastsafetensor uses GPU direct storage to load the weights directly from storage into GPU memory, making it faster as compared to reading parameters iteratively from file. See https://github.com/foundation-model-stack/fastsafetensors for more details.

Following are the experimental results for the time taken to load the weights from NVMe using the existing safetensor loader v/s fastsafetensor loader.

Model	Tensor_parallel	GPU	Time with safetensor loader(sec)	Time with fastsafetensor loader(sec)
Llama-2-7b-hf	1	A100	7.29	3.67
Llama-2-13b-hf	1	A100	16.04	6.88
Llama-2-13b-hf	2	A100	14.35	6.02
Llama-2-13b-hf	4	L40S	12.39	4.74

• The model files are loaded from storage, i.e., not present in the filesystem cache
• The time is for only loading the model weights within vLLM model initialization path

Additional performance number comparing three different loaders:

Tensor_parallel	safetensor loader(sec)	runai_streamer loader(sec)	fastsafetensor loader(sec)
1	50.05	36.63	42.40
2	55.98	45.86	42.19
4	60.08	59.98	53.50

• The above numbers are with Model Llama-2-13b-hf on L40S GPU
• The timing includes end-to-end timing from loading and initializing the model to the generation of the first token with eager-mode
• Like before, the model files are loaded from storage, i.e., not present in the filesystem cache

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @manish-sethi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[manoelmarques]

Hi @manish-sethi,

This PR is very useful thank you for creating it.
Do you intend to finish it ? It seems to me it is basically done right ? I used it in some tests.
Also, there are no reviewers added, adding reviewers would accelerate the process.

[njhill]

Thanks @manish-sethi, this looks great!

It would be awesome to see how this compares to tensorizer and runai-model-streamer performance-wise. Maybe you could add a couple more columns to the table?

For measuring the performance (of all of theses) it would probably be best to have a test that loads the model and then immediately performs a generate request of a single token (i.e. time to first token). To make sure there isn't "hidden" latent loading time (of course also ensuring the file system cache is cold too as normal).

Could you also add a CI test - you can look at what there is for the other model loaders as an example.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @manish-sethi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[njhill]

Thanks @manish-sethi these updates look really good.

Just one minor inline comment, and also I think you will need to add the dependency to this file for the tests.

If you take the PR out of draft after the above changes we can start to run the tests in the CI.

Also this is mainly a curiosity but for completeness any chance of also adding a column for tensorizer to the new table?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @manish-sethi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[manish-sethi]

Thanks @manish-sethi these updates look really good.

Just one minor inline comment, and also I think you will need to add the dependency to this file for the tests.

Added the dependency on test.ini. Can it be assumed that the test machine would have GPUs?

If you take the PR out of draft after the above changes we can start to run the tests in the CI.

Took out of draft

Also this is mainly a curiosity but for completeness any chance of also adding a column for tensorizer to the new table?

I skipped the tensorizer earlier as it was not an apple-to-apple comparison. Still, gave it a try now, but the code seems to be broken. I ran into an error when used this example for serializing the model.

 File "/nvme/manish/repos/vllm/vllm/model_executor/model_loader/tensorizer.py", line 465, in tensorize_vllm_model
    engine.model_executor.collective_rpc(
    ^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'LLMEngine' object has no attribute 'model_executor'

[manoelmarques]

In my tests integrating this PR to vLLM code and running, it is way slower than safetensor when GDS is not installed.
Is there any test being done for the case where GDS is not installed ?
It would be good to expose the nogds flag in the loader in a next PR. The default is False.

loader = SafeTensorsFileLoader(pg, device, nogds=nogds)

[njhill]

Thanks @manish-sethi for the great work!

WDYT about @manoelmarques's comment? Would it be worth incorporating a check for GDS?

[manish-sethi]

Thanks @manish-sethi for the great work!

WDYT about @manoelmarques's comment? Would it be worth incorporating a check for GDS?

I am not sure if there is a straight forward way to test that programmatically. I guess better would be to open an issue against fastsafetensors repo so it can be exposed as an API. It may require some investigation and also, the fastsafetensors repo is a better place to maintain that code instead of vllm repo from reusability and maintenance purposes.

[zhewenl]

@manish-sethi / @njhill This PR needs to skip tests on AMD, fixing in #27612