| 1 | //! LoRA adapter merge module. |
| 2 | //! |
| 3 | //! Merges LoRA adapter weights (produced by the fine-tuner) back into the base |
| 4 | //! model safetensors, producing a single merged `model.safetensors` alongside |
| 5 | //! copied tokenizer and config files. |
| 6 | |
| 7 | use std::collections::HashMap; |
| 8 | use std::path::PathBuf; |
| 9 | |
| 10 | use anyhow::{Context, Result}; |
| 11 | use candle_core::safetensors::MmapedSafetensors; |
| 12 | use candle_core::{DType, Device, Tensor}; |
| 13 | |
| 14 | // --------------------------------------------------------------------------- |
| 15 | // Public types |
| 16 | // --------------------------------------------------------------------------- |
| 17 | |
| 18 | pub struct MergeConfig { |
| 19 | /// Snapshot directory containing `model.safetensors` (or sharded files) and |
| 20 | /// `config.json`. |
| 21 | pub base_dir: PathBuf, |
| 22 | /// Path to the `lora_adapter.safetensors` produced by the fine-tuner. |
| 23 | pub adapter_path: PathBuf, |
| 24 | /// Directory where the merged model and copied config/tokenizer files will |
| 25 | /// be written. |
| 26 | pub output_dir: PathBuf, |
| 27 | } |
| 28 | |
| 29 | pub enum MergeProgress { |
| 30 | Loading, |
| 31 | Merging { layer: usize, total: usize }, |
| 32 | Saving, |
| 33 | Done { output_path: PathBuf }, |
| 34 | Failed(String), |
| 35 | } |
| 36 | |
| 37 | // --------------------------------------------------------------------------- |
| 38 | // Public entry point |
| 39 | // --------------------------------------------------------------------------- |
| 40 | |
| 41 | pub fn start_merge(config: MergeConfig, tx: tokio::sync::mpsc::UnboundedSender<MergeProgress>) { |
| 42 | std::thread::spawn(move || { |
| 43 | if let Err(e) = run_merge(&config, &tx) { |
| 44 | eprintln!("[merge] error: {e:#}"); |
| 45 | let _ = tx.send(MergeProgress::Failed(format!("{e:#}"))); |
| 46 | } |
| 47 | }); |
| 48 | } |
| 49 | |
| 50 | // --------------------------------------------------------------------------- |
| 51 | // Private helpers |
| 52 | // --------------------------------------------------------------------------- |
| 53 | |
| 54 | #[derive(serde::Deserialize)] |
| 55 | struct MinimalConfig { |
| 56 | num_hidden_layers: usize, |
| 57 | } |
| 58 | |
| 59 | #[derive(serde::Deserialize)] |
| 60 | struct IndexJson { |
| 61 | weight_map: HashMap<String, String>, |
| 62 | } |
| 63 | |
| 64 | fn run_merge( |
| 65 | config: &MergeConfig, |
| 66 | tx: &tokio::sync::mpsc::UnboundedSender<MergeProgress>, |
| 67 | ) -> Result<()> { |
| 68 | let _ = tx.send(MergeProgress::Loading); |
| 69 | |
| 70 | // 1. Read num_hidden_layers from config.json |
| 71 | let config_path = config.base_dir.join("config.json"); |
| 72 | let config_text = std::fs::read_to_string(&config_path) |
| 73 | .with_context(|| format!("failed to read {}", config_path.display()))?; |
| 74 | let model_cfg: MinimalConfig = |
| 75 | serde_json::from_str(&config_text).context("failed to parse config.json")?; |
| 76 | let num_layers = model_cfg.num_hidden_layers; |
| 77 | |
| 78 | // 2. Load base safetensors (single file or sharded) |
| 79 | let index_path = config.base_dir.join("model.safetensors.index.json"); |
| 80 | let base_paths: Vec<PathBuf> = if index_path.exists() { |
| 81 | let index_text = std::fs::read_to_string(&index_path) |
| 82 | .with_context(|| format!("failed to read {}", index_path.display()))?; |
| 83 | let index: IndexJson = serde_json::from_str(&index_text) |
| 84 | .context("failed to parse model.safetensors.index.json")?; |
| 85 | let mut filenames: Vec<String> = index.weight_map.into_values().collect(); |
| 86 | filenames.sort(); |
| 87 | filenames.dedup(); |
| 88 | filenames |
| 89 | .into_iter() |
| 90 | .map(|f| config.base_dir.join(f)) |
| 91 | .collect() |
| 92 | } else { |
| 93 | vec![config.base_dir.join("model.safetensors")] |
| 94 | }; |
| 95 | |
| 96 | // SAFETY: the files are memory-mapped read-only. |
| 97 | let base = if base_paths.len() == 1 { |
| 98 | unsafe { MmapedSafetensors::new(&base_paths[0]) } |
| 99 | } else { |
| 100 | unsafe { MmapedSafetensors::multi(&base_paths) } |
| 101 | } |
| 102 | .context("failed to mmap base safetensors")?; |
| 103 | |
| 104 | // 3. Load adapter tensors |
| 105 | let adapter = |
| 106 | candle_core::safetensors::load(&config.adapter_path, &Device::Cpu).with_context(|| { |
| 107 | format!( |
| 108 | "failed to load adapter from {}", |
| 109 | config.adapter_path.display() |
| 110 | ) |
| 111 | })?; |
| 112 | |
| 113 | // 4. Infer rank from any lora_a tensor |
| 114 | let rank = adapter |
| 115 | .iter() |
| 116 | .find(|(k, _)| k.ends_with("lora_a")) |
| 117 | .map(|(_, t)| t.dim(0)) |
| 118 | .context("no lora_a tensor found in adapter")? |
| 119 | .context("failed to read lora_a dimension")?; |
| 120 | |
| 121 | // 5. Scale = alpha / rank = (rank * 2) / rank = 2.0 |
| 122 | let scale = 2.0_f64; |
| 123 | |
| 124 | // 6. Iterate base tensors and merge where adapters exist |
| 125 | let base_tensor_info = base.tensors(); |
| 126 | let mut merged_tensors: HashMap<String, Tensor> = HashMap::new(); |
| 127 | let mut merge_layer_idx: usize = 0; |
| 128 | |
| 129 | // Count how many layers actually have adapters for progress reporting |
| 130 | let total_merge_layers = { |
| 131 | let mut count = 0usize; |
| 132 | for i in 0..num_layers { |
| 133 | let q_a = format!("model.layers.{i}.self_attn.q_proj.lora_a"); |
| 134 | let v_a = format!("model.layers.{i}.self_attn.v_proj.lora_a"); |
| 135 | if adapter.contains_key(&q_a) || adapter.contains_key(&v_a) { |
| 136 | count += 1; |
| 137 | } |
| 138 | } |
| 139 | count |
| 140 | }; |
| 141 | |
| 142 | // Track which layers we've already sent progress for |
| 143 | let mut reported_layers: std::collections::HashSet<usize> = std::collections::HashSet::new(); |
| 144 | |
| 145 | for (name, _tv) in &base_tensor_info { |
| 146 | let base_tensor = base |
| 147 | .load(name, &Device::Cpu) |
| 148 | .with_context(|| format!("failed to load base tensor {name}"))?; |
| 149 | let base_dtype = base_tensor.dtype(); |
| 150 | |
| 151 | let stem = name.strip_suffix(".weight").unwrap_or(name); |
| 152 | let lora_a_key = format!("{stem}.lora_a"); |
| 153 | let lora_b_key = format!("{stem}.lora_b"); |
| 154 | |
| 155 | let tensor = if let (Some(lora_a), Some(lora_b)) = |
| 156 | (adapter.get(&lora_a_key), adapter.get(&lora_b_key)) |
| 157 | { |
| 158 | // Report progress per-layer |
| 159 | // Extract layer index from the stem if possible |
| 160 | if let Some(layer_num) = extract_layer_index(stem) |
| 161 | && reported_layers.insert(layer_num) |
| 162 | { |
| 163 | merge_layer_idx += 1; |
| 164 | let _ = tx.send(MergeProgress::Merging { |
| 165 | layer: merge_layer_idx, |
| 166 | total: total_merge_layers, |
| 167 | }); |
| 168 | } |
| 169 | |
| 170 | let base_f32 = base_tensor.to_dtype(DType::F32)?; |
| 171 | let a_f32 = lora_a.to_dtype(DType::F32)?; |
| 172 | let b_f32 = lora_b.to_dtype(DType::F32)?; |
| 173 | let delta = (b_f32.matmul(&a_f32)? * scale)?; |
| 174 | let merged = (base_f32 + delta)?; |
| 175 | merged.to_dtype(base_dtype)? |
| 176 | } else { |
| 177 | base_tensor |
| 178 | }; |
| 179 | |
| 180 | merged_tensors.insert(name.clone(), tensor); |
| 181 | } |
| 182 | |
| 183 | // Ensure rank is used (suppress unused warning via a debug log) |
| 184 | let _ = rank; |
| 185 | |
| 186 | // 7. Save merged model |
| 187 | let _ = tx.send(MergeProgress::Saving); |
| 188 | std::fs::create_dir_all(&config.output_dir).with_context(|| { |
| 189 | format!( |
| 190 | "failed to create output dir {}", |
| 191 | config.output_dir.display() |
| 192 | ) |
| 193 | })?; |
| 194 | |
| 195 | let output_model = config.output_dir.join("model.safetensors"); |
| 196 | candle_core::safetensors::save(&merged_tensors, &output_model) |
| 197 | .with_context(|| format!("failed to save merged model to {}", output_model.display()))?; |
| 198 | |
| 199 | // 8. Copy config/tokenizer files (skip missing) |
| 200 | let copy_files = [ |
| 201 | "config.json", |
| 202 | "tokenizer.json", |
| 203 | "tokenizer_config.json", |
| 204 | "generation_config.json", |
| 205 | ]; |
| 206 | for filename in ©_files { |
| 207 | let src = config.base_dir.join(filename); |
| 208 | if src.exists() { |
| 209 | let dst = config.output_dir.join(filename); |
| 210 | if let Err(e) = std::fs::copy(&src, &dst) { |
| 211 | eprintln!("[merge] warning: failed to copy {}: {e:#}", src.display()); |
| 212 | } |
| 213 | } |
| 214 | } |
| 215 | |
| 216 | // 9. Done |
| 217 | let _ = tx.send(MergeProgress::Done { |
| 218 | output_path: output_model, |
| 219 | }); |
| 220 | Ok(()) |
| 221 | } |
| 222 | |
| 223 | /// Extract a numeric layer index from a tensor name stem such as |
| 224 | /// `model.layers.5.self_attn.q_proj`. |
| 225 | fn extract_layer_index(stem: &str) -> Option<usize> { |
| 226 | let parts: Vec<&str> = stem.split('.').collect(); |
| 227 | for (i, part) in parts.iter().enumerate() { |
| 228 | if *part == "layers" { |
| 229 | return parts.get(i + 1).and_then(|s| s.parse::<usize>().ok()); |
| 230 | } |
| 231 | } |
| 232 | None |
| 233 | } |