| 1 | //! LoRA fine-tuning for Qwen-family models (Qwen2, Qwen2.5, Qwen3). |
| 2 | //! |
| 3 | //! Loads a safetensors base model from a HuggingFace cache directory, injects |
| 4 | //! trainable LoRA adapters into the Q and V projection layers of every |
| 5 | //! transformer block, trains for a configurable number of epochs on a JSONL |
| 6 | //! dataset, and saves the resulting adapter weights to a safetensors file. |
| 7 | |
| 8 | use std::collections::{HashMap, HashSet}; |
| 9 | use std::io::{BufRead, BufReader}; |
| 10 | use std::path::{Path, PathBuf}; |
| 11 | |
| 12 | use anyhow::{Context, Result}; |
| 13 | use candle_core::{D, DType, Device, Tensor, Var}; |
| 14 | use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder}; |
| 15 | use tokenizers::Tokenizer; |
| 16 | |
| 17 | // --------------------------------------------------------------------------- |
| 18 | // Public types |
| 19 | // --------------------------------------------------------------------------- |
| 20 | |
| 21 | /// Configuration for a LoRA fine-tuning run. |
| 22 | pub struct FineTuneConfig { |
| 23 | /// Directory that contains `config.json`, `tokenizer.json`, and the |
| 24 | /// safetensors weight files (single-file or sharded). |
| 25 | pub model_dir: PathBuf, |
| 26 | /// JSONL file where every line is `{"text": "…"}`. |
| 27 | pub data_path: PathBuf, |
| 28 | /// Directory where `lora_adapter.safetensors` will be written. |
| 29 | pub output_dir: PathBuf, |
| 30 | /// LoRA rank (default 8). |
| 31 | pub lora_rank: usize, |
| 32 | /// LoRA alpha, used to compute the scaling factor `alpha / rank` (default 16.0). |
| 33 | pub lora_alpha: f32, |
| 34 | /// Learning rate for AdamW (default 1e-4). |
| 35 | pub learning_rate: f64, |
| 36 | /// Number of full passes over the dataset (default 3). |
| 37 | pub epochs: usize, |
| 38 | /// Maximum number of tokens per training example (default 512). |
| 39 | pub max_seq_len: usize, |
| 40 | } |
| 41 | |
| 42 | /// Progress events emitted by the background fine-tuning thread. |
| 43 | pub enum FineTuneProgress { |
| 44 | /// Checking that the model directory and data file exist. |
| 45 | Validating, |
| 46 | /// Loading and memory-mapping the model weights. |
| 47 | LoadingModel, |
| 48 | /// Tokenizing the dataset. |
| 49 | Tokenizing { done: usize, total: usize }, |
| 50 | /// One optimizer step has completed. |
| 51 | Training { |
| 52 | epoch: usize, |
| 53 | total_epochs: usize, |
| 54 | step: usize, |
| 55 | total_steps: usize, |
| 56 | loss: f32, |
| 57 | }, |
| 58 | /// Writing the adapter file to disk. |
| 59 | Saving, |
| 60 | /// Training finished successfully. |
| 61 | Done { adapter_path: PathBuf }, |
| 62 | /// An unrecoverable error occurred. |
| 63 | Failed(String), |
| 64 | } |
| 65 | |
| 66 | // --------------------------------------------------------------------------- |
| 67 | // Public API |
| 68 | // --------------------------------------------------------------------------- |
| 69 | |
| 70 | /// Spawn a background OS thread that runs the full fine-tuning pipeline and |
| 71 | /// streams [`FineTuneProgress`] events over `tx`. |
| 72 | /// |
| 73 | /// The thread is intentionally a blocking `std::thread` (not a tokio task) so |
| 74 | /// that heavy tensor operations do not starve the async runtime. |
| 75 | pub fn start_finetune( |
| 76 | config: FineTuneConfig, |
| 77 | tx: tokio::sync::mpsc::UnboundedSender<FineTuneProgress>, |
| 78 | ) { |
| 79 | std::thread::spawn(move || { |
| 80 | if let Err(e) = run_finetune(&config, &tx) { |
| 81 | eprintln!("[finetune] error: {e:#}"); |
| 82 | let _ = tx.send(FineTuneProgress::Failed(format!("{e:#}"))); |
| 83 | } |
| 84 | }); |
| 85 | } |
| 86 | |
| 87 | // --------------------------------------------------------------------------- |
| 88 | // Private: top-level orchestration |
| 89 | // --------------------------------------------------------------------------- |
| 90 | |
| 91 | fn run_finetune( |
| 92 | config: &FineTuneConfig, |
| 93 | tx: &tokio::sync::mpsc::UnboundedSender<FineTuneProgress>, |
| 94 | ) -> Result<()> { |
| 95 | // ------------------------------------------------------------------ |
| 96 | // Step A — Validate inputs |
| 97 | // ------------------------------------------------------------------ |
| 98 | let _ = tx.send(FineTuneProgress::Validating); |
| 99 | |
| 100 | let model_dir = &config.model_dir; |
| 101 | let config_json_path = model_dir.join("config.json"); |
| 102 | |
| 103 | if !config_json_path.exists() { |
| 104 | anyhow::bail!( |
| 105 | "model config not found at {:?}; is this a valid HuggingFace model directory?", |
| 106 | config_json_path |
| 107 | ); |
| 108 | } |
| 109 | if !config.data_path.exists() { |
| 110 | anyhow::bail!("training data file not found at {:?}", config.data_path); |
| 111 | } |
| 112 | |
| 113 | // ------------------------------------------------------------------ |
| 114 | // Step B — Load model config |
| 115 | // ------------------------------------------------------------------ |
| 116 | let config_text = std::fs::read_to_string(&config_json_path) |
| 117 | .with_context(|| format!("reading {:?}", config_json_path))?; |
| 118 | let model_cfg: ModelConfig = serde_json::from_str(&config_text) |
| 119 | .with_context(|| format!("parsing {:?}", config_json_path))?; |
| 120 | |
| 121 | // ------------------------------------------------------------------ |
| 122 | // Step C — Choose compute device |
| 123 | // ------------------------------------------------------------------ |
| 124 | // Training runs on CPU even on macOS: candle 0.10.2's Metal *backward* |
| 125 | // pass produces garbage gradient magnitudes (1e6–1e29 at pristine weights, |
| 126 | // where the CPU backward yields norms of ~40–65 for identical losses), and |
| 127 | // with Accelerate the CPU backward is also ~2x faster per step. Metal |
| 128 | // remains opt-in via ONDE_FINETUNE_METAL for re-testing on newer candle. |
| 129 | // Forward-only paths (inference, merge, export) are unaffected. |
| 130 | #[cfg(target_os = "macos")] |
| 131 | let device = if std::env::var("ONDE_FINETUNE_METAL").is_ok() { |
| 132 | Device::new_metal(0).unwrap_or(Device::Cpu) |
| 133 | } else { |
| 134 | Device::Cpu |
| 135 | }; |
| 136 | #[cfg(not(target_os = "macos"))] |
| 137 | let device = Device::Cpu; |
| 138 | |
| 139 | // ------------------------------------------------------------------ |
| 140 | // Step D — Load frozen base weights |
| 141 | // ------------------------------------------------------------------ |
| 142 | let _ = tx.send(FineTuneProgress::LoadingModel); |
| 143 | |
| 144 | let index_path = model_dir.join("model.safetensors.index.json"); |
| 145 | let vb = if index_path.exists() { |
| 146 | // Sharded model: read the weight-map JSON and collect unique shard paths. |
| 147 | let index_text = std::fs::read_to_string(&index_path) |
| 148 | .with_context(|| format!("reading {:?}", index_path))?; |
| 149 | let index: IndexJson = serde_json::from_str(&index_text) |
| 150 | .with_context(|| format!("parsing {:?}", index_path))?; |
| 151 | |
| 152 | let mut seen: HashSet<String> = HashSet::new(); |
| 153 | let mut shard_paths: Vec<PathBuf> = Vec::new(); |
| 154 | for filename in index.weight_map.values() { |
| 155 | if seen.insert(filename.clone()) { |
| 156 | shard_paths.push(model_dir.join(filename)); |
| 157 | } |
| 158 | } |
| 159 | shard_paths.sort(); |
| 160 | |
| 161 | // SAFETY: the files are memory-mapped read-only. |
| 162 | unsafe { |
| 163 | VarBuilder::from_mmaped_safetensors(&shard_paths, DType::F32, &device) |
| 164 | .context("loading sharded safetensors")? |
| 165 | } |
| 166 | } else { |
| 167 | let single = model_dir.join("model.safetensors"); |
| 168 | // SAFETY: the file is memory-mapped read-only. |
| 169 | unsafe { |
| 170 | VarBuilder::from_mmaped_safetensors(&[single], DType::F32, &device) |
| 171 | .context("loading model.safetensors")? |
| 172 | } |
| 173 | }; |
| 174 | |
| 175 | // ------------------------------------------------------------------ |
| 176 | // Step E — Build LoRA model |
| 177 | // ------------------------------------------------------------------ |
| 178 | let model = LoraQwenModel::load(vb, &model_cfg, config.lora_rank, config.lora_alpha, &device) |
| 179 | .context("building LoRA model")?; |
| 180 | |
| 181 | // ------------------------------------------------------------------ |
| 182 | // Step F — Load and tokenize dataset |
| 183 | // ------------------------------------------------------------------ |
| 184 | let tokenizer_path = model_dir.join("tokenizer.json"); |
| 185 | let tokenizer = Tokenizer::from_file(&tokenizer_path) |
| 186 | .map_err(|e| anyhow::anyhow!("loading tokenizer from {:?}: {e}", tokenizer_path))?; |
| 187 | |
| 188 | let file = std::fs::File::open(&config.data_path) |
| 189 | .with_context(|| format!("opening {:?}", config.data_path))?; |
| 190 | let reader = BufReader::new(file); |
| 191 | let raw_lines: Vec<String> = reader |
| 192 | .lines() |
| 193 | .collect::<std::io::Result<Vec<_>>>() |
| 194 | .context("reading training data")?; |
| 195 | |
| 196 | let total_lines = raw_lines.len(); |
| 197 | |
| 198 | struct TrainingExample { |
| 199 | tokens: Vec<u32>, |
| 200 | mask: Vec<bool>, |
| 201 | } |
| 202 | |
| 203 | let mut examples: Vec<TrainingExample> = Vec::new(); |
| 204 | |
| 205 | for (i, line) in raw_lines.iter().enumerate() { |
| 206 | let line = line.trim(); |
| 207 | if line.is_empty() { |
| 208 | let _ = tx.send(FineTuneProgress::Tokenizing { |
| 209 | done: i + 1, |
| 210 | total: total_lines, |
| 211 | }); |
| 212 | continue; |
| 213 | } |
| 214 | |
| 215 | let entry: DataEntry = |
| 216 | serde_json::from_str(line).with_context(|| format!("parsing JSONL line {}", i + 1))?; |
| 217 | |
| 218 | let (rendered_text, assistant_spans) = match &entry { |
| 219 | DataEntry::Raw { text } => (text.clone(), Vec::new()), |
| 220 | DataEntry::Messages { messages } => render_chatml(messages), |
| 221 | }; |
| 222 | |
| 223 | let encoding = tokenizer |
| 224 | .encode(rendered_text.as_str(), false) |
| 225 | .map_err(|e| anyhow::anyhow!("tokenizing line {}: {e}", i + 1))?; |
| 226 | |
| 227 | let ids = encoding.get_ids(); |
| 228 | let offsets = encoding.get_offsets(); |
| 229 | |
| 230 | // Truncate ids and offsets together to max_seq_len |
| 231 | let (ids, offsets) = if ids.len() > config.max_seq_len { |
| 232 | ( |
| 233 | ids[..config.max_seq_len].to_vec(), |
| 234 | offsets[..config.max_seq_len].to_vec(), |
| 235 | ) |
| 236 | } else { |
| 237 | (ids.to_vec(), offsets.to_vec()) |
| 238 | }; |
| 239 | |
| 240 | // Need at least 2 tokens for next-token prediction. |
| 241 | if ids.len() < 2 { |
| 242 | eprintln!( |
| 243 | "[finetune] WARNING: skipping line {} (too short: {} tokens)", |
| 244 | i + 1, |
| 245 | ids.len() |
| 246 | ); |
| 247 | let _ = tx.send(FineTuneProgress::Tokenizing { |
| 248 | done: i + 1, |
| 249 | total: total_lines, |
| 250 | }); |
| 251 | continue; |
| 252 | } |
| 253 | |
| 254 | let mask = if assistant_spans.is_empty() { |
| 255 | // Raw format: full-sequence loss (backward compatible) |
| 256 | vec![true; ids.len()] |
| 257 | } else { |
| 258 | // Messages format: mask only assistant completions |
| 259 | compute_loss_mask(&offsets, &assistant_spans) |
| 260 | }; |
| 261 | |
| 262 | // For Messages entries, check that the shifted mask has at least one true |
| 263 | if !assistant_spans.is_empty() && mask.len() > 1 { |
| 264 | let shifted_mask = &mask[1..]; |
| 265 | if !shifted_mask.iter().any(|&m| m) { |
| 266 | eprintln!( |
| 267 | "[finetune] WARNING: skipping line {} (assistant content truncated away)", |
| 268 | i + 1 |
| 269 | ); |
| 270 | let _ = tx.send(FineTuneProgress::Tokenizing { |
| 271 | done: i + 1, |
| 272 | total: total_lines, |
| 273 | }); |
| 274 | continue; |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | examples.push(TrainingExample { tokens: ids, mask }); |
| 279 | |
| 280 | let _ = tx.send(FineTuneProgress::Tokenizing { |
| 281 | done: i + 1, |
| 282 | total: total_lines, |
| 283 | }); |
| 284 | } |
| 285 | |
| 286 | if examples.is_empty() { |
| 287 | anyhow::bail!("no valid training examples found in {:?}", config.data_path); |
| 288 | } |
| 289 | |
| 290 | // ------------------------------------------------------------------ |
| 291 | // Step G — Training loop |
| 292 | // ------------------------------------------------------------------ |
| 293 | let lora_vars = model.lora_vars(); |
| 294 | let adamw_params = ParamsAdamW { |
| 295 | lr: config.learning_rate, |
| 296 | ..ParamsAdamW::default() |
| 297 | }; |
| 298 | let mut optimizer = |
| 299 | AdamW::new(lora_vars.clone(), adamw_params).context("creating AdamW optimizer")?; |
| 300 | |
| 301 | let total_steps = examples.len(); |
| 302 | let vocab_size = model_cfg.vocab_size; |
| 303 | let max_grad_norm: f64 = 1.0; |
| 304 | |
| 305 | for epoch in 0..config.epochs { |
| 306 | for (step, example) in examples.iter().enumerate() { |
| 307 | let seq_len = example.tokens.len(); |
| 308 | |
| 309 | let input_ids = Tensor::from_slice(example.tokens.as_slice(), (1, seq_len), &device) |
| 310 | .context("building input_ids tensor")?; |
| 311 | |
| 312 | let logits = model.forward(&input_ids).context("model forward pass")?; |
| 313 | |
| 314 | // Next-token prediction: shift logits and targets by 1. |
| 315 | let logits_shifted = logits |
| 316 | .narrow(1, 0, seq_len - 1) |
| 317 | .and_then(|t| t.reshape((seq_len - 1, vocab_size))) |
| 318 | .context("preparing shifted logits")?; |
| 319 | |
| 320 | let targets = Tensor::from_slice(&example.tokens[1..], (seq_len - 1,), &device) |
| 321 | .context("building target tensor")?; |
| 322 | |
| 323 | let shifted_mask = &example.mask[1..]; |
| 324 | |
| 325 | // Fast path: if all tokens are active (typical for Raw entries), skip index_select |
| 326 | let (logits_masked, targets_masked) = if shifted_mask.iter().all(|&m| m) { |
| 327 | (logits_shifted.clone(), targets.clone()) |
| 328 | } else { |
| 329 | // Masked path: keep only the positions where mask is true |
| 330 | let keep: Vec<u32> = shifted_mask |
| 331 | .iter() |
| 332 | .enumerate() |
| 333 | .filter_map(|(i, &m)| if m { Some(i as u32) } else { None }) |
| 334 | .collect(); |
| 335 | |
| 336 | if keep.is_empty() { |
| 337 | eprintln!( |
| 338 | "[finetune] WARNING: no masked tokens at epoch {} step {}, skipping", |
| 339 | epoch + 1, |
| 340 | step + 1 |
| 341 | ); |
| 342 | let _ = tx.send(FineTuneProgress::Training { |
| 343 | epoch: epoch + 1, |
| 344 | total_epochs: config.epochs, |
| 345 | step: step + 1, |
| 346 | total_steps, |
| 347 | loss: 0.0, |
| 348 | }); |
| 349 | continue; |
| 350 | } |
| 351 | |
| 352 | let keep_idx = Tensor::from_slice(keep.as_slice(), (keep.len(),), &device) |
| 353 | .context("building keep indices")?; |
| 354 | let logits_masked = logits_shifted |
| 355 | .index_select(&keep_idx, 0) |
| 356 | .context("indexing logits")?; |
| 357 | let targets_masked = targets |
| 358 | .index_select(&keep_idx, 0) |
| 359 | .context("indexing targets")?; |
| 360 | (logits_masked, targets_masked) |
| 361 | }; |
| 362 | |
| 363 | let loss = candle_nn::loss::cross_entropy(&logits_masked, &targets_masked) |
| 364 | .context("computing cross-entropy loss")?; |
| 365 | |
| 366 | let loss_val = loss.to_scalar::<f32>().context("reading loss scalar")?; |
| 367 | |
| 368 | // Guard against NaN/Inf loss — skip this step to avoid poisoning |
| 369 | // the optimizer state. |
| 370 | if loss_val.is_nan() || loss_val.is_infinite() { |
| 371 | eprintln!( |
| 372 | "[finetune] WARNING: loss is {} at epoch {} step {}, skipping", |
| 373 | loss_val, |
| 374 | epoch + 1, |
| 375 | step + 1 |
| 376 | ); |
| 377 | let _ = tx.send(FineTuneProgress::Training { |
| 378 | epoch: epoch + 1, |
| 379 | total_epochs: config.epochs, |
| 380 | step: step + 1, |
| 381 | total_steps, |
| 382 | loss: loss_val, |
| 383 | }); |
| 384 | continue; |
| 385 | } |
| 386 | |
| 387 | // Manual backward + gradient clipping + optimizer step. |
| 388 | // |
| 389 | // We avoid `backward_step` so we can clip gradients between the |
| 390 | // backward pass and the AdamW update. This preserves AdamW's |
| 391 | // momentum / velocity state (unlike manually applying lr * grad). |
| 392 | let mut grads = loss.backward().context("backward pass")?; |
| 393 | |
| 394 | // Sanitize gradients before clipping. Backprop through the deep |
| 395 | // stack can produce non-finite (NaN/Inf) gradient *elements* even |
| 396 | // when the forward loss is finite; a single such element makes the |
| 397 | // whole tensor's L2 norm non-finite, which would defeat norm-based |
| 398 | // clipping (NaN > threshold is false → unclipped) and poison the |
| 399 | // AdamW moment estimates, corrupting every weight from that step on. |
| 400 | // Replace non-finite elements with zero so the remaining finite |
| 401 | // gradient still trains those parameters this step. |
| 402 | // Pass 1 — sanitize each gradient and find the largest |g| element |
| 403 | // across all tensors. sanitize_grad admits elements up to 1e30, but |
| 404 | // squaring anything above ~1.8e19 overflows the f32 sum below, so the |
| 405 | // norm must be computed in a scaled space. |
| 406 | let mut cleaned: Vec<(&Var, Tensor)> = Vec::with_capacity(lora_vars.len()); |
| 407 | let mut global_max_abs: f64 = 0.0; |
| 408 | for var in &lora_vars { |
| 409 | if let Some(grad) = grads.get(var.as_tensor()) { |
| 410 | let clean = sanitize_grad(grad).context("sanitizing gradient")?; |
| 411 | let max_abs = clean |
| 412 | .abs() |
| 413 | .and_then(|t| t.max_all()) |
| 414 | .and_then(|t| t.to_scalar::<f32>()) |
| 415 | .unwrap_or(0.0) as f64; |
| 416 | global_max_abs = global_max_abs.max(max_abs); |
| 417 | cleaned.push((var, clean)); |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | // Pass 2 — L2 norm of the gradients scaled by 1/global_max_abs. |
| 422 | // Scaled elements are ≤ 1 so the f32 accumulation cannot overflow; |
| 423 | // the true norm is recovered in f64. Exploding-but-finite gradients |
| 424 | // thus get a finite norm and are tamed by clipping below instead of |
| 425 | // forcing the whole step to be skipped. |
| 426 | let mut scaled_norm_sq: f64 = 0.0; |
| 427 | if global_max_abs > 0.0 { |
| 428 | let inv = 1.0 / global_max_abs; |
| 429 | for (_, clean) in &cleaned { |
| 430 | let norm_sq = (clean * inv) |
| 431 | .and_then(|t| t.sqr()) |
| 432 | .and_then(|t| t.sum_all()) |
| 433 | .and_then(|t| t.to_scalar::<f32>()) |
| 434 | .unwrap_or(0.0) as f64; |
| 435 | scaled_norm_sq += norm_sq; |
| 436 | } |
| 437 | } |
| 438 | let total_norm = global_max_abs * scaled_norm_sq.sqrt(); |
| 439 | if std::env::var("ONDE_FINETUNE_GRAD_DEBUG").is_ok() { |
| 440 | eprintln!( |
| 441 | "[finetune] grad-debug epoch {} step {} seq_len {} loss {:.4} max_abs {:.3e} norm {:.3e}", |
| 442 | epoch + 1, |
| 443 | step + 1, |
| 444 | seq_len, |
| 445 | loss_val, |
| 446 | global_max_abs, |
| 447 | total_norm |
| 448 | ); |
| 449 | } |
| 450 | for (var, clean) in cleaned { |
| 451 | grads.insert(var.as_tensor(), clean); |
| 452 | } |
| 453 | |
| 454 | // Nothing finite to learn from this step — skip the update so AdamW |
| 455 | // state stays clean. |
| 456 | if total_norm == 0.0 || !total_norm.is_finite() { |
| 457 | eprintln!( |
| 458 | "[finetune] WARNING: global grad norm is {} at epoch {} step {} (seq_len {}), skipping update", |
| 459 | total_norm, |
| 460 | epoch + 1, |
| 461 | step + 1, |
| 462 | seq_len |
| 463 | ); |
| 464 | continue; |
| 465 | } |
| 466 | |
| 467 | // If the global norm exceeds the threshold, scale every gradient |
| 468 | // down and insert the clipped version back into the GradStore. |
| 469 | if total_norm > max_grad_norm { |
| 470 | let clip_coef = max_grad_norm / (total_norm + 1e-6); |
| 471 | for var in &lora_vars { |
| 472 | if let Some(grad) = grads.get(var.as_tensor()) { |
| 473 | let clipped = (grad * clip_coef).context("clipping gradient")?; |
| 474 | grads.insert(var.as_tensor(), clipped); |
| 475 | } |
| 476 | } |
| 477 | } |
| 478 | |
| 479 | // AdamW step with sanitized, clipped gradients — preserves |
| 480 | // first-moment and second-moment estimates. |
| 481 | optimizer.step(&grads).context("optimizer step")?; |
| 482 | |
| 483 | let _ = tx.send(FineTuneProgress::Training { |
| 484 | epoch: epoch + 1, |
| 485 | total_epochs: config.epochs, |
| 486 | step: step + 1, |
| 487 | total_steps, |
| 488 | loss: loss_val, |
| 489 | }); |
| 490 | } |
| 491 | } |
| 492 | |
| 493 | // ------------------------------------------------------------------ |
| 494 | // Step H — Save LoRA adapter |
| 495 | // ------------------------------------------------------------------ |
| 496 | let _ = tx.send(FineTuneProgress::Saving); |
| 497 | |
| 498 | let adapter_path = model |
| 499 | .save_lora(&config.output_dir) |
| 500 | .context("saving LoRA adapter")?; |
| 501 | |
| 502 | let _ = tx.send(FineTuneProgress::Done { adapter_path }); |
| 503 | |
| 504 | Ok(()) |
| 505 | } |
| 506 | |
| 507 | // --------------------------------------------------------------------------- |
| 508 | // Private: serde types for JSON deserialization |
| 509 | // --------------------------------------------------------------------------- |
| 510 | |
| 511 | #[derive(serde::Deserialize)] |
| 512 | #[allow(dead_code)] |
| 513 | struct ModelConfig { |
| 514 | hidden_size: usize, |
| 515 | num_hidden_layers: usize, |
| 516 | num_attention_heads: usize, |
| 517 | num_key_value_heads: Option<usize>, |
| 518 | head_dim: Option<usize>, |
| 519 | intermediate_size: usize, |
| 520 | vocab_size: usize, |
| 521 | rms_norm_eps: Option<f64>, |
| 522 | rope_theta: Option<f64>, |
| 523 | #[serde(default)] |
| 524 | tie_word_embeddings: bool, |
| 525 | } |
| 526 | |
| 527 | #[derive(serde::Deserialize)] |
| 528 | #[serde(untagged)] |
| 529 | enum DataEntry { |
| 530 | Messages { messages: Vec<DatasetMessage> }, |
| 531 | Raw { text: String }, |
| 532 | } |
| 533 | |
| 534 | #[derive(serde::Deserialize)] |
| 535 | struct DatasetMessage { |
| 536 | role: String, |
| 537 | content: String, |
| 538 | } |
| 539 | |
| 540 | /// Renders `messages` as literal Qwen ChatML, matching the format written by |
| 541 | /// hand in existing `Raw` datasets. Returns the rendered string plus byte-offset |
| 542 | /// spans of every assistant turn (content start through the end of that turn's |
| 543 | /// closing `<|im_end|>`, which is included so the model learns to emit it; the |
| 544 | /// trailing `\n` is not) for loss-mask computation. |
| 545 | fn render_chatml(messages: &[DatasetMessage]) -> (String, Vec<(usize, usize)>) { |
| 546 | let mut output = String::new(); |
| 547 | let mut assistant_spans = Vec::new(); |
| 548 | |
| 549 | for msg in messages { |
| 550 | output.push_str("<|im_start|>"); |
| 551 | output.push_str(&msg.role); |
| 552 | output.push('\n'); |
| 553 | |
| 554 | let content_start = output.len(); |
| 555 | output.push_str(&msg.content); |
| 556 | output.push_str("<|im_end|>"); |
| 557 | |
| 558 | if msg.role == "assistant" { |
| 559 | // Span end is *after* `<|im_end|>` so the closing marker is masked |
| 560 | // in for loss: the model has to learn to emit it, otherwise |
| 561 | // generations never terminate. Whether the marker tokenizes as one |
| 562 | // special token or splits into several is irrelevant here — every |
| 563 | // resulting token's byte offset falls inside this span. |
| 564 | let span_end = output.len(); |
| 565 | assistant_spans.push((content_start, span_end)); |
| 566 | } |
| 567 | |
| 568 | output.push('\n'); |
| 569 | } |
| 570 | |
| 571 | (output, assistant_spans) |
| 572 | } |
| 573 | |
| 574 | /// Pure function: true where a token's byte-offset range overlaps any |
| 575 | /// assistant span. Decoupled from tokenizer types so it's unit-testable |
| 576 | /// with synthetic data. |
| 577 | fn compute_loss_mask(offsets: &[(usize, usize)], assistant_spans: &[(usize, usize)]) -> Vec<bool> { |
| 578 | offsets |
| 579 | .iter() |
| 580 | .map(|&(token_start, token_end)| { |
| 581 | assistant_spans |
| 582 | .iter() |
| 583 | .any(|&(span_start, span_end)| token_start < span_end && token_end > span_start) |
| 584 | }) |
| 585 | .collect() |
| 586 | } |
| 587 | |
| 588 | /// Top-level structure of `model.safetensors.index.json`. |
| 589 | #[derive(serde::Deserialize)] |
| 590 | struct IndexJson { |
| 591 | weight_map: HashMap<String, String>, |
| 592 | } |
| 593 | |
| 594 | // --------------------------------------------------------------------------- |
| 595 | // Private: tensor helpers |
| 596 | // --------------------------------------------------------------------------- |
| 597 | |
| 598 | /// Replace every non-finite element (NaN, ±Inf) of a gradient tensor with zero. |
| 599 | /// |
| 600 | /// Finite values satisfy `|x| <= 1e30` (well below `f32::MAX ≈ 3.4e38`); NaN |
| 601 | /// fails the comparison (NaN compares false) and ±Inf exceeds the bound, so |
| 602 | /// both map to zero while genuine gradients pass through unchanged. |
| 603 | fn sanitize_grad(grad: &Tensor) -> candle_core::Result<Tensor> { |
| 604 | let finite = grad.abs()?.le(1e30f64)?; |
| 605 | let zeros = grad.zeros_like()?; |
| 606 | finite.where_cond(grad, &zeros) |
| 607 | } |
| 608 | |
| 609 | /// RMS layer normalisation: `x / rms(x) * weight`. |
| 610 | /// |
| 611 | /// `weight` has shape `[hidden_size]`; `x` has shape `[batch, seq, hidden_size]`. |
| 612 | fn rms_norm(x: &Tensor, weight: &Tensor, eps: f64) -> candle_core::Result<Tensor> { |
| 613 | // mean of squares across the last (hidden) dimension → [batch, seq, 1] |
| 614 | let mean_sq = x.sqr()?.mean_keepdim(D::Minus1)?; |
| 615 | // x / sqrt(mean_sq + eps) |
| 616 | let normed = x.broadcast_div(&mean_sq.affine(1.0, eps)?.sqrt()?)?; |
| 617 | // scale by the learned weight |
| 618 | normed.broadcast_mul(weight) |
| 619 | } |
| 620 | |
| 621 | /// Apply rotary position embeddings to `x`. |
| 622 | /// |
| 623 | /// `x` has shape `[batch, heads, seq, head_dim]`. |
| 624 | /// `cos` has shape `[seq, head_dim]` (already narrowed to the current sequence length). |
| 625 | /// `sin` has shape `[seq, head_dim]`. |
| 626 | fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> candle_core::Result<Tensor> { |
| 627 | let head_dim = x.dim(D::Minus1)?; |
| 628 | let half = head_dim / 2; |
| 629 | |
| 630 | let x1 = x.narrow(D::Minus1, 0, half)?; // first half |
| 631 | let x2 = x.narrow(D::Minus1, half, half)?; // second half |
| 632 | |
| 633 | // rotate_half: [-x2, x1] |
| 634 | let neg_x2 = x2.neg()?; |
| 635 | let rotated = Tensor::cat(&[&neg_x2, &x1], D::Minus1)?; |
| 636 | |
| 637 | // Unsqueeze cos/sin to [1, 1, seq, head_dim] for broadcasting. |
| 638 | let cos = cos.unsqueeze(0)?.unsqueeze(0)?; |
| 639 | let sin = sin.unsqueeze(0)?.unsqueeze(0)?; |
| 640 | |
| 641 | x.broadcast_mul(&cos)? + rotated.broadcast_mul(&sin)? |
| 642 | } |
| 643 | |
| 644 | /// Precompute cosine and sine tables for RoPE. |
| 645 | /// |
| 646 | /// Returns `(cos, sin)` each with shape `[max_seq_len, head_dim]`. |
| 647 | fn precompute_rope( |
| 648 | head_dim: usize, |
| 649 | max_seq_len: usize, |
| 650 | theta: f64, |
| 651 | device: &Device, |
| 652 | ) -> candle_core::Result<(Tensor, Tensor)> { |
| 653 | let half = head_dim / 2; |
| 654 | |
| 655 | // Inverse frequencies: 1 / theta^(2i / head_dim) for i in 0..half |
| 656 | let freqs: Vec<f32> = (0..half) |
| 657 | .map(|i| 1.0f32 / (theta as f32).powf(2.0 * i as f32 / head_dim as f32)) |
| 658 | .collect(); |
| 659 | |
| 660 | // freqs: [1, half] |
| 661 | let freqs = Tensor::from_slice(freqs.as_slice(), (1, half), device)?; |
| 662 | |
| 663 | // position indices: [max_seq_len, 1] |
| 664 | let pos: Vec<f32> = (0..max_seq_len).map(|i| i as f32).collect(); |
| 665 | let pos = Tensor::from_slice(pos.as_slice(), (max_seq_len, 1), device)?; |
| 666 | |
| 667 | // outer product → [max_seq_len, half] |
| 668 | let freqs = pos.broadcast_mul(&freqs)?; |
| 669 | |
| 670 | // Duplicate halves so both halves receive the same rotation angle. |
| 671 | // Result shape: [max_seq_len, head_dim] |
| 672 | let cos = Tensor::cat(&[&freqs.cos()?, &freqs.cos()?], 1)?; |
| 673 | let sin = Tensor::cat(&[&freqs.sin()?, &freqs.sin()?], 1)?; |
| 674 | |
| 675 | Ok((cos, sin)) |
| 676 | } |
| 677 | |
| 678 | // --------------------------------------------------------------------------- |
| 679 | // Private: LoRA-augmented linear layer |
| 680 | // --------------------------------------------------------------------------- |
| 681 | |
| 682 | struct LoraLinear { |
| 683 | /// Frozen base weight `[out_features, in_features]`. |
| 684 | weight: Tensor, |
| 685 | /// Optional frozen base bias `[out_features]`. |
| 686 | bias: Option<Tensor>, |
| 687 | /// Trainable LoRA A matrix `[rank, in_features]`, initialised N(0, 1/√rank). |
| 688 | lora_a: Var, |
| 689 | /// Trainable LoRA B matrix `[out_features, rank]`, initialised to zeros. |
| 690 | lora_b: Var, |
| 691 | /// Scaling factor `alpha / rank`. |
| 692 | scale: f64, |
| 693 | } |
| 694 | |
| 695 | impl LoraLinear { |
| 696 | fn new( |
| 697 | weight: Tensor, |
| 698 | bias: Option<Tensor>, |
| 699 | rank: usize, |
| 700 | alpha: f32, |
| 701 | device: &Device, |
| 702 | ) -> candle_core::Result<Self> { |
| 703 | let shape = weight.shape().dims(); |
| 704 | // weight is [out_features, in_features] |
| 705 | let out_features = shape[0]; |
| 706 | let in_features = shape[1]; |
| 707 | |
| 708 | let std_val = (1.0 / (rank as f64).sqrt()) as f32; |
| 709 | let lora_a = Var::randn(0.0f32, std_val, (rank, in_features), device)?; |
| 710 | let lora_b = Var::zeros((out_features, rank), DType::F32, device)?; |
| 711 | let scale = alpha as f64 / rank as f64; |
| 712 | |
| 713 | Ok(Self { |
| 714 | weight, |
| 715 | bias, |
| 716 | lora_a, |
| 717 | lora_b, |
| 718 | scale, |
| 719 | }) |
| 720 | } |
| 721 | |
| 722 | /// `y = x W^T + b + (x A^T) B^T * scale` |
| 723 | /// |
| 724 | /// Flattens to 2D before matmul because the published candle-core 0.10.2 |
| 725 | /// Metal kernel does not support 3D × 2D matmul. |
| 726 | fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> { |
| 727 | let (b, seq, in_f) = x.dims3()?; |
| 728 | let x_flat = x.reshape((b * seq, in_f))?; |
| 729 | |
| 730 | // Base projection |
| 731 | let base = x_flat.matmul(&self.weight.t()?)?; |
| 732 | |
| 733 | // LoRA path: x A^T → x B^T scaled |
| 734 | let lora = x_flat |
| 735 | .matmul(&self.lora_a.as_tensor().t()?)? |
| 736 | .matmul(&self.lora_b.as_tensor().t()?)? |
| 737 | .affine(self.scale, 0.0)?; |
| 738 | |
| 739 | let out_f = self.weight.dim(0)?; |
| 740 | let combined = (base + lora)?.reshape((b, seq, out_f))?; |
| 741 | |
| 742 | match &self.bias { |
| 743 | Some(bias) => combined.broadcast_add(bias), |
| 744 | None => Ok(combined), |
| 745 | } |
| 746 | } |
| 747 | |
| 748 | /// Return the two trainable LoRA variables. |
| 749 | fn vars(&self) -> Vec<Var> { |
| 750 | vec![self.lora_a.clone(), self.lora_b.clone()] |
| 751 | } |
| 752 | } |
| 753 | |
| 754 | // --------------------------------------------------------------------------- |
| 755 | // Private: frozen linear layer (no LoRA) |
| 756 | // --------------------------------------------------------------------------- |
| 757 | |
| 758 | struct FrozenLinear { |
| 759 | weight: Tensor, |
| 760 | bias: Option<Tensor>, |
| 761 | } |
| 762 | |
| 763 | impl FrozenLinear { |
| 764 | fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> { |
| 765 | let (b, seq, in_f) = x.dims3()?; |
| 766 | let x_flat = x.reshape((b * seq, in_f))?; |
| 767 | let out_flat = x_flat.matmul(&self.weight.t()?)?; |
| 768 | let out_f = self.weight.dim(0)?; |
| 769 | let out = out_flat.reshape((b, seq, out_f))?; |
| 770 | match &self.bias { |
| 771 | Some(bias) => out.broadcast_add(bias), |
| 772 | None => Ok(out), |
| 773 | } |
| 774 | } |
| 775 | } |
| 776 | |
| 777 | // --------------------------------------------------------------------------- |
| 778 | // Private: single transformer block |
| 779 | // --------------------------------------------------------------------------- |
| 780 | |
| 781 | struct TransformerLayer { |
| 782 | // Attention projections |
| 783 | q_proj: LoraLinear, // LoRA applied |
| 784 | k_proj: FrozenLinear, // frozen |
| 785 | v_proj: LoraLinear, // LoRA applied |
| 786 | o_proj: FrozenLinear, // frozen |
| 787 | // QK-norm (Qwen3 only — None for Qwen2/2.5) |
| 788 | q_norm: Option<Tensor>, // [head_dim] |
| 789 | k_norm: Option<Tensor>, // [head_dim] |
| 790 | // MLP projections (SwiGLU) |
| 791 | gate_proj: FrozenLinear, |
| 792 | up_proj: FrozenLinear, |
| 793 | down_proj: FrozenLinear, |
| 794 | // LayerNorm weights (RMSNorm, no bias) |
| 795 | input_layernorm: Tensor, // [hidden_size] |
| 796 | post_attention_layernorm: Tensor, // [hidden_size] |
| 797 | // Hyperparameters |
| 798 | num_heads: usize, |
| 799 | num_kv_heads: usize, |
| 800 | head_dim: usize, |
| 801 | rms_norm_eps: f64, |
| 802 | } |
| 803 | |
| 804 | impl TransformerLayer { |
| 805 | /// Full pre-norm transformer block forward pass. |
| 806 | /// |
| 807 | /// `x` has shape `[batch, seq, hidden_size]`. |
| 808 | /// `cos` / `sin` have shape `[seq, head_dim]` (sliced from the precomputed tables). |
| 809 | fn forward(&self, x: &Tensor, cos: &Tensor, sin: &Tensor) -> candle_core::Result<Tensor> { |
| 810 | let (b, seq, _hidden) = x.dims3()?; |
| 811 | |
| 812 | // ----- Attention sub-layer ----- |
| 813 | |
| 814 | // Pre-attention RMSNorm |
| 815 | let normed = rms_norm(x, &self.input_layernorm, self.rms_norm_eps)?; |
| 816 | |
| 817 | // QKV projections |
| 818 | let q = self.q_proj.forward(&normed)?; // [b, seq, num_heads * head_dim] |
| 819 | let k = self.k_proj.forward(&normed)?; // [b, seq, num_kv_heads * head_dim] |
| 820 | let v = self.v_proj.forward(&normed)?; // [b, seq, num_kv_heads * head_dim] |
| 821 | |
| 822 | // Reshape to [b, heads, seq, head_dim] |
| 823 | let q = q |
| 824 | .reshape((b, seq, self.num_heads, self.head_dim))? |
| 825 | .transpose(1, 2)?; |
| 826 | let k = k |
| 827 | .reshape((b, seq, self.num_kv_heads, self.head_dim))? |
| 828 | .transpose(1, 2)?; |
| 829 | let v = v |
| 830 | .reshape((b, seq, self.num_kv_heads, self.head_dim))? |
| 831 | .transpose(1, 2)?; |
| 832 | |
| 833 | // Qwen3 QK-Norm: apply per-head RMSNorm to Q and K before RoPE. |
| 834 | let q = if let Some(w) = &self.q_norm { |
| 835 | rms_norm(&q, w, self.rms_norm_eps)? |
| 836 | } else { |
| 837 | q |
| 838 | }; |
| 839 | let k = if let Some(w) = &self.k_norm { |
| 840 | rms_norm(&k, w, self.rms_norm_eps)? |
| 841 | } else { |
| 842 | k |
| 843 | }; |
| 844 | |
| 845 | // Apply RoPE (requires contiguous storage for the custom kernel) |
| 846 | let q = apply_rope(&q.contiguous()?, cos, sin)?; |
| 847 | let k = apply_rope(&k.contiguous()?, cos, sin)?; |
| 848 | |
| 849 | // Grouped Query Attention: expand K/V heads if num_heads > num_kv_heads |
| 850 | let groups = self.num_heads / self.num_kv_heads; |
| 851 | let (k, v) = if groups == 1 { |
| 852 | (k.contiguous()?, v.contiguous()?) |
| 853 | } else { |
| 854 | let k_exp = k |
| 855 | .unsqueeze(2)? |
| 856 | .expand((b, self.num_kv_heads, groups, seq, self.head_dim))? |
| 857 | .reshape((b, self.num_heads, seq, self.head_dim))? |
| 858 | .contiguous()?; |
| 859 | let v_exp = v |
| 860 | .unsqueeze(2)? |
| 861 | .expand((b, self.num_kv_heads, groups, seq, self.head_dim))? |
| 862 | .reshape((b, self.num_heads, seq, self.head_dim))? |
| 863 | .contiguous()?; |
| 864 | (k_exp, v_exp) |
| 865 | }; |
| 866 | |
| 867 | // Scaled dot-product attention scores |
| 868 | let scale = 1.0 / (self.head_dim as f64).sqrt(); |
| 869 | let attn_weights = (q.matmul(&k.transpose(2, 3)?)? * scale)?; |
| 870 | |
| 871 | // Causal mask: add a large negative bias to positions the query must |
| 872 | // not attend to. |
| 873 | let device = x.device(); |
| 874 | let tril = |
| 875 | Tensor::tril2(seq, DType::F32, device)?.broadcast_as((b, self.num_heads, seq, seq))?; |
| 876 | // (1 - tril) is 1 on upper triangle (blocked), 0 on lower triangle (allowed). |
| 877 | let attn_bias = (1.0f64 - &tril)?.affine(-1e9, 0.0)?; |
| 878 | let attn_weights = (attn_weights + attn_bias)?; |
| 879 | // Use the generic softmax (broadcast_sub/broadcast_div internally) — the |
| 880 | // CustomOp-based softmax_last_dim has no Metal implementation in candle 0.10.2. |
| 881 | let attn_weights = candle_nn::ops::softmax(&attn_weights, D::Minus1)?; |
| 882 | |
| 883 | // Weighted sum of values |
| 884 | let attn_out = attn_weights.matmul(&v)?; |
| 885 | |
| 886 | // Merge heads: [b, heads, seq, head_dim] → [b, seq, hidden_size] |
| 887 | let attn_out = attn_out.transpose(1, 2)?.contiguous()?.reshape(( |
| 888 | b, |
| 889 | seq, |
| 890 | self.num_heads * self.head_dim, |
| 891 | ))?; |
| 892 | |
| 893 | // Output projection |
| 894 | let attn_out = self.o_proj.forward(&attn_out)?; |
| 895 | |
| 896 | // First residual connection |
| 897 | let h = (x + &attn_out)?; |
| 898 | |
| 899 | // ----- MLP sub-layer (SwiGLU) ----- |
| 900 | |
| 901 | // Pre-MLP RMSNorm |
| 902 | let normed2 = rms_norm(&h, &self.post_attention_layernorm, self.rms_norm_eps)?; |
| 903 | |
| 904 | let gate = candle_nn::ops::silu(&self.gate_proj.forward(&normed2)?)?; |
| 905 | let up = self.up_proj.forward(&normed2)?; |
| 906 | let mlp_out = self.down_proj.forward(&(gate * up)?)?; |
| 907 | |
| 908 | // Second residual connection |
| 909 | h + mlp_out |
| 910 | } |
| 911 | |
| 912 | /// All trainable LoRA variables in this layer. |
| 913 | fn lora_vars(&self) -> Vec<Var> { |
| 914 | let mut v = self.q_proj.vars(); |
| 915 | v.extend(self.v_proj.vars()); |
| 916 | v |
| 917 | } |
| 918 | } |
| 919 | |
| 920 | // --------------------------------------------------------------------------- |
| 921 | // Private: full model |
| 922 | // --------------------------------------------------------------------------- |
| 923 | |
| 924 | struct LoraQwenModel { |
| 925 | embed_tokens: Tensor, // [vocab_size, hidden_size] |
| 926 | layers: Vec<TransformerLayer>, |
| 927 | norm: Tensor, // [hidden_size] — final RMSNorm weight |
| 928 | lm_head: Tensor, // [vocab_size, hidden_size] |
| 929 | rope_cos: Tensor, // [max_seq_len, head_dim] |
| 930 | rope_sin: Tensor, // [max_seq_len, head_dim] |
| 931 | hidden_size: usize, |
| 932 | #[allow(dead_code)] |
| 933 | vocab_size: usize, |
| 934 | rms_norm_eps: f64, |
| 935 | } |
| 936 | |
| 937 | impl LoraQwenModel { |
| 938 | /// Load base weights from a `VarBuilder` and inject LoRA adapters. |
| 939 | fn load( |
| 940 | vb: VarBuilder, |
| 941 | cfg: &ModelConfig, |
| 942 | rank: usize, |
| 943 | alpha: f32, |
| 944 | device: &Device, |
| 945 | ) -> candle_core::Result<Self> { |
| 946 | let hidden_size = cfg.hidden_size; |
| 947 | let vocab_size = cfg.vocab_size; |
| 948 | let num_heads = cfg.num_attention_heads; |
| 949 | let num_kv_heads = cfg.num_key_value_heads.unwrap_or(num_heads); |
| 950 | let head_dim = cfg.head_dim.unwrap_or(hidden_size / num_heads); |
| 951 | let intermediate_size = cfg.intermediate_size; |
| 952 | let rms_norm_eps = cfg.rms_norm_eps.unwrap_or(1e-6); |
| 953 | let rope_theta = cfg.rope_theta.unwrap_or(10_000.0); |
| 954 | |
| 955 | // Maximum sequence length for the RoPE tables. |
| 956 | // 2048 covers all models targeted by this module; increase if needed. |
| 957 | let max_seq_len = 2048usize; |
| 958 | |
| 959 | // Embedding table |
| 960 | let embed_tokens = vb |
| 961 | .pp("model") |
| 962 | .pp("embed_tokens") |
| 963 | .get((vocab_size, hidden_size), "weight")?; |
| 964 | |
| 965 | // Transformer layers |
| 966 | let mut layers: Vec<TransformerLayer> = Vec::with_capacity(cfg.num_hidden_layers); |
| 967 | |
| 968 | for i in 0..cfg.num_hidden_layers { |
| 969 | let vb_layer = vb.pp("model").pp("layers").pp(i.to_string()); |
| 970 | |
| 971 | let vb_attn = vb_layer.pp("self_attn"); |
| 972 | |
| 973 | // --- Q projection (LoRA) --- |
| 974 | let q_out = num_heads * head_dim; |
| 975 | let q_weight = vb_attn.pp("q_proj").get((q_out, hidden_size), "weight")?; |
| 976 | // Qwen2 has bias on q/k/v; Qwen3 does not — use .ok() to handle both. |
| 977 | let q_bias = vb_attn.pp("q_proj").get((q_out,), "bias").ok(); |
| 978 | let q_proj = LoraLinear::new(q_weight, q_bias, rank, alpha, device)?; |
| 979 | |
| 980 | // --- K projection (frozen) --- |
| 981 | let kv_out = num_kv_heads * head_dim; |
| 982 | let k_weight = vb_attn.pp("k_proj").get((kv_out, hidden_size), "weight")?; |
| 983 | let k_bias = vb_attn.pp("k_proj").get((kv_out,), "bias").ok(); |
| 984 | let k_proj = FrozenLinear { |
| 985 | weight: k_weight, |
| 986 | bias: k_bias, |
| 987 | }; |
| 988 | |
| 989 | // --- V projection (LoRA) --- |
| 990 | let v_weight = vb_attn.pp("v_proj").get((kv_out, hidden_size), "weight")?; |
| 991 | let v_bias = vb_attn.pp("v_proj").get((kv_out,), "bias").ok(); |
| 992 | let v_proj = LoraLinear::new(v_weight, v_bias, rank, alpha, device)?; |
| 993 | |
| 994 | // --- O projection (frozen) --- |
| 995 | let o_weight = vb_attn |
| 996 | .pp("o_proj") |
| 997 | .get((hidden_size, num_heads * head_dim), "weight")?; |
| 998 | let o_proj = FrozenLinear { |
| 999 | weight: o_weight, |
| 1000 | bias: None, |
| 1001 | }; |
| 1002 | |
| 1003 | // --- QK-Norm (Qwen3 only) --- |
| 1004 | let q_norm = vb_attn.pp("q_norm").get((head_dim,), "weight").ok(); |
| 1005 | let k_norm = vb_attn.pp("k_norm").get((head_dim,), "weight").ok(); |
| 1006 | |
| 1007 | // --- MLP projections (all frozen) --- |
| 1008 | let vb_mlp = vb_layer.pp("mlp"); |
| 1009 | |
| 1010 | let gate_weight = vb_mlp |
| 1011 | .pp("gate_proj") |
| 1012 | .get((intermediate_size, hidden_size), "weight")?; |
| 1013 | let up_weight = vb_mlp |
| 1014 | .pp("up_proj") |
| 1015 | .get((intermediate_size, hidden_size), "weight")?; |
| 1016 | let down_weight = vb_mlp |
| 1017 | .pp("down_proj") |
| 1018 | .get((hidden_size, intermediate_size), "weight")?; |
| 1019 | |
| 1020 | // --- LayerNorm weights --- |
| 1021 | let input_layernorm = vb_layer |
| 1022 | .pp("input_layernorm") |
| 1023 | .get((hidden_size,), "weight")?; |
| 1024 | let post_attention_layernorm = vb_layer |
| 1025 | .pp("post_attention_layernorm") |
| 1026 | .get((hidden_size,), "weight")?; |
| 1027 | |
| 1028 | layers.push(TransformerLayer { |
| 1029 | q_proj, |
| 1030 | k_proj, |
| 1031 | v_proj, |
| 1032 | o_proj, |
| 1033 | q_norm, |
| 1034 | k_norm, |
| 1035 | gate_proj: FrozenLinear { |
| 1036 | weight: gate_weight, |
| 1037 | bias: None, |
| 1038 | }, |
| 1039 | up_proj: FrozenLinear { |
| 1040 | weight: up_weight, |
| 1041 | bias: None, |
| 1042 | }, |
| 1043 | down_proj: FrozenLinear { |
| 1044 | weight: down_weight, |
| 1045 | bias: None, |
| 1046 | }, |
| 1047 | input_layernorm, |
| 1048 | post_attention_layernorm, |
| 1049 | num_heads, |
| 1050 | num_kv_heads, |
| 1051 | head_dim, |
| 1052 | rms_norm_eps, |
| 1053 | }); |
| 1054 | } |
| 1055 | |
| 1056 | // Final RMSNorm and language model head |
| 1057 | let norm = vb.pp("model").pp("norm").get((hidden_size,), "weight")?; |
| 1058 | let lm_head = vb |
| 1059 | .pp("lm_head") |
| 1060 | .get((vocab_size, hidden_size), "weight") |
| 1061 | .unwrap_or_else(|_| embed_tokens.clone()); |
| 1062 | |
| 1063 | // Precompute RoPE tables |
| 1064 | let (rope_cos, rope_sin) = precompute_rope(head_dim, max_seq_len, rope_theta, device)?; |
| 1065 | |
| 1066 | Ok(Self { |
| 1067 | embed_tokens, |
| 1068 | layers, |
| 1069 | norm, |
| 1070 | lm_head, |
| 1071 | rope_cos, |
| 1072 | rope_sin, |
| 1073 | hidden_size, |
| 1074 | vocab_size, |
| 1075 | rms_norm_eps, |
| 1076 | }) |
| 1077 | } |
| 1078 | |
| 1079 | /// Full forward pass. |
| 1080 | /// |
| 1081 | /// `input_ids` has shape `[batch, seq]`. |
| 1082 | /// Returns logits with shape `[batch, seq, vocab_size]`. |
| 1083 | fn forward(&self, input_ids: &Tensor) -> candle_core::Result<Tensor> { |
| 1084 | let (b, seq) = input_ids.dims2()?; |
| 1085 | |
| 1086 | // Token embedding lookup: [batch, seq, hidden_size] |
| 1087 | let ids_flat = input_ids.flatten_all()?; // [b * seq] |
| 1088 | let mut hidden = self |
| 1089 | .embed_tokens |
| 1090 | .embedding(&ids_flat)? // [b * seq, hidden_size] |
| 1091 | .reshape((b, seq, self.hidden_size))?; // [b, seq, hidden_size] |
| 1092 | |
| 1093 | // Slice RoPE tables to the current sequence length. |
| 1094 | let cos = self.rope_cos.narrow(0, 0, seq)?; // [seq, head_dim] |
| 1095 | let sin = self.rope_sin.narrow(0, 0, seq)?; // [seq, head_dim] |
| 1096 | |
| 1097 | // Transformer layers |
| 1098 | for layer in &self.layers { |
| 1099 | hidden = layer.forward(&hidden, &cos, &sin)?; |
| 1100 | } |
| 1101 | |
| 1102 | // Final RMSNorm |
| 1103 | let hidden = rms_norm(&hidden, &self.norm, self.rms_norm_eps)?; |
| 1104 | |
| 1105 | // Language model head: flatten → matmul → reshape (Metal 3D×2D workaround) |
| 1106 | let hidden_flat = hidden.reshape((b * seq, self.hidden_size))?; |
| 1107 | let logits_flat = hidden_flat.matmul(&self.lm_head.t()?)?; |
| 1108 | logits_flat.reshape((b, seq, self.vocab_size)) |
| 1109 | } |
| 1110 | |
| 1111 | /// Collect all trainable LoRA variables across every layer. |
| 1112 | fn lora_vars(&self) -> Vec<Var> { |
| 1113 | self.layers |
| 1114 | .iter() |
| 1115 | .flat_map(|layer| layer.lora_vars()) |
| 1116 | .collect() |
| 1117 | } |
| 1118 | |
| 1119 | /// Serialise the LoRA adapter weights to `output_dir/lora_adapter.safetensors`. |
| 1120 | /// |
| 1121 | /// Saves one `lora_a` and one `lora_b` tensor per projection (Q and V) per layer. |
| 1122 | fn save_lora(&self, output_dir: &Path) -> Result<PathBuf> { |
| 1123 | std::fs::create_dir_all(output_dir) |
| 1124 | .with_context(|| format!("creating output directory {:?}", output_dir))?; |
| 1125 | |
| 1126 | let mut tensors: HashMap<String, Tensor> = HashMap::new(); |
| 1127 | |
| 1128 | for (i, layer) in self.layers.iter().enumerate() { |
| 1129 | tensors.insert( |
| 1130 | format!("model.layers.{i}.self_attn.q_proj.lora_a"), |
| 1131 | layer.q_proj.lora_a.as_tensor().clone(), |
| 1132 | ); |
| 1133 | tensors.insert( |
| 1134 | format!("model.layers.{i}.self_attn.q_proj.lora_b"), |
| 1135 | layer.q_proj.lora_b.as_tensor().clone(), |
| 1136 | ); |
| 1137 | tensors.insert( |
| 1138 | format!("model.layers.{i}.self_attn.v_proj.lora_a"), |
| 1139 | layer.v_proj.lora_a.as_tensor().clone(), |
| 1140 | ); |
| 1141 | tensors.insert( |
| 1142 | format!("model.layers.{i}.self_attn.v_proj.lora_b"), |
| 1143 | layer.v_proj.lora_b.as_tensor().clone(), |
| 1144 | ); |
| 1145 | } |
| 1146 | |
| 1147 | let path = output_dir.join("lora_adapter.safetensors"); |
| 1148 | candle_core::safetensors::save(&tensors, &path) |
| 1149 | .with_context(|| format!("writing {:?}", path))?; |
| 1150 | |
| 1151 | Ok(path) |
| 1152 | } |
| 1153 | } |
| 1154 | |
| 1155 | // --------------------------------------------------------------------------- |
| 1156 | // Tests |
| 1157 | // --------------------------------------------------------------------------- |
| 1158 | |
| 1159 | #[cfg(test)] |
| 1160 | mod tests { |
| 1161 | use super::*; |
| 1162 | |
| 1163 | #[test] |
| 1164 | fn test_dataentry_raw_deserialization() { |
| 1165 | let json = r#"{"text":"hello world"}"#; |
| 1166 | let entry: DataEntry = serde_json::from_str(json).unwrap(); |
| 1167 | match entry { |
| 1168 | DataEntry::Raw { text } => assert_eq!(text, "hello world"), |
| 1169 | _ => panic!("expected Raw variant"), |
| 1170 | } |
| 1171 | } |
| 1172 | |
| 1173 | #[test] |
| 1174 | fn test_dataentry_messages_deserialization() { |
| 1175 | let json = r#"{"messages":[{"role":"user","content":"hi"}]}"#; |
| 1176 | let entry: DataEntry = serde_json::from_str(json).unwrap(); |
| 1177 | match entry { |
| 1178 | DataEntry::Messages { messages } => { |
| 1179 | assert_eq!(messages.len(), 1); |
| 1180 | assert_eq!(messages[0].role, "user"); |
| 1181 | assert_eq!(messages[0].content, "hi"); |
| 1182 | } |
| 1183 | _ => panic!("expected Messages variant"), |
| 1184 | } |
| 1185 | } |
| 1186 | |
| 1187 | #[test] |
| 1188 | fn test_dataentry_invalid() { |
| 1189 | let json = r#"{"invalid":"field"}"#; |
| 1190 | let result: Result<DataEntry, _> = serde_json::from_str(json); |
| 1191 | assert!(result.is_err(), "should reject unrecognized format"); |
| 1192 | } |
| 1193 | |
| 1194 | #[test] |
| 1195 | fn test_render_chatml_basic() { |
| 1196 | let messages = vec![ |
| 1197 | DatasetMessage { |
| 1198 | role: "system".to_string(), |
| 1199 | content: "You are helpful".to_string(), |
| 1200 | }, |
| 1201 | DatasetMessage { |
| 1202 | role: "user".to_string(), |
| 1203 | content: "Hello".to_string(), |
| 1204 | }, |
| 1205 | DatasetMessage { |
| 1206 | role: "assistant".to_string(), |
| 1207 | content: "Hi there".to_string(), |
| 1208 | }, |
| 1209 | ]; |
| 1210 | |
| 1211 | let (rendered, spans) = render_chatml(&messages); |
| 1212 | |
| 1213 | assert!( |
| 1214 | rendered.contains("<|im_start|>system"), |
| 1215 | "should contain system marker" |
| 1216 | ); |
| 1217 | assert!( |
| 1218 | rendered.contains("You are helpful"), |
| 1219 | "should contain system content" |
| 1220 | ); |
| 1221 | assert!( |
| 1222 | rendered.contains("<|im_start|>user"), |
| 1223 | "should contain user marker" |
| 1224 | ); |
| 1225 | assert!(rendered.contains("Hello"), "should contain user content"); |
| 1226 | assert!( |
| 1227 | rendered.contains("<|im_start|>assistant"), |
| 1228 | "should contain assistant marker" |
| 1229 | ); |
| 1230 | assert!( |
| 1231 | rendered.contains("Hi there"), |
| 1232 | "should contain assistant content" |
| 1233 | ); |
| 1234 | |
| 1235 | assert_eq!(spans.len(), 1, "should have exactly one assistant span"); |
| 1236 | let (start, end) = spans[0]; |
| 1237 | // The span covers the assistant content and its closing `<|im_end|>`, |
| 1238 | // so the model is trained to emit the stop marker. |
| 1239 | assert!(rendered[start..end].contains("Hi there")); |
| 1240 | assert!(rendered[start..end].ends_with("<|im_end|>")); |
| 1241 | // Only the trailing newline remains after the span. |
| 1242 | assert!(rendered[end..].starts_with('\n')); |
| 1243 | } |
| 1244 | |
| 1245 | #[test] |
| 1246 | fn test_render_chatml_multi_assistant() { |
| 1247 | let messages = vec![ |
| 1248 | DatasetMessage { |
| 1249 | role: "assistant".to_string(), |
| 1250 | content: "First".to_string(), |
| 1251 | }, |
| 1252 | DatasetMessage { |
| 1253 | role: "user".to_string(), |
| 1254 | content: "Second".to_string(), |
| 1255 | }, |
| 1256 | DatasetMessage { |
| 1257 | role: "assistant".to_string(), |
| 1258 | content: "Third".to_string(), |
| 1259 | }, |
| 1260 | ]; |
| 1261 | |
| 1262 | let (_rendered, spans) = render_chatml(&messages); |
| 1263 | assert_eq!(spans.len(), 2, "should have exactly two assistant spans"); |
| 1264 | } |
| 1265 | |
| 1266 | #[test] |
| 1267 | fn test_render_chatml_no_assistant() { |
| 1268 | let messages = vec![DatasetMessage { |
| 1269 | role: "user".to_string(), |
| 1270 | content: "Just a question".to_string(), |
| 1271 | }]; |
| 1272 | |
| 1273 | let (_rendered, spans) = render_chatml(&messages); |
| 1274 | assert_eq!(spans.len(), 0, "should have no assistant spans"); |
| 1275 | } |
| 1276 | |
| 1277 | #[test] |
| 1278 | fn test_render_chatml_utf8_indonesian() { |
| 1279 | let messages = vec![ |
| 1280 | DatasetMessage { |
| 1281 | role: "user".to_string(), |
| 1282 | content: "Apa kabar? Ini ada diacritik: é".to_string(), |
| 1283 | }, |
| 1284 | DatasetMessage { |
| 1285 | role: "assistant".to_string(), |
| 1286 | content: "Baik-baik saja. Jawaban dengan: ñ dan ü".to_string(), |
| 1287 | }, |
| 1288 | ]; |
| 1289 | |
| 1290 | let (rendered, spans) = render_chatml(&messages); |
| 1291 | assert_eq!(spans.len(), 1, "should have one assistant span"); |
| 1292 | |
| 1293 | let (start, end) = spans[0]; |
| 1294 | let assistant_content = &rendered[start..end]; |
| 1295 | assert!(assistant_content.contains("Baik-baik saja")); |
| 1296 | assert!(assistant_content.contains("ñ")); |
| 1297 | assert!(assistant_content.contains("ü")); |
| 1298 | } |
| 1299 | |
| 1300 | #[test] |
| 1301 | fn test_compute_loss_mask_all_inside() { |
| 1302 | let offsets = vec![(10, 15), (15, 20), (20, 25)]; |
| 1303 | let spans = vec![(0, 30)]; |
| 1304 | let mask = compute_loss_mask(&offsets, &spans); |
| 1305 | |
| 1306 | assert_eq!(mask, vec![true, true, true]); |
| 1307 | } |
| 1308 | |
| 1309 | #[test] |
| 1310 | fn test_compute_loss_mask_all_outside() { |
| 1311 | let offsets = vec![(0, 5), (5, 10), (40, 45)]; |
| 1312 | let spans = vec![(20, 30)]; |
| 1313 | let mask = compute_loss_mask(&offsets, &spans); |
| 1314 | |
| 1315 | assert_eq!(mask, vec![false, false, false]); |
| 1316 | } |
| 1317 | |
| 1318 | #[test] |
| 1319 | fn test_compute_loss_mask_straddling_start() { |
| 1320 | let offsets = vec![(15, 25), (25, 35)]; |
| 1321 | let spans = vec![(20, 40)]; |
| 1322 | let mask = compute_loss_mask(&offsets, &spans); |
| 1323 | |
| 1324 | // First token [15, 25) overlaps [20, 40) → true |
| 1325 | // Second token [25, 35) overlaps [20, 40) → true |
| 1326 | assert_eq!(mask, vec![true, true]); |
| 1327 | } |
| 1328 | |
| 1329 | #[test] |
| 1330 | fn test_compute_loss_mask_straddling_end() { |
| 1331 | let offsets = vec![(10, 25), (25, 35)]; |
| 1332 | let spans = vec![(15, 30)]; |
| 1333 | let mask = compute_loss_mask(&offsets, &spans); |
| 1334 | |
| 1335 | // First token [10, 25) overlaps [15, 30) → true |
| 1336 | // Second token [25, 35) overlaps [15, 30) → true |
| 1337 | assert_eq!(mask, vec![true, true]); |
| 1338 | } |
| 1339 | |
| 1340 | #[test] |
| 1341 | fn test_compute_loss_mask_zero_length_span() { |
| 1342 | // A zero-length span occurs when assistant content is empty — the |
| 1343 | // span collapses to a single point where content_start == end_tag_start. |
| 1344 | // Any token whose range contains that point (e.g. the immediately |
| 1345 | // following `<|im_end|>` token) must still be masked true, since the |
| 1346 | // model needs to learn to emit `<|im_end|>` right after the role |
| 1347 | // header even for an empty response. |
| 1348 | let offsets = vec![(10, 20), (20, 30)]; |
| 1349 | let spans = vec![(15, 15)]; // zero-length at byte 15 |
| 1350 | let mask = compute_loss_mask(&offsets, &spans); |
| 1351 | |
| 1352 | // Token [10, 20) contains the zero-length span's point (15) → true. |
| 1353 | // Token [20, 30) does not → false. |
| 1354 | assert_eq!(mask, vec![true, false]); |
| 1355 | } |
| 1356 | |
| 1357 | #[test] |
| 1358 | fn test_compute_loss_mask_multiple_spans() { |
| 1359 | let offsets = vec![(0, 5), (5, 10), (10, 15), (15, 20), (20, 25)]; |
| 1360 | let spans = vec![(2, 12), (18, 22)]; |
| 1361 | let mask = compute_loss_mask(&offsets, &spans); |
| 1362 | |
| 1363 | // [0, 5) overlaps [2, 12) → true |
| 1364 | // [5, 10) overlaps [2, 12) → true |
| 1365 | // [10, 15) overlaps [2, 12) → true |
| 1366 | // [15, 20) overlaps [18, 22) → true |
| 1367 | // [20, 25) overlaps [18, 22) → true |
| 1368 | assert_eq!(mask, vec![true, true, true, true, true]); |
| 1369 | } |
| 1370 | |
| 1371 | #[test] |
| 1372 | fn test_assistant_span_masks_closing_im_end() { |
| 1373 | // The closing `<|im_end|>` of an assistant turn must be masked in for |
| 1374 | // loss so the model learns to stop. Its byte range must fall inside the |
| 1375 | // reported span, and every token it produces must be masked true — |
| 1376 | // whether the marker tokenizes as one special token or splits into |
| 1377 | // several literal pieces (which can happen with add_special_tokens=false). |
| 1378 | let messages = vec![DatasetMessage { |
| 1379 | role: "assistant".to_string(), |
| 1380 | content: "Hi".to_string(), |
| 1381 | }]; |
| 1382 | |
| 1383 | let (rendered, spans) = render_chatml(&messages); |
| 1384 | assert_eq!(spans.len(), 1); |
| 1385 | let (span_start, span_end) = spans[0]; |
| 1386 | |
| 1387 | let marker = "<|im_end|>"; |
| 1388 | let marker_start = rendered.find(marker).expect("marker present"); |
| 1389 | let marker_end = marker_start + marker.len(); |
| 1390 | |
| 1391 | // The marker's whole byte range sits within the assistant span. |
| 1392 | assert!(marker_start >= span_start && marker_end <= span_end); |
| 1393 | |
| 1394 | // Single-token marker → masked. |
| 1395 | let single = compute_loss_mask(&[(marker_start, marker_end)], &spans); |
| 1396 | assert_eq!(single, vec![true]); |
| 1397 | |
| 1398 | // Split marker (e.g. "<|", "im", "_end", "|>") → every piece masked. |
| 1399 | let split = compute_loss_mask( |
| 1400 | &[ |
| 1401 | (marker_start, marker_start + 2), |
| 1402 | (marker_start + 2, marker_start + 4), |
| 1403 | (marker_start + 4, marker_start + 8), |
| 1404 | (marker_start + 8, marker_end), |
| 1405 | ], |
| 1406 | &spans, |
| 1407 | ); |
| 1408 | assert_eq!(split, vec![true, true, true, true]); |
| 1409 | } |
| 1410 | |
| 1411 | /// Resolve a locally cached Qwen3 `tokenizer.json`, if present, so the |
| 1412 | /// real-tokenizer test below can skip gracefully on machines without the |
| 1413 | /// model. Honours `ONDE_TEST_MODEL_DIR` (a dir containing `tokenizer.json`) |
| 1414 | /// and otherwise falls back to the Onde App Group container cache. |
| 1415 | fn qwen3_tokenizer_path() -> Option<PathBuf> { |
| 1416 | if let Ok(dir) = std::env::var("ONDE_TEST_MODEL_DIR") { |
| 1417 | let p = PathBuf::from(dir).join("tokenizer.json"); |
| 1418 | if p.exists() { |
| 1419 | return Some(p); |
| 1420 | } |
| 1421 | } |
| 1422 | let snapshots = dirs::home_dir()? |
| 1423 | .join("Library/Group Containers/group.com.ondeinference.apps") |
| 1424 | .join("models/hub/models--Qwen--Qwen3-0.6B/snapshots"); |
| 1425 | std::fs::read_dir(&snapshots) |
| 1426 | .ok()? |
| 1427 | .flatten() |
| 1428 | .map(|e| e.path().join("tokenizer.json")) |
| 1429 | .find(|p| p.exists()) |
| 1430 | } |
| 1431 | |
| 1432 | /// Regression guard for the offset-based loss mask against the *real* Qwen3 |
| 1433 | /// tokenizer (siGit review, PR #6): confirms that with |
| 1434 | /// `add_special_tokens=false` the ChatML markers still tokenize as single |
| 1435 | /// special tokens carrying real byte offsets — not synthesized `(0, 0)` / |
| 1436 | /// zero-length spans that would let the closing `<|im_end|>` escape the |
| 1437 | /// mask. Skips when the model is not cached locally. |
| 1438 | #[test] |
| 1439 | fn test_real_tokenizer_masks_assistant_and_im_end() { |
| 1440 | let Some(tok_path) = qwen3_tokenizer_path() else { |
| 1441 | eprintln!("Qwen3 tokenizer not found locally; skipping"); |
| 1442 | return; |
| 1443 | }; |
| 1444 | let tokenizer = Tokenizer::from_file(&tok_path).expect("load tokenizer"); |
| 1445 | |
| 1446 | let messages = vec![ |
| 1447 | DatasetMessage { |
| 1448 | role: "user".to_string(), |
| 1449 | content: "hi".to_string(), |
| 1450 | }, |
| 1451 | DatasetMessage { |
| 1452 | role: "assistant".to_string(), |
| 1453 | content: "Hello there!".to_string(), |
| 1454 | }, |
| 1455 | ]; |
| 1456 | let (rendered, spans) = render_chatml(&messages); |
| 1457 | |
| 1458 | // add_special_tokens=false mirrors the training tokenization path. |
| 1459 | let encoding = tokenizer.encode(rendered.as_str(), false).unwrap(); |
| 1460 | let offsets = encoding.get_offsets(); |
| 1461 | let tokens = encoding.get_tokens(); |
| 1462 | |
| 1463 | // `<|im_end|>` must survive as a single token with a real, non-empty |
| 1464 | // offset — not a split into literal pieces and not a `(0, 0)` offset. |
| 1465 | let im_end_positions: Vec<usize> = tokens |
| 1466 | .iter() |
| 1467 | .enumerate() |
| 1468 | .filter(|(_, t)| t.as_str() == "<|im_end|>") |
| 1469 | .map(|(i, _)| i) |
| 1470 | .collect(); |
| 1471 | assert_eq!( |
| 1472 | im_end_positions.len(), |
| 1473 | 2, |
| 1474 | "expected one <|im_end|> per turn as a single token" |
| 1475 | ); |
| 1476 | for &i in &im_end_positions { |
| 1477 | let (s, e) = offsets[i]; |
| 1478 | assert!( |
| 1479 | e > s, |
| 1480 | "<|im_end|> at token {i} has a zero-length offset {:?}", |
| 1481 | (s, e) |
| 1482 | ); |
| 1483 | } |
| 1484 | |
| 1485 | let mask = compute_loss_mask(offsets, &spans); |
| 1486 | |
| 1487 | // The assistant's closing `<|im_end|>` (the second one) must be masked |
| 1488 | // in; if its offset were synthesized it would fall outside the span and |
| 1489 | // this would fail — exactly the escape the review flagged. |
| 1490 | let assistant_im_end = *im_end_positions.last().unwrap(); |
| 1491 | assert!( |
| 1492 | mask[assistant_im_end], |
| 1493 | "assistant closing <|im_end|> escaped the loss mask" |
| 1494 | ); |
| 1495 | |
| 1496 | // The assistant content tokens are masked; the user turn and the |
| 1497 | // trailing newline after the assistant turn are not. |
| 1498 | let hello = tokens.iter().position(|t| t.contains("Hello")).unwrap(); |
| 1499 | assert!(mask[hello], "assistant content should be masked in"); |
| 1500 | let user_hi = tokens.iter().position(|t| t.as_str() == "hi").unwrap(); |
| 1501 | assert!(!mask[user_hi], "user content must not be masked in"); |
| 1502 | // Newline token immediately after the assistant <|im_end|>. |
| 1503 | if let Some(nl) = offsets.get(assistant_im_end + 1) { |
| 1504 | assert!( |
| 1505 | !compute_loss_mask(&[*nl], &spans)[0], |
| 1506 | "trailing newline after assistant turn must not be masked" |
| 1507 | ); |
| 1508 | } |
| 1509 | } |
| 1510 | } |