Skip to content

Commit 48ba441

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

22 files changed

Lines changed: 469 additions & 147 deletions

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

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,13 @@
22
use crate::{Arg, Dim, Edge, NNError, NNGraph, NuralNetwork, ctx::name::Namespace, op::OpError};
33
use graph::{GraphTopo, TopoNode};
44
use mem::{External, Node, Operator};
5-
use std::{cell::RefCell, collections::HashMap, fmt::Display, ops::Range, rc::Rc};
5+
use std::{cell::RefCell, clone::Clone, collections::HashMap, fmt::Display, ops::Range, rc::Rc};
66
use tensor::digit_layout::DigitLayout;
77

88
pub struct Context<T>(Rc<RefCell<Internal<T>>>);
99

1010
impl GraphBuilder {
11-
pub fn build<T, NN: NuralNetwork<T>>(
11+
pub fn build<T: Clone, NN: NuralNetwork<T>>(
1212
&self,
1313
nn: NN,
1414
inputs: impl IntoIterator<Item = TensorMeta>,
@@ -18,7 +18,7 @@ impl GraphBuilder {
1818
Ok(ctx.into_graph(outputs))
1919
}
2020

21-
fn new_context<T>(
21+
fn new_context<T: Clone>(
2222
&self,
2323
global_inputs: impl IntoIterator<Item = TensorMeta>,
2424
) -> (Context<T>, Vec<Tensor<T>>) {
@@ -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) {
@@ -188,7 +193,7 @@ impl<T> Context<T> {
188193
}
189194
}
190195

191-
impl<T> Context<T> {
196+
impl<T: Clone> Context<T> {
192197
pub(super) fn clone(&self) -> Self {
193198
Self(self.0.clone())
194199
}

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

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,16 @@
11
use super::Context;
22
use crate::{NNError, macros::destruct};
33
use arg::{Arg, Dim};
4+
use std::clone::Clone;
45
use tensor::digit_layout::DigitLayout;
56

67
/// 计算图层张量
7-
pub struct Tensor<T> {
8+
pub struct Tensor<T: Clone> {
89
pub(super) idx: usize,
910
pub(super) ctx: Context<T>,
1011
}
1112

12-
impl<T> Clone for Tensor<T> {
13+
impl<T: Clone> Clone for Tensor<T> {
1314
fn clone(&self) -> Self {
1415
Self {
1516
idx: self.idx,
@@ -18,7 +19,7 @@ impl<T> Clone for Tensor<T> {
1819
}
1920
}
2021

21-
impl<T> Tensor<T> {
22+
impl<T: Clone> Tensor<T> {
2223
#[inline]
2324
pub fn dt(&self) -> DigitLayout {
2425
self.meta().dt
@@ -34,7 +35,7 @@ impl<T> Tensor<T> {
3435
}
3536
}
3637

37-
impl<T> Tensor<T> {
38+
impl<T: Clone> Tensor<T> {
3839
pub fn split(
3940
self,
4041
name: impl ToString,
@@ -101,14 +102,29 @@ pub struct TensorMeta {
101102

102103
impl TensorMeta {
103104
pub fn new(dt: DigitLayout, shape: impl IntoIterator<Item = Dim>) -> Self {
105+
let shape = shape.into_iter().collect::<Box<_>>();
106+
Self { dt, shape }
107+
}
108+
109+
pub fn load_external(dt: DigitLayout, shape: impl IntoIterator<Item = Dim>) -> Vec<Self> {
104110
let mut shape = shape.into_iter().collect::<Box<_>>();
105111
let group = dt.group_size();
106112
if group > 1 {
107113
if let Some(dim) = shape.last_mut() {
108114
*dim = std::mem::replace(dim, Dim::from(0)) / group
109115
}
110116
}
111-
Self { dt, shape }
117+
match dt.to_string().as_str() {
118+
// TODO: 量化类型构图时拆分多个tensor
119+
"q4k" | "q6k" => vec![
120+
Self {
121+
dt,
122+
shape: shape.clone(),
123+
},
124+
Self { dt, shape },
125+
],
126+
_ => vec![Self { dt, shape }],
127+
}
112128
}
113129

114130
#[inline]

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

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

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

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

Lines changed: 56 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,22 +3,32 @@
33
macros::*,
44
};
55
use crate::{
6-
TPAction,
7-
weight_types::{AttnQKV, RowTPWeight},
6+
Arg, TPAction,
7+
weight_types::{AttnQKV, ColumnTPWeight, RowTPWeight},
88
};
99
use tensor::digit_layout::types;
1010

1111
#[derive(Clone)]
12-
pub struct Attention<T> {
12+
pub struct Attention<T: Clone> {
1313
pub nh: usize,
1414
pub nkvh: usize,
15-
pub qkv: Linear<T>,
15+
pub qkv: QKVFormat<T>,
1616
pub q_norm: Option<Normalization<T>>,
1717
pub k_norm: Option<Normalization<T>>,
1818
pub rope: Option<RoPE<T>>,
1919
pub output: Linear<T>,
2020
}
2121

22+
#[derive(Clone)]
23+
pub enum QKVFormat<T: Clone> {
24+
Combined(Linear<T>),
25+
Separated {
26+
q: Linear<T>,
27+
k: Linear<T>,
28+
v: Linear<T>,
29+
},
30+
}
31+
2232
#[derive(Clone)]
2333
pub struct RoPE<T> {
2434
pub multimodal: bool,
@@ -27,7 +37,7 @@ pub struct RoPE<T> {
2737
pub cos: T,
2838
}
2939

30-
impl<T> Attention<T> {
40+
impl<T: Clone> Attention<T> {
3141
pub fn tensor_parallel(self, dist: Distribution) -> Attention<TPTensor<T>> {
3242
let Self {
3343
nh,
@@ -43,7 +53,16 @@ impl<T> Attention<T> {
4353
Attention {
4454
nh: nh / dist.total * dist.len,
4555
nkvh: nkvh / dist.total * dist.len,
46-
qkv: qkv.parallel(TPAction::new(AttnQKV(nh / nkvh), dist)),
56+
qkv: match qkv {
57+
QKVFormat::Combined(qkv) => {
58+
QKVFormat::Combined(qkv.parallel(TPAction::new(AttnQKV(nh / nkvh), dist)))
59+
}
60+
QKVFormat::Separated { q, k, v } => QKVFormat::Separated {
61+
q: q.parallel(TPAction::new(ColumnTPWeight, dist)),
62+
k: k.parallel(TPAction::new(ColumnTPWeight, dist)),
63+
v: v.parallel(TPAction::new(ColumnTPWeight, dist)),
64+
},
65+
},
4766
q_norm: q_norm.map(|norm| norm.tensor_parallel()),
4867
k_norm: k_norm.map(|norm| norm.tensor_parallel()),
4968
rope: rope.map(
@@ -64,7 +83,7 @@ impl<T> Attention<T> {
6483
}
6584
}
6685

67-
impl<T> NuralNetwork<T> for Attention<T> {
86+
impl<T: Clone> NuralNetwork<T> for Attention<T> {
6887
fn launch(
6988
self,
7089
inputs: impl IntoIterator<Item = Tensor<T>>,
@@ -81,12 +100,36 @@ impl<T> NuralNetwork<T> for Attention<T> {
81100
rope,
82101
output,
83102
} = self;
84-
destruct!([x] = ctx.trap("attn-qkv", qkv, [x])?);
85-
dims!([_, dqkv] = x);
86-
let dh = dqkv.clone() / (nh + nkvh + nkvh);
87103

88-
destruct!([q, k, v] = x.split("split-qkv", 1, [nh.into(), nkvh.into(), nkvh.into()])?);
104+
dims!([_, d] = x);
105+
let dh = d.clone() / nh;
89106

107+
let [q, k, v] = match qkv {
108+
QKVFormat::Combined(qkv) => {
109+
destruct!([x] = ctx.trap("attn-qkv", qkv, [x])?);
110+
destruct!(
111+
[q, k, v] = ctx.call(
112+
"split-qkv",
113+
"split",
114+
Some(Arg::dict([
115+
("axis".into(), Arg::int(1)),
116+
(
117+
"parts".into(),
118+
Arg::arr([Arg::dim(nh), Arg::dim(nkvh), Arg::dim(nkvh)])
119+
)
120+
])),
121+
[x],
122+
)?
123+
);
124+
[q, k, v]
125+
}
126+
QKVFormat::Separated { q, k, v } => {
127+
destruct!([q] = ctx.trap("attn-q", q, [x.clone()])?);
128+
destruct!([k] = ctx.trap("attn-k", k, [x.clone()])?);
129+
destruct!([v] = ctx.trap("attn-v", v, [x])?);
130+
[q, k, v]
131+
}
132+
};
90133
// Apply normalization to q and k if they exist
91134
let q = match q_norm {
92135
Some(norm) => {
@@ -114,9 +157,8 @@ impl<T> NuralNetwork<T> for Attention<T> {
114157
cos,
115158
}) => {
116159
let shape = [nctx.into(), dh.clone() / 2];
117-
let sin = ctx.load_external("rope.sin", types::F32, shape.clone(), sin);
118-
let cos = ctx.load_external("rope.cos", types::F32, shape, cos);
119-
160+
destruct!([sin] = ctx.load_external("rope.sin", types::F32, shape.clone(), sin)?);
161+
destruct!([cos] = ctx.load_external("rope.cos", types::F32, shape, cos)?);
120162
let op = if multimodal { "mrope" } else { "rope" };
121163
destruct!(
122164
[q_] = ctx.call(

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,14 @@ use super::{
44
};
55

66
#[derive(Clone)]
7-
pub struct CogVLM<T> {
7+
pub struct CogVLM<T: Clone> {
88
pub patch_embd: PatchEmbd<T>,
99
pub vision_blks: Box<[TransformerBlk<T>]>,
1010
pub glu_proj: Mlp<T>,
1111
pub merger: Merger<T>,
1212
}
1313

14-
impl<T> CogVLM<T> {
14+
impl<T: Clone> CogVLM<T> {
1515
pub fn tensor_parallel(self, dist: Distribution) -> CogVLM<TPTensor<T>> {
1616
let Self {
1717
patch_embd,
@@ -31,7 +31,7 @@ impl<T> CogVLM<T> {
3131
}
3232
}
3333

34-
impl<T> NuralNetwork<T> for CogVLM<T> {
34+
impl<T: Clone> NuralNetwork<T> for CogVLM<T> {
3535
fn launch(
3636
self,
3737
inputs: impl IntoIterator<Item = Tensor<T>>,

‎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)