Skip to content

Commit a626cdb

Browse files
Ceng23333Ceng23333
committed
feat(nn): load_external张量支持多对一映射
Signed-off-by: Ceng <441651826@qq.com>
1 parent fa4ec06 commit a626cdb

22 files changed

Lines changed: 457 additions & 149 deletions

‎1_nn/src/ctx/graph.rs‎

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ struct Tensor_<T> {
7272
external: Option<T>,
7373
}
7474

75-
impl<T> Context<T> {
75+
impl<T: Clone> Context<T> {
7676
pub fn path(&self) -> String {
7777
self.0.borrow().namespace.top().path().to_string()
7878
}
@@ -97,23 +97,28 @@ impl<T> Context<T> {
9797
dt: DigitLayout,
9898
shape: impl IntoIterator<Item = Dim>,
9999
item: T,
100-
) -> Tensor<T> {
100+
) -> Result<Vec<Tensor<T>>, NNError> {
101101
let mut internal = self.0.borrow_mut();
102102

103103
let top = internal.namespace.top_mut();
104104
assert!(top.tensor.check(&name));
105105
let name = format!("{}.{}", top.path(), name);
106106

107-
let idx = internal.tensors.len();
108-
internal.tensors.push(Tensor_ {
109-
name,
110-
meta: TensorMeta::new(dt, shape),
111-
external: Some(item),
112-
});
113-
Tensor {
114-
idx,
115-
ctx: Context(self.0.clone()),
107+
let external_meta = TensorMeta::load_external(dt, shape);
108+
let mut tensors = Vec::with_capacity(external_meta.len());
109+
for meta in external_meta {
110+
let idx = internal.tensors.len();
111+
internal.tensors.push(Tensor_ {
112+
name: name.clone(),
113+
meta,
114+
external: Some(item.clone()),
115+
});
116+
tensors.push(Tensor {
117+
idx,
118+
ctx: Context(self.0.clone()),
119+
});
116120
}
121+
Ok(tensors)
117122
}
118123

119124
pub fn bind_external(&mut self, tensor: Tensor<T>, item: T) {

‎1_nn/src/ctx/tensor.rs‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,29 @@ pub struct TensorMeta {
4141

4242
impl TensorMeta {
4343
pub fn new(dt: DigitLayout, shape: impl IntoIterator<Item = Dim>) -> Self {
44+
let shape = shape.into_iter().collect::<Box<_>>();
45+
Self { dt, shape }
46+
}
47+
48+
pub fn load_external(dt: DigitLayout, shape: impl IntoIterator<Item = Dim>) -> Vec<Self> {
4449
let mut shape = shape.into_iter().collect::<Box<_>>();
4550
let group = dt.group_size();
4651
if group > 1 {
4752
if let Some(dim) = shape.last_mut() {
4853
*dim = std::mem::replace(dim, Dim::from(0)) / group
4954
}
5055
}
51-
Self { dt, shape }
56+
match dt.to_string().as_str() {
57+
// TODO: 量化类型构图时拆分多个tensor
58+
"q4k" | "q6k" => vec![
59+
Self {
60+
dt,
61+
shape: shape.clone(),
62+
},
63+
Self { dt, shape },
64+
],
65+
_ => vec![Self { dt, shape }],
66+
}
5267
}
5368

5469
#[inline]

‎1_nn/src/nn/activation.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ pub enum Activation {
77
GeLU,
88
}
99

10-
impl<T> NuralNetwork<T> for Activation {
10+
impl<T: Clone> NuralNetwork<T> for Activation {
1111
fn launch(
1212
self,
1313
inputs: impl IntoIterator<Item = Tensor<T>>,

‎1_nn/src/nn/attention.rs‎

Lines changed: 55 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,29 @@
11
use super::{Context, Distribution, Linear, NNError, NuralNetwork, TPTensor, Tensor, macros::*};
22
use crate::{
33
Arg, TPAction,
4-
weight_types::{AttnQKV, RowTPWeight},
4+
weight_types::{AttnQKV, ColumnTPWeight, RowTPWeight},
55
};
66
use tensor::digit_layout::types;
77

88
#[derive(Clone)]
9-
pub struct Attention<T> {
9+
pub struct Attention<T: Clone> {
1010
pub nh: usize,
1111
pub nkvh: usize,
12-
pub qkv: Linear<T>,
12+
pub qkv: QKVFormat<T>,
1313
pub rope: Option<RoPE<T>>,
1414
pub output: Linear<T>,
1515
}
1616

17+
#[derive(Clone)]
18+
pub enum QKVFormat<T: Clone> {
19+
Combined(Linear<T>),
20+
Separated {
21+
q: Linear<T>,
22+
k: Linear<T>,
23+
v: Linear<T>,
24+
},
25+
}
26+
1727
#[derive(Clone)]
1828
pub struct RoPE<T> {
1929
pub multimodal: bool,
@@ -22,7 +32,7 @@ pub struct RoPE<T> {
2232
pub cos: T,
2333
}
2434

25-
impl<T> Attention<T> {
35+
impl<T: Clone> Attention<T> {
2636
pub fn tensor_parallel(self, dist: Distribution) -> Attention<TPTensor<T>> {
2737
let Self {
2838
nh,
@@ -36,7 +46,16 @@ impl<T> Attention<T> {
3646
Attention {
3747
nh: nh / dist.total * dist.len,
3848
nkvh: nkvh / dist.total * dist.len,
39-
qkv: qkv.parallel(TPAction::new(AttnQKV(nh / nkvh), dist)),
49+
qkv: match qkv {
50+
QKVFormat::Combined(qkv) => {
51+
QKVFormat::Combined(qkv.parallel(TPAction::new(AttnQKV(nh / nkvh), dist)))
52+
}
53+
QKVFormat::Separated { q, k, v } => QKVFormat::Separated {
54+
q: q.parallel(TPAction::new(ColumnTPWeight, dist)),
55+
k: k.parallel(TPAction::new(ColumnTPWeight, dist)),
56+
v: v.parallel(TPAction::new(ColumnTPWeight, dist)),
57+
},
58+
},
4059
rope: rope.map(
4160
|RoPE {
4261
multimodal,
@@ -55,7 +74,7 @@ impl<T> Attention<T> {
5574
}
5675
}
5776

58-
impl<T> NuralNetwork<T> for Attention<T> {
77+
impl<T: Clone> NuralNetwork<T> for Attention<T> {
5978
fn launch(
6079
self,
6180
inputs: impl IntoIterator<Item = Tensor<T>>,
@@ -71,24 +90,35 @@ impl<T> NuralNetwork<T> for Attention<T> {
7190
output,
7291
} = self;
7392

74-
destruct!([x] = ctx.trap("attn-qkv", qkv, [x])?);
75-
dims!([_, dqkv] = x);
76-
let dh = dqkv.clone() / (nh + nkvh + nkvh);
93+
dims!([_, d] = x);
94+
let dh = d.clone() / nh;
7795

78-
destruct!(
79-
[q, k, v] = ctx.call(
80-
"split-qkv",
81-
"split",
82-
Some(Arg::dict([
83-
("axis".into(), Arg::int(1)),
84-
(
85-
"parts".into(),
86-
Arg::arr([Arg::dim(nh), Arg::dim(nkvh), Arg::dim(nkvh)])
87-
)
88-
])),
89-
[x],
90-
)?
91-
);
96+
let [q, k, v] = match qkv {
97+
QKVFormat::Combined(qkv) => {
98+
destruct!([x] = ctx.trap("attn-qkv", qkv, [x])?);
99+
destruct!(
100+
[q, k, v] = ctx.call(
101+
"split-qkv",
102+
"split",
103+
Some(Arg::dict([
104+
("axis".into(), Arg::int(1)),
105+
(
106+
"parts".into(),
107+
Arg::arr([Arg::dim(nh), Arg::dim(nkvh), Arg::dim(nkvh)])
108+
)
109+
])),
110+
[x],
111+
)?
112+
);
113+
[q, k, v]
114+
}
115+
QKVFormat::Separated { q, k, v } => {
116+
destruct!([q] = ctx.trap("attn-q", q, [x.clone()])?);
117+
destruct!([k] = ctx.trap("attn-k", k, [x.clone()])?);
118+
destruct!([v] = ctx.trap("attn-v", v, [x])?);
119+
[q, k, v]
120+
}
121+
};
92122

93123
let [q, k] = match rope {
94124
Some(RoPE {
@@ -98,9 +128,8 @@ impl<T> NuralNetwork<T> for Attention<T> {
98128
cos,
99129
}) => {
100130
let shape = [nctx.into(), dh.clone() / 2];
101-
let sin = ctx.load_external("rope.sin", types::F32, shape.clone(), sin);
102-
let cos = ctx.load_external("rope.cos", types::F32, shape, cos);
103-
131+
destruct!([sin] = ctx.load_external("rope.sin", types::F32, shape.clone(), sin)?);
132+
destruct!([cos] = ctx.load_external("rope.cos", types::F32, shape, cos)?);
104133
let op = if multimodal { "mrope" } else { "rope" };
105134
destruct!(
106135
[q_] = ctx.call(

‎1_nn/src/nn/embedding.rs‎

Lines changed: 33 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1-
use super::{Context, NNError, NuralNetwork, TPTensor, Tensor};
1+
use crate::macros::destruct;
2+
3+
use super::{Context, NNError, NuralNetwork, TPTensor, Tensor};
24
use tensor::digit_layout::DigitLayout;
35

46
#[derive(Clone)]
5-
pub struct Embedding<T> {
7+
pub struct Embedding<T: Clone> {
68
pub dt: DigitLayout,
79
pub d: usize,
810
pub wte: Table<T>,
@@ -15,7 +17,7 @@ pub struct Table<T> {
1517
pub weight: T,
1618
}
1719

18-
impl<T> Embedding<T> {
20+
impl<T: Clone> Embedding<T> {
1921
pub fn tensor_parallel(self) -> Embedding<TPTensor<T>> {
2022
let Self { dt, d, wte, wpe } = self;
2123
Embedding {
@@ -33,7 +35,7 @@ impl<T> Embedding<T> {
3335
}
3436
}
3537

36-
impl<T> NuralNetwork<T> for Embedding<T> {
38+
impl<T: Clone> NuralNetwork<T> for Embedding<T> {
3739
fn launch(
3840
self,
3941
inputs: impl IntoIterator<Item = Tensor<T>>,
@@ -43,19 +45,36 @@ impl<T> NuralNetwork<T> for Embedding<T> {
4345
let mut inputs = inputs.into_iter();
4446

4547
let Table { row, weight } = wte;
46-
let wte = ctx.load_external("wte", dt, [row.into(), d.into()], weight);
48+
4749
let tokens = inputs.next().unwrap();
4850

49-
let outputs = match wpe {
50-
Some(wpe) => {
51-
let Table { row, weight } = wpe;
52-
let wpe = ctx.load_external("wpe", dt, [row.into(), d.into()], weight);
53-
let pos = inputs.next().unwrap();
54-
ctx.call("", "embedding", None, [wte, tokens, wpe, pos])
51+
let outputs = if dt.group_size() > 1 {
52+
let w = ctx.load_external("weight", dt, [row.into(), d.into()], weight)?;
53+
match wpe {
54+
Some(_) => {
55+
todo!()
56+
}
57+
None => {
58+
let inputs = w.into_iter().chain([tokens]).collect::<Vec<_>>();
59+
ctx.call("", "quant-embedding", Some(false.into()), inputs)
60+
}
5561
}
56-
None => {
57-
// format
58-
ctx.call("", "embedding", None, [wte, tokens])
62+
} else {
63+
destruct!([wte] = ctx.load_external("wte", dt, [row.into(), d.into()], weight)?);
64+
65+
match wpe {
66+
Some(wpe) => {
67+
let Table { row, weight } = wpe;
68+
destruct!(
69+
[wpe] = ctx.load_external("wpe", dt, [row.into(), d.into()], weight)?
70+
);
71+
let pos = inputs.next().unwrap();
72+
ctx.call("", "embedding", None, [wte, tokens, wpe, pos])
73+
}
74+
None => {
75+
// format
76+
ctx.call("", "embedding", None, [wte, tokens])
77+
}
5978
}
6079
};
6180

0 commit comments

Comments
 (0)