Adds parallel model weight loading for runai_streamer

[bbartels]

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

The runai streamer library published a new way to load files in parallel as opposed to one at a time. This will allow for a much greater level of parallelism, especially when files are small.

Test Plan

The existing tests cover the code change

Test Result

(Optional) Documentation Update

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

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 either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request aims to add parallel model weight loading for runai_streamer. The implementation introduces a new version of runai_safetensors_weights_iterator that uses streamer.stream_files for parallel processing.

However, there are a couple of significant issues:

1. The old runai_safetensors_weights_iterator function was not removed, resulting in a duplicate function definition in the same file.
2. The tqdm progress bar in the new function is misconfigured, which will lead to an incorrect progress display and issues in distributed environments.

I've provided comments with suggestions to fix these issues.

[gemini-code-assist]

This pull request introduces a second definition for the function runai_safetensors_weights_iterator. The original implementation exists at lines 479-492, and this new implementation will shadow it. This is likely unintentional and makes the code confusing.

Please remove the old implementation and keep this new one, which seems to correctly implement parallel weight loading as intended by the PR.

I cannot suggest the deletion directly as the original function is not part of this diff, but it should be removed to resolve this issue.

[gemini-code-assist]

The tqdm progress bar is misconfigured here, which can lead to incorrect behavior:

1. Incorrect total: The total is set to len(hf_weights_files), which is the number of files. However, tqdm is iterating over tensors from streamer.get_tensors(). Since there are usually multiple tensors per file, the progress bar will show incorrect percentages (e.g., > 100%). If the total number of tensors is not available, it's better to omit the total argument. tqdm will then display the iteration count instead of a percentage.

2. Incorrect disable logic for distributed environments: The disable argument should be not enable_tqdm(use_tqdm_on_load) to ensure the progress bar is only displayed on the main process (rank 0) in a distributed setting. The current implementation not use_tqdm_on_load will cause progress bars to be printed from all worker processes.

Here is a suggested fix:

        tensor_iter = tqdm(
            streamer.get_tensors(),
            desc="Loading safetensors using Runai Model Streamer",
            bar_format=_BAR_FORMAT,
            disable=not enable_tqdm(use_tqdm_on_load),
        )

[bbartels]

@simon-mo @khluu I don't believe I have perms to kick off the full build. Could you kick it off for me? Or allow me to run them?

[bbartels]

@DarkLight1337 Not sure whom to ping but you have reviewed existing runai related prs. Pipeline seems to have failed for unrelated reason. Any chance a full build can be kicked off?

[DarkLight1337]

Triggering full CI

[DarkLight1337]

The test failures are unrelated to this PR. Can you address my previous comment? Then we can merge the PR

[bbartels]

The test failures are unrelated to this PR. Can you address my previous comment? Then we can merge the PR

All done :)