runai_model_streamer.md 4.81 KB
Newer Older
1
# Loading models with Run:ai Model Streamer
2
3
4
5
6
7
8

Run:ai Model Streamer is a library to read tensors in concurrency, while streaming it to GPU memory.
Further reading can be found in [Run:ai Model Streamer Documentation](https://github.com/run-ai/runai-model-streamer/blob/master/docs/README.md).

vLLM supports loading weights in Safetensors format using the Run:ai Model Streamer.
You first need to install vLLM RunAI optional dependency:

9
```bash
10
pip3 install vllm[runai]
11
12
13
14
```

To run it as an OpenAI-compatible server, add the `--load-format runai_streamer` flag:

15
```bash
Reid's avatar
Reid committed
16
17
vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer
18
19
20
21
```

To run model from AWS S3 object store run:

22
```bash
Reid's avatar
Reid committed
23
vllm serve s3://core-llm/Llama-3-8b \
24
25
26
27
28
29
30
    --load-format runai_streamer
```

To run model from Google Cloud Storage run:

```bash
vllm serve gs://core-llm/Llama-3-8b \
Reid's avatar
Reid committed
31
    --load-format runai_streamer
32
33
```

34
35
36
37
38
39
40
41
42
43
To run model from Azure Blob Storage run:

```bash
AZURE_STORAGE_ACCOUNT_NAME=<account> \
vllm serve az://<container>/<model-path> \
    --load-format runai_streamer
```

Authentication uses `DefaultAzureCredential`, which supports `az login`, managed identity, environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`), and other methods.

44
45
To run model from a S3 compatible object store run:

46
```bash
Reid's avatar
Reid committed
47
48
49
50
51
RUNAI_STREAMER_S3_USE_VIRTUAL_ADDRESSING=0 \
AWS_EC2_METADATA_DISABLED=true \
AWS_ENDPOINT_URL=https://storage.googleapis.com \
vllm serve s3://core-llm/Llama-3-8b \
    --load-format runai_streamer
52
53
54
55
56
57
```

## Tunable parameters

You can tune parameters using `--model-loader-extra-config`:

58
59
60
61
62
63
64
65
66
You can tune `distributed` that controls whether distributed streaming should be used. This is currently only possible on CUDA and ROCM devices. This can significantly improve loading times from object storage or high-throughput network fileshares.
You can read further about Distributed streaming [here](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/usage.md#distributed-streaming)

```bash
vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer \
    --model-loader-extra-config '{"distributed":true}'
```

67
68
69
You can tune `concurrency` that controls the level of concurrency and number of OS threads reading tensors from the file to the CPU buffer.
For reading from S3, it will be the number of client instances the host is opening to the S3 server.

70
```bash
Reid's avatar
Reid committed
71
72
73
vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer \
    --model-loader-extra-config '{"concurrency":16}'
74
75
```

76
You can control the size of the CPU Memory buffer to which tensors are read from the file, and limit this size.
77
78
You can read further about CPU buffer memory limiting [here](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/env-vars.md#runai_streamer_memory_limit).

79
```bash
Reid's avatar
Reid committed
80
81
82
vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer \
    --model-loader-extra-config '{"memory_limit":5368709120}'
83
84
```

85
86
!!! note
    For further instructions about tunable parameters and additional parameters configurable through environment variables, read the [Environment Variables Documentation](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/env-vars.md).
87
88
89
90
91

## Sharded Model Loading

vLLM also supports loading sharded models using Run:ai Model Streamer. This is particularly useful for large models that are split across multiple files. To use this feature, use the `--load-format runai_streamer_sharded` flag:

92
```bash
93
94
95
96
97
vllm serve /path/to/sharded/model --load-format runai_streamer_sharded
```

The sharded loader expects model files to follow the same naming pattern as the regular sharded state loader: `model-rank-{rank}-part-{part}.safetensors`. You can customize this pattern using the `pattern` parameter in `--model-loader-extra-config`:

98
```bash
Reid's avatar
Reid committed
99
100
101
vllm serve /path/to/sharded/model \
    --load-format runai_streamer_sharded \
    --model-loader-extra-config '{"pattern":"custom-model-rank-{rank}-part-{part}.safetensors"}'
102
103
```

104
To create sharded model files, you can use the script provided in [examples/offline_inference/save_sharded_state.py](../../../examples/offline_inference/save_sharded_state.py). This script demonstrates how to save a model in the sharded format that is compatible with the Run:ai Model Streamer sharded loader.
105
106
107

The sharded loader supports all the same tunable parameters as the regular Run:ai Model Streamer, including `concurrency` and `memory_limit`. These can be configured in the same way:

108
```bash
Reid's avatar
Reid committed
109
110
111
vllm serve /path/to/sharded/model \
    --load-format runai_streamer_sharded \
    --model-loader-extra-config '{"concurrency":16, "memory_limit":5368709120}'
112
113
```

114
115
!!! note
    The sharded loader is particularly efficient for tensor or pipeline parallel models where each worker only needs to read its own shard rather than the entire checkpoint.