Skip to content

Commit c33dbd6

Browse files
YdrMasterYdrMaster
committed
perf(llama): 优化分布式切分和参数加载
Signed-off-by: YdrMaster <[email protected]>
1 parent 8729b3d commit c33dbd6

File tree

15 files changed

+602
-643
lines changed

15 files changed

+602
-643
lines changed

‎Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ itertools = "0.13"
3434
env_logger = "0.11"
3535
build-script-cfg = "0.0"
3636

37-
operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "df027a4", default-features = false }
37+
operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "17ff77a", default-features = false }
3838

3939
search-cl-tools = { git = "https://github.com/InfiniTensor/clrt", rev = "9b6289d" }
4040
search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "f40bcb5" }

‎common/src/lib.rs‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,9 @@
1-
use std::{borrow::Borrow, collections::HashMap, hash::Hash, ops::Deref};
1+
use std::{
2+
borrow::Borrow,
3+
collections::HashMap,
4+
hash::Hash,
5+
ops::{Deref, Range},
6+
};
27

38
pub enum Contiguous<'a, T> {
49
Borrowed(&'a [u8]),
@@ -52,3 +57,42 @@ impl<K: Eq + Hash, V> Slab<K, V> {
5257
self.0.entry(key).or_default().push(value);
5358
}
5459
}
60+
61+
#[derive(Clone, Copy, Debug)]
62+
pub struct Distribution {
63+
pub start: usize,
64+
pub len: usize,
65+
pub total: usize,
66+
}
67+
68+
impl Distribution {
69+
pub const MONO: Self = Self {
70+
start: 0,
71+
len: 1,
72+
total: 1,
73+
};
74+
}
75+
76+
pub struct WeightMemCalculator {
77+
align: usize,
78+
size: usize,
79+
}
80+
81+
impl WeightMemCalculator {
82+
#[inline]
83+
pub const fn new(align: usize) -> Self {
84+
Self { align, size: 0 }
85+
}
86+
87+
#[inline]
88+
pub const fn size(&self) -> usize {
89+
self.size
90+
}
91+
92+
#[inline]
93+
pub fn push(&mut self, size: usize) -> Range<usize> {
94+
let start = self.size.div_ceil(self.align) * self.align;
95+
self.size = start + size;
96+
start..self.size
97+
}
98+
}

‎models/llama/common-cpu/src/infer.rs‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use crate::{Operators, RandomSample, Weights};
2+
use common::Distribution;
23
use gguf::GGufModel;
34
use llama::{ext::ggml_quants::f16, LlamaRequest, LlamaStorage, LlamaWorker, Tensor};
45
use operators::{
@@ -8,7 +9,13 @@ use operators::{
89
Blob,
910
};
1011
use regex::Regex;
11-
use std::{iter::zip, ptr::copy_nonoverlapping, slice::from_raw_parts_mut, thread};
12+
use std::{
13+
iter::zip,
14+
ptr::copy_nonoverlapping,
15+
slice::from_raw_parts_mut,
16+
sync::{Arc, Barrier},
17+
thread,
18+
};
1219
use test_utils::{test_infer_paralle, Inference, Task, TokenizerAndPrompt, WorkerSeed};
1320

1421
type Worker<'w> = LlamaWorker<Operators<InprocNode<usize>, AllReduce>, Weights<'w>>;
@@ -51,24 +58,28 @@ fn test_infer() {
5158
.collect()
5259
})
5360
.unwrap_or_else(|| vec![1]);
54-
let count = lens.iter().sum();
61+
let dist = lens.iter().sum();
5562
println!("distribution: {lens:?}");
5663

5764
let (seeds, senders) = WorkerSeed::new(InprocNode::new(lens.len()));
65+
let barrier = Arc::new(Barrier::new(dist + 1));
5866
thread::scope(|s| {
5967
let _workers = zip(lens, seeds)
6068
.enumerate()
6169
.scan(0, |start, (id, (len, seed))| {
62-
let range = *start..*start + len;
63-
*start = range.end;
64-
65-
let mut meta = model.meta.clone();
66-
meta.distribute(range.clone(), count);
70+
let dist = Distribution {
71+
start: *start,
72+
len,
73+
total: dist,
74+
};
75+
*start += len;
6776

77+
let meta = model.meta.distribute(dist);
6878
let model = &model;
79+
let barrier = barrier.clone();
6980
Some(s.spawn(move || {
7081
let WorkerSeed { node, tasks } = seed;
71-
let weights = Weights::new(model, range, count);
82+
let weights = Weights::new(model, dist);
7283
let mut worker = Worker::new(id, &node, meta.clone(), weights);
7384
let mut cache = meta.kv_cache(meta.nctx).map(Blob::new);
7485
let sin_cos = <Operators as llama::Operators>::build_sin_cos(
@@ -85,6 +96,7 @@ fn test_infer() {
8596
from_raw_parts_mut(&mut pair as *mut _ as *mut u8, size_of_val(&pair))
8697
});
8798

99+
barrier.wait();
88100
for task in tasks {
89101
let Task {
90102
nt,
@@ -137,6 +149,7 @@ fn test_infer() {
137149
.collect::<Vec<_>>();
138150

139151
let senders = senders.into_boxed_slice();
152+
barrier.wait();
140153
test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
141154
})
142155
}

‎models/llama/common-cpu/src/lib.rs‎

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
use common::Contiguous;
1+
use common::{Contiguous, Distribution};
22
use llama::{
33
ext::ggml_quants::{self, digit_layout::DigitLayout, f16, DataBlock, QuantExt},
4-
BlkWeight, LlamaBlkStorage, LlamaStorage, Tensor,
4+
LlamaBlkStorage, LlamaBlkWeight, LlamaStorage, Tensor,
55
TensorUsage::Computation,
66
WeightLoader,
77
};
@@ -16,7 +16,7 @@ use std::{
1616
cell::{Ref, RefCell},
1717
marker::PhantomData,
1818
mem::size_of,
19-
ops::{Deref, Range, RangeBounds},
19+
ops::{Deref, Range},
2020
ptr::copy_nonoverlapping,
2121
slice::{from_raw_parts, from_raw_parts_mut},
2222
};
@@ -41,7 +41,7 @@ pub struct Weights<'w> {
4141

4242
pub struct WeightCache {
4343
cache: Blob,
44-
cached_weight: BlkWeight,
44+
cached_weight: LlamaBlkWeight,
4545
cached_weight_iblk: usize,
4646
}
4747

@@ -85,11 +85,7 @@ where
8585
}
8686

8787
impl<'w> Weights<'w> {
88-
pub fn new(
89-
model: &'w LlamaStorage<&'w [u8]>,
90-
range: impl RangeBounds<usize> + Clone,
91-
count: usize,
92-
) -> Self {
88+
pub fn new(model: &'w LlamaStorage<&'w [u8]>, dist: Distribution) -> Self {
9389
let LlamaStorage {
9490
meta,
9591
output_norm,
@@ -100,11 +96,18 @@ impl<'w> Weights<'w> {
10096

10197
let blks = blocks
10298
.iter()
103-
.map(|blk| blk.distribute(meta, range.clone(), count, Blob::new))
99+
.map(|blk| {
100+
blk.clone()
101+
.into_vec()
102+
.into_iter()
103+
.map(|(which, data)| {
104+
(which, meta.distribute_data(which, data, dist, Blob::new))
105+
})
106+
.collect::<LlamaBlkStorage<_>>()
107+
})
104108
.collect::<Box<_>>();
105109

106-
let mut meta = meta.clone();
107-
meta.distribute(range.clone(), count);
110+
let meta = meta.distribute(dist);
108111
let size_qkv = meta.attn_qkv(Computation).take();
109112
let size_o = meta.attn_o(Computation).take();
110113
let size_gate_up = meta.ffn_gate_up(Computation).take();
@@ -113,7 +116,7 @@ impl<'w> Weights<'w> {
113116
let weight_cache = if meta.dt_embd == meta.dt_linear {
114117
RefCell::new(WeightCache {
115118
cache: Blob::new(0),
116-
cached_weight: BlkWeight::AttnQKV,
119+
cached_weight: LlamaBlkWeight::AttnQKV,
117120
cached_weight_iblk: 0,
118121
})
119122
} else {
@@ -131,7 +134,7 @@ impl<'w> Weights<'w> {
131134

132135
RefCell::new(WeightCache {
133136
cache,
134-
cached_weight: BlkWeight::AttnQKV,
137+
cached_weight: LlamaBlkWeight::AttnQKV,
135138
cached_weight_iblk: 0,
136139
})
137140
};
@@ -207,7 +210,7 @@ impl WeightLoader for Weights<'_> {
207210
#[inline]
208211
fn load_blk(
209212
&self,
210-
which: BlkWeight,
213+
which: LlamaBlkWeight,
211214
iblk: usize,
212215
_queue: &QueueOf<Self::Hardware>,
213216
) -> Self::Weight<'_> {
@@ -233,10 +236,10 @@ impl WeightLoader for Weights<'_> {
233236
ffn_down,
234237
} = &blks[iblk];
235238

236-
use BlkWeight::{
239+
use Dequant::{Borrowed, Cached};
240+
use LlamaBlkWeight::{
237241
AttnNorm, AttnO, AttnQKV, AttnQKVBias, FfnDown, FfnGateInp, FfnGateUp, FfnNorm,
238242
};
239-
use Dequant::{Borrowed, Cached};
240243

241244
#[rustfmt::skip]
242245
match which {
@@ -301,7 +304,7 @@ impl WeightLoader for Weights<'_> {
301304

302305
fn load_moe<'a>(
303306
&'a self,
304-
which: BlkWeight,
307+
which: LlamaBlkWeight,
305308
iblk: usize,
306309
iexp: usize,
307310
_queue: &'a QueueOf<Self::Hardware>,
@@ -315,8 +318,8 @@ impl WeightLoader for Weights<'_> {
315318
} = self;
316319
assert_eq!(dt_embd, dt_mat);
317320
let w = match which {
318-
BlkWeight::FfnGateUp => &*blks[iblk].ffn_gate_up,
319-
BlkWeight::FfnDown => &*blks[iblk].ffn_down,
321+
LlamaBlkWeight::FfnGateUp => &*blks[iblk].ffn_gate_up,
322+
LlamaBlkWeight::FfnDown => &*blks[iblk].ffn_down,
320323
_ => unreachable!(),
321324
};
322325
let one = w.len() / nexp;

‎models/llama/common/src/compute.rs‎

Lines changed: 15 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::{args::Args, LlamaMeta};
1+
use super::{args::Args, LlamaBlkWeight, LlamaMeta};
22
use gguf::ggml_quants::{
33
digit_layout::{types as ty, DigitLayout},
44
f16,
@@ -53,18 +53,6 @@ pub trait Operators {
5353
}
5454
}
5555

56-
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
57-
pub enum BlkWeight {
58-
AttnNorm,
59-
AttnQKV,
60-
AttnQKVBias,
61-
AttnO,
62-
FfnNorm,
63-
FfnGateInp,
64-
FfnGateUp,
65-
FfnDown,
66-
}
67-
6856
pub trait WeightLoader {
6957
type Hardware: Hardware;
7058
type Weight<'s>: Deref<Target = [ByteOf<Self::Hardware>]> + 's
@@ -73,14 +61,14 @@ pub trait WeightLoader {
7361

7462
fn load_blk<'a>(
7563
&'a self,
76-
which: BlkWeight,
64+
which: LlamaBlkWeight,
7765
iblk: usize,
7866
queue: &'a QueueOf<Self::Hardware>,
7967
) -> Self::Weight<'a>;
8068

8169
fn load_moe<'a>(
8270
&'a self,
83-
which: BlkWeight,
71+
which: LlamaBlkWeight,
8472
iblk: usize,
8573
iexp: usize,
8674
queue: &'a QueueOf<Self::Hardware>,
@@ -638,7 +626,7 @@ impl<W: WeightLoader> WeightDecorator<W> {
638626
iblk: usize,
639627
queue: &'a QueueOf<W::Hardware>,
640628
) -> Tensor<W::Weight<'a>> {
641-
let w = self.weights.load_blk(BlkWeight::AttnNorm, iblk, queue);
629+
let w = self.weights.load_blk(LlamaBlkWeight::AttnNorm, iblk, queue);
642630
self.norm.clone().map(|_| w)
643631
}
644632

@@ -648,7 +636,7 @@ impl<W: WeightLoader> WeightDecorator<W> {
648636
iblk: usize,
649637
queue: &'a QueueOf<W::Hardware>,
650638
) -> Tensor<W::Weight<'a>> {
651-
let w = self.weights.load_blk(BlkWeight::AttnQKV, iblk, queue);
639+
let w = self.weights.load_blk(LlamaBlkWeight::AttnQKV, iblk, queue);
652640
self.attn_qkv.clone().map(|_| w)
653641
}
654642

@@ -658,7 +646,9 @@ impl<W: WeightLoader> WeightDecorator<W> {
658646
iblk: usize,
659647
queue: &'a QueueOf<W::Hardware>,
660648
) -> Tensor<W::Weight<'a>> {
661-
let w = self.weights.load_blk(BlkWeight::AttnQKVBias, iblk, queue);
649+
let w = self
650+
.weights
651+
.load_blk(LlamaBlkWeight::AttnQKVBias, iblk, queue);
662652
self.attn_qkv_bias.clone().map(|_| w)
663653
}
664654

@@ -668,7 +658,7 @@ impl<W: WeightLoader> WeightDecorator<W> {
668658
iblk: usize,
669659
queue: &'a QueueOf<W::Hardware>,
670660
) -> Tensor<W::Weight<'a>> {
671-
let w = self.weights.load_blk(BlkWeight::AttnO, iblk, queue);
661+
let w = self.weights.load_blk(LlamaBlkWeight::AttnO, iblk, queue);
672662
self.attn_o.clone().map(|_| w)
673663
}
674664

@@ -678,7 +668,7 @@ impl<W: WeightLoader> WeightDecorator<W> {
678668
iblk: usize,
679669
queue: &'a QueueOf<W::Hardware>,
680670
) -> Tensor<W::Weight<'a>> {
681-
let w = self.weights.load_blk(BlkWeight::FfnNorm, iblk, queue);
671+
let w = self.weights.load_blk(LlamaBlkWeight::FfnNorm, iblk, queue);
682672
self.norm.clone().map(|_| w)
683673
}
684674

@@ -688,7 +678,9 @@ impl<W: WeightLoader> WeightDecorator<W> {
688678
iblk: usize,
689679
queue: &'a QueueOf<W::Hardware>,
690680
) -> Tensor<W::Weight<'a>> {
691-
let w = self.weights.load_blk(BlkWeight::FfnGateInp, iblk, queue);
681+
let w = self
682+
.weights
683+
.load_blk(LlamaBlkWeight::FfnGateInp, iblk, queue);
692684
self.ffn_gate_inp.clone().map(|_| w)
693685
}
694686

@@ -699,7 +691,7 @@ impl<W: WeightLoader> WeightDecorator<W> {
699691
iexp: usize,
700692
queue: &'a QueueOf<W::Hardware>,
701693
) -> Tensor<W::Weight<'a>> {
702-
const WHICH: BlkWeight = BlkWeight::FfnGateUp;
694+
const WHICH: LlamaBlkWeight = LlamaBlkWeight::FfnGateUp;
703695
let w = if self.is_moe {
704696
self.weights.load_moe(WHICH, iblk, iexp, queue)
705697
} else {
@@ -715,7 +707,7 @@ impl<W: WeightLoader> WeightDecorator<W> {
715707
iexp: usize,
716708
queue: &'a QueueOf<W::Hardware>,
717709
) -> Tensor<W::Weight<'a>> {
718-
const WHICH: BlkWeight = BlkWeight::FfnDown;
710+
const WHICH: LlamaBlkWeight = LlamaBlkWeight::FfnDown;
719711
let w = if self.is_moe {
720712
self.weights.load_moe(WHICH, iblk, iexp, queue)
721713
} else {

0 commit comments

Comments
 (0)