llm_interface 0.0.3

llm_interface: The backend for the llm_client crate
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

use llm_devices::DeviceConfig;
use mistralrs::{Device, DeviceLayerMapMetadata, DeviceMapMetadata};

pub fn mistral_rs_device_map(
    generic_device_map: &DeviceConfig,
) -> crate::Result<(Device, DeviceMapMetadata)> {
    match generic_device_map.gpu_count() {
        0 => new_only_cpu(generic_device_map),
        1 => new_single_gpu(generic_device_map),
        _ => new_multiple_gpu(generic_device_map),
    }
}

fn new_only_cpu(generic_device_map: &DeviceConfig) -> crate::Result<(Device, DeviceMapMetadata)> {
    std::env::set_var(
        "RAYON_NUM_THREADS",
        generic_device_map
            .cpu_config
            .thread_count_or_default()
            .to_string(),
    );
    Ok((Device::Cpu, DeviceMapMetadata::dummy()))
}

fn new_single_gpu(generic_device_map: &DeviceConfig) -> crate::Result<(Device, DeviceMapMetadata)> {
    let _gpu_devices = generic_device_map.allocate_layers_to_gpus(0, 0)?;
    let main_gpu = generic_device_map.main_gpu()?;
    // let layer_count = gpu_devices
    //     .iter()
    //     .map(|d| d.allocated_layers as usize)
    //     .sum();

    Ok((
        Device::cuda_if_available(main_gpu as usize)?,
        DeviceMapMetadata::dummy(),
    ))
    // Ok((
    //     Device::cuda_if_available(main_gpu as usize)?,
    //     DeviceMapMetadata::from_num_device_layers(vec![DeviceLayerMapMetadata {
    //         ordinal: main_gpu as usize,
    //         layers: layer_count,
    //     }]),
    // ))
}

fn new_multiple_gpu(
    generic_device_map: &DeviceConfig,
) -> crate::Result<(Device, DeviceMapMetadata)> {
    let gpu_devices = generic_device_map.allocate_layers_to_gpus(0, 0)?;
    let main_gpu = generic_device_map.main_gpu()?;

    let mut device_map_metadata: Vec<DeviceLayerMapMetadata> = Vec::new();
    for gpu in &gpu_devices {
        device_map_metadata.push(DeviceLayerMapMetadata {
            ordinal: gpu.ordinal as usize,
            layers: gpu.allocated_layers as usize,
        });
    }
    Ok((
        Device::cuda_if_available(main_gpu as usize)?,
        DeviceMapMetadata::from_num_device_layers(device_map_metadata),
    ))
}