mistralrs_core/pipeline/loaders/
auto_device_map.rs1use std::fmt::{self, Display};
2
3use crate::paged_attention::{
4 calculate_cache_config, ModelConfigLike, DEFAULT_PAGED_ATTENTION_BLOCK_SIZE,
5};
6use crate::utils::debug::DeviceRepr;
7use crate::{DeviceLayerMapMetadata, DeviceMapMetadata, MemoryUsage, PagedAttentionConfig};
8use anyhow::{Context, Result};
9use candle_core::{DType, Device};
10use itertools::Itertools;
11use tracing::{info, warn};
12
13use super::DeviceMappedModelLoader;
14
15#[derive(Clone, Debug)]
16pub(crate) enum NonMappedSubModel {
17 Vision,
18}
19
20impl Display for NonMappedSubModel {
21 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
22 match self {
23 NonMappedSubModel::Vision => write!(f, "vision"),
24 }
25 }
26}
27
28#[derive(Debug, Clone)]
29pub enum AutoDeviceMapParams {
30 Text {
31 max_seq_len: usize,
32 max_batch_size: usize,
33 },
34 Vision {
35 max_seq_len: usize,
36 max_batch_size: usize,
37 max_image_shape: (usize, usize),
38 max_num_images: usize,
39 },
40}
41
42impl AutoDeviceMapParams {
43 pub fn maybe_promote_to_vision(&self) -> Self {
44 match *self {
45 Self::Text {
46 max_seq_len,
47 max_batch_size,
48 } => Self::Vision {
49 max_seq_len,
50 max_batch_size,
51 max_image_shape: (
52 Self::DEFAULT_MAX_IMAGE_LENGTH,
53 Self::DEFAULT_MAX_IMAGE_LENGTH,
54 ),
55 max_num_images: Self::DEFAULT_MAX_NUM_IMAGES,
56 },
57 Self::Vision {
58 max_seq_len,
59 max_batch_size,
60 max_image_shape,
61 max_num_images,
62 } => Self::Vision {
63 max_seq_len,
64 max_batch_size,
65 max_image_shape,
66 max_num_images,
67 },
68 }
69 }
70
71 pub fn max_seq_len(&self) -> usize {
72 match self {
73 Self::Text { max_seq_len, .. } | Self::Vision { max_seq_len, .. } => *max_seq_len,
74 }
75 }
76}
77
78impl Display for AutoDeviceMapParams {
79 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80 match self {
81 Self::Text {
82 max_seq_len,
83 max_batch_size,
84 } => write!(
85 f,
86 "text[max_seq_len: {max_seq_len}, max_batch_size: {max_batch_size}]"
87 ),
88 Self::Vision {
89 max_seq_len,
90 max_batch_size,
91 max_image_shape,
92 max_num_images,
93 } => write!(
94 f,
95 "vision[max_seq_len: {max_seq_len}, max_batch_size: {max_batch_size}, max_image_shape: {max_image_shape:?}, max_num_images: {max_num_images}]"
96 ),
97 }
98 }
99}
100
101impl AutoDeviceMapParams {
102 pub const DEFAULT_MAX_SEQ_LEN: usize = 4 * 1024;
104 pub const DEFAULT_MAX_BATCH_SIZE: usize = 1;
105 pub const DEFAULT_MAX_NUM_IMAGES: usize = 1;
106 pub const DEFAULT_MAX_IMAGE_LENGTH: usize = 1024;
107
108 pub fn default_text() -> Self {
109 Self::Text {
110 max_seq_len: Self::DEFAULT_MAX_SEQ_LEN,
111 max_batch_size: Self::DEFAULT_MAX_BATCH_SIZE,
112 }
113 }
114
115 pub fn default_vision() -> Self {
116 Self::Vision {
117 max_seq_len: Self::DEFAULT_MAX_SEQ_LEN,
118 max_batch_size: Self::DEFAULT_MAX_BATCH_SIZE,
119 max_num_images: Self::DEFAULT_MAX_NUM_IMAGES,
120 max_image_shape: (
121 Self::DEFAULT_MAX_IMAGE_LENGTH,
122 Self::DEFAULT_MAX_IMAGE_LENGTH,
123 ),
124 }
125 }
126}
127
128fn calculate_key_block_shape(
129 model_config: &dyn ModelConfigLike,
130 dtype: DType,
131 block_size: usize,
132) -> (usize, usize, usize, usize) {
133 let element_size = dtype.size_in_bytes();
134 let x = 16 / element_size;
135 (
136 model_config.num_kv_heads(),
137 model_config.k_head_dim() / x,
138 block_size,
139 x,
140 )
141}
142
143fn calculate_value_block_shape(
144 model_config: &dyn ModelConfigLike,
145 block_size: usize,
146) -> (usize, usize, usize) {
147 (
148 model_config.num_kv_heads(),
149 model_config.v_head_dim(),
150 block_size,
151 )
152}
153
154macro_rules! b_to_mb {
155 ($x:expr) => {
156 $x / (1024 * 1024)
157 };
158}
159
160#[allow(clippy::too_many_arguments)]
161pub fn get_device_layers(
163 loader: &dyn DeviceMappedModelLoader,
164 config: &str,
165 num_layers: usize,
166 mut layer_sizes_in_bytes: Vec<usize>,
167 non_mapped_size_in_bytes: usize,
168 total_model_size_in_bytes: usize,
169 devices: &[Device],
170 dtype: DType,
171 params: &AutoDeviceMapParams,
172 prompt_chunksize: usize,
173 paged_attn_config: Option<&PagedAttentionConfig>,
174) -> Result<DeviceMapMetadata> {
175 let mapped_max =
176 loader.mapped_max_act_size_elems(config, params, prompt_chunksize)? * dtype.size_in_bytes();
177 let non_mapped_max =
178 loader.non_mapped_max_act_size_elems(config, params)? * dtype.size_in_bytes();
179
180 let mut remaining = total_model_size_in_bytes;
181 let max_seq_len = match params {
182 AutoDeviceMapParams::Text { max_seq_len, .. }
183 | AutoDeviceMapParams::Vision { max_seq_len, .. } => *max_seq_len,
184 };
185 let max_batch_size = match params {
186 AutoDeviceMapParams::Text { max_batch_size, .. }
187 | AutoDeviceMapParams::Vision { max_batch_size, .. } => *max_batch_size,
188 };
189
190 let model_cfg = loader.model_config(config)?;
191 let kv_cache_elems = match paged_attn_config {
192 Some(cfg) => {
193 let cache = calculate_cache_config(
194 cfg.mem_gpu,
195 cfg.mem_cpu,
196 Some(cfg.block_size.unwrap_or(DEFAULT_PAGED_ATTENTION_BLOCK_SIZE)),
197 dtype,
198 paged_attn_config
199 .map(|cfg| cfg.cache_type)
200 .unwrap_or_default(),
201 &*model_cfg,
202 &devices[0],
203 &devices.iter().map(|d| Some(d.clone())).collect::<Vec<_>>(),
204 true,
205 )?;
206 let key_shape = calculate_key_block_shape(&*model_cfg, dtype, cache.block_size);
207 let key_sz =
208 cache.num_gpu_blocks * key_shape.0 * key_shape.1 * key_shape.2 * key_shape.3;
209 let val_shape = calculate_value_block_shape(&*model_cfg, cache.block_size);
210 let val_sz = cache.num_gpu_blocks * val_shape.0 * val_shape.1 * val_shape.2;
211 key_sz + val_sz
212 }
213 None => {
214 let key_shape = [
215 max_batch_size,
216 model_cfg.num_kv_heads(),
217 max_seq_len,
218 model_cfg.k_head_dim(),
219 ];
220 let val_shape = [
221 max_batch_size,
222 model_cfg.num_kv_heads(),
223 max_seq_len,
224 model_cfg.v_head_dim(),
225 ];
226 key_shape.iter().product::<usize>() + val_shape.iter().product::<usize>()
227 }
228 };
229 let kv_cache_bytes = kv_cache_elems * dtype.size_in_bytes();
230
231 let mut avail = Vec::new();
233 for dev in [devices, &[Device::Cpu]].concat() {
234 let a = MemoryUsage.get_memory_available(&dev)?;
235 avail.push((a, dev));
236 }
237 avail.reverse();
238 layer_sizes_in_bytes.reverse();
239
240 let mut mappings = Vec::new();
241 info!("Using automatic device mapping parameters: {params}.");
242 if let Some(subs) = loader.non_mapped_sub_models() {
243 let (_, last) = avail.last().unwrap();
244 info!(
245 "The following sub-models will not be device mapped and will be loaded on {}: {}",
246 last.device_pretty_repr(),
247 subs.iter().map(|x| x.to_string()).join(", ")
248 );
249 }
250
251 let mut ordinal = 0;
252 let mut layer = 0;
253 let avail_copy = avail.clone();
254 let mut includes_cpu = false;
255 while remaining > 0 && !avail.is_empty() {
256 let (cap, dev) = avail
257 .pop()
258 .context("No more devices to map to. The model does not fit on this system.")?;
259
260 #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss)]
262 let cap = (cap as f64 * 0.90) as usize;
263
264 let required_whole_capacity = if ordinal == 0 {
271 remaining
272 + non_mapped_max.max(mapped_max)
273 + non_mapped_size_in_bytes
274 + kv_cache_bytes * (num_layers - layer)
275 } else {
276 remaining + mapped_max + kv_cache_bytes * (num_layers - layer)
277 };
278
279 let layers_on_dev = if cap >= required_whole_capacity {
280 remaining = 0;
281 num_layers - layer
282 } else {
283 let mut used = mapped_max;
284 let mut used_no_act = 0;
285 let mut count = 0;
286 if ordinal == 0 {
287 used = used.max(non_mapped_max) + non_mapped_size_in_bytes;
288 used_no_act += non_mapped_size_in_bytes;
289 }
290 while let Some(&sz) = layer_sizes_in_bytes.last() {
291 let delta = sz + kv_cache_bytes;
292 if used + delta > cap {
293 break;
294 }
295 layer_sizes_in_bytes.pop();
296 used += delta;
297 used_no_act += delta;
298 count += 1;
299 }
300 if count > 0 {
301 remaining = remaining.saturating_sub(used_no_act);
302 } else {
303 warn!(
304 "Device {} can fit 0 layers. Consider reducing auto map params from current: {params} (ex. reducing max seq len or max num images)",
305 dev.device_pretty_repr(),
306 );
307 ordinal += 1;
308 continue;
309 }
310 count
311 };
312 if !dev.is_cpu() {
313 mappings.push(DeviceLayerMapMetadata {
314 ordinal,
315 layers: layers_on_dev,
316 });
317 ordinal += 1;
318 } else {
319 includes_cpu = true;
320 }
321 layer += layers_on_dev;
322 }
323 if remaining > 0 {
324 let over = b_to_mb!(remaining);
325 anyhow::bail!(
326 "This model does not fit on the devices {:?}, and exceeds total capacity by {}MB. Auto device mapping params: {params}",
327 avail_copy.iter().rev().map(|(a, d)| format!("{} (avail: {}MB)", d.device_pretty_repr(), b_to_mb!(a))).collect::<Vec<_>>(),
328 over
329 );
330 }
331 if paged_attn_config.is_some_and(|_| includes_cpu) {
332 return get_device_layers(
333 loader,
334 config,
335 num_layers,
336 layer_sizes_in_bytes,
337 non_mapped_size_in_bytes,
338 total_model_size_in_bytes,
339 devices,
340 dtype,
341 params,
342 prompt_chunksize,
343 None,
344 );
345 }
346 Ok(DeviceMapMetadata::from_num_device_layers(mappings))
347}