Function split_model

cake-core/src/utils/split.rs:155–234

Split a model into per-worker bundles. Each bundle contains a reduced safetensors file with only the worker's assigned tensors, a matching index file, and the worker's topology.

(
    model_path: &Path,
    topology_path: &str,
    worker: Option<&str>,
    output: &Path,
)

/// Each bundle contains a reduced safetensors file with only the worker's assigned tensors,
/// a matching index file, and the worker's topology.
pub fn split_model(
    model_path: &Path,
    topology_path: &str,
    worker: Option<&str>,
    output: &Path,
) -> Result<()> {
    let topology = Topology::from_path(topology_path, &ModelType::TextModel)?;
    let index = load_index(model_path)?;

    log::info!("index has {} tensors", index.weight_map.len());

    let selected_workers: Vec<String> = if let Some(name) = worker {
        vec![name.to_string()]
    } else {
        topology.keys().map(|s| s.to_string()).collect()
    };

    log::info!("processing {} workers", selected_workers.len());

    for worker_name in &selected_workers {
        log::info!("processing worker {worker_name} ...");

        let worker_node = topology
            .get(worker_name)
            .ok_or_else(|| anyhow!("can't find worker '{}' in topology", worker_name))?;

        let (new_index, reduced) = reduce_for_worker(&index, worker_node)?;

        log::info!("compacting {} tensors ...", new_index.weight_map.len());

        let metadata = create_new_metadata(model_path, &reduced)?;

        let bundle_name = format!("{worker_name}-node");
        let output_path = output.join(&bundle_name);
        let model_output_path = output_path.join("model");
        if !output_path.exists() {
            log::info!("creating {}", model_output_path.display());
            std::fs::create_dir_all(&model_output_path)?;
        } else {
            log::info!("saving model to {}", model_output_path.display());
        }

        let new_index_path = model_output_path.join("model.safetensors.index.json");

        log::info!("saving new index to {} ...", new_index_path.display());

        let new_index_data = serde_json::to_string_pretty(&new_index)?;
        std::fs::write(&new_index_path, new_index_data)?;

        let new_tensors_path = model_output_path.join("reduced.safetensors");

        log::info!(
            "saving reduced tensors to {} ...",
            new_tensors_path.display()
        );

        safetensors::serialize_to_file(metadata, None, &new_tensors_path)?;