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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
use anyhow::Result;
use burn::optim::{AdamWConfig, GradientsParams, Optimizer};
use burn::prelude::*;
use burn::tensor::activation::log_softmax;
use burn::tensor::backend::AutodiffBackend;
use burn::tensor::TensorData;
use crate::data::{collate, Example};
use crate::model::Llama;
/// Causal cross-entropy with a loss mask.
///
/// The shift is the key part: position `i` predicts token `i+1`. And `mask`
/// makes only the assistant answer count, so the prompt adds no gradient.
/// This is exactly what llama.cpp's `llama-finetune` does not do.
pub fn masked_ce_loss<B: Backend>(
logits: Tensor<B, 3>,
ids: Tensor<B, 2, Int>,
mask: Tensor<B, 2>,
) -> Tensor<B, 1> {
let [b, t, v] = logits.dims();
let n = b * (t - 1);
let logits = logits.slice([0..b, 0..t - 1, 0..v]).reshape([n, v]);
let targets = ids.slice([0..b, 1..t]).reshape([n, 1]);
let mask = mask.slice([0..b, 1..t]).reshape([n, 1]);
let logp = log_softmax(logits, 1).gather(1, targets);
let denom = mask.clone().sum().clamp_min(1.0);
-(logp * mask).sum() / denom
}
pub struct TrainCfg {
pub steps: usize,
pub batch_size: usize,
pub lr: f64,
}
pub fn train<B: AutodiffBackend>(
mut model: Llama<B>,
examples: &[Example],
cfg: TrainCfg,
device: &B::Device,
) -> Result<Llama<B>> {
let trainable = model.trainable_ids();
println!(
"training {} LoRA tensors on {} examples\n",
trainable.len(),
examples.len()
);
// B starts at zero. If it is still zero at the end, the gradients never
// reached LoRA; if W_base changes, freezing failed.
let norm_base = |m: &Llama<B>| -> f32 {
m.layers[0]
.attn
.q
.base
.weight
.val()
.powf_scalar(2.0)
.sum()
.sqrt()
.into_scalar()
.elem()
};
println!(
"before -> |B| = {:.6} |W_base| = {:.4}",
model.lora_norm(),
norm_base(&model)
);
let mut optim = AdamWConfig::new().init::<B, Llama<B>>();
let mut cursor = 0usize;
for step in 1..=cfg.steps {
if cursor + cfg.batch_size > examples.len() {
cursor = 0;
}
let batch = collate(&examples[cursor..cursor + cfg.batch_size]);
cursor += cfg.batch_size;
let dims = [batch.batch, batch.seq];
let ids = Tensor::<B, 2, Int>::from_data(TensorData::new(batch.ids, dims), device);
let mask = Tensor::<B, 2>::from_data(TensorData::new(batch.mask, dims), device);
let logits = model.forward(ids.clone());
let loss = masked_ce_loss(logits, ids, mask);
let value: f32 = loss.clone().into_scalar().elem();
let mut grads = loss.backward();
// The heart of LoRA: only the gradients of the A and B tensors are
// collected. The rest of the model is never touched.
let grads = GradientsParams::from_params(&mut grads, &model, &trainable);
model = optim.step(cfg.lr, model, grads);
println!(
"step {step:4}/{} loss {value:.4} seq {}",
cfg.steps, batch.seq
);
}
println!(
"after -> |B| = {:.6} |W_base| = {:.4}",
model.lora_norm(),
norm_base(&model)
);
Ok(model)
}
|
