tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

use linear_b

Browse Source
This commit is contained in:
jhqxxx
2026-01-15 23:17:34 +08:00
parent d9b803d27e
commit f4e95e3937
5 changed files with 91 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/common/mod.rs
+20 -61
View File
@@ -3,8 +3,8 @@ use candle_core::{D, Tensor};
use candle_nn::{ use candle_nn::{
Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, Embedding, Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, Embedding,
LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, VarBuilder, batch_norm, conv1d, LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, VarBuilder, batch_norm, conv1d,
conv1d_no_bias, conv2d, conv2d_no_bias, embedding, layer_norm, linear, linear_no_bias, conv1d_no_bias, conv2d, conv2d_no_bias, embedding, layer_norm, linear_b,
rms_norm, linear_no_bias, rms_norm,
}; };
use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator}; use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
@@ -29,19 +29,9 @@ impl GateUpDownMLP {
act_fn: Activation, act_fn: Activation,
bias: bool, bias: bool,
) -> Result<Self> { ) -> Result<Self> {
let (gate_proj, up_proj, down_proj) = if bias { let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("gate_proj"))?;
( let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("up_proj"))?;
linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?, let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp("down_proj"))?;
linear(hidden_size, intermediate_size, vb.pp("up_proj"))?,
linear(intermediate_size, hidden_size, vb.pp("down_proj"))?,
)
} else {
(
linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?,
linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?,
)
};
Ok(Self { Ok(Self {
gate_proj, gate_proj,
up_proj, up_proj,
@@ -78,17 +68,9 @@ impl TwoLinearMLP {
linear1_pp_name: &str, linear1_pp_name: &str,
linear2_pp_name: &str, linear2_pp_name: &str,
) -> Result<Self> { ) -> Result<Self> {
let (linear1, linear2) = if bias { let linear1 = linear_b(in_dim, middle_dim, bias, vb.pp(linear1_pp_name))?;
( let linear2 = linear_b(middle_dim, out_dim, bias, vb.pp(linear2_pp_name))?;
linear(in_dim, middle_dim, vb.pp(linear1_pp_name))?,
linear(middle_dim, out_dim, vb.pp(linear2_pp_name))?,
)
} else {
(
linear_no_bias(in_dim, middle_dim, vb.pp(linear1_pp_name))?,
linear_no_bias(middle_dim, out_dim, vb.pp(linear2_pp_name))?,
)
};
Ok(Self { Ok(Self {
linear1, linear1,
linear2, linear2,
@@ -141,53 +123,30 @@ impl NaiveAttention {
let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj"); let k_proj_pp_name = k_proj_pp_name.unwrap_or("k_proj");
let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj"); let v_proj_pp_name = v_proj_pp_name.unwrap_or("v_proj");
let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj"); let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj");
let (q_proj, k_proj, v_proj, o_proj) = if bias { let q_proj = linear_b(
(
linear(
hidden_size, hidden_size,
num_attention_heads * head_dim, num_attention_heads * head_dim,
bias,
vb.pp(q_proj_pp_name), vb.pp(q_proj_pp_name),
)?, )?;
linear( let k_proj = linear_b(
hidden_size, hidden_size,
num_key_value_heads * head_dim, num_key_value_heads * head_dim,
bias,
vb.pp(k_proj_pp_name), vb.pp(k_proj_pp_name),
)?, )?;
linear( let v_proj = linear_b(
hidden_size, hidden_size,
num_key_value_heads * head_dim, num_key_value_heads * head_dim,
bias,
vb.pp(v_proj_pp_name), vb.pp(v_proj_pp_name),
)?, )?;
linear( let o_proj = linear_b(
num_attention_heads * head_dim, num_attention_heads * head_dim,
hidden_size, hidden_size,
bias,
vb.pp(o_proj_pp_name), vb.pp(o_proj_pp_name),
)?, )?;
)
} else {
(
linear_no_bias(
hidden_size,
num_attention_heads * head_dim,
vb.pp(q_proj_pp_name),
)?,
linear_no_bias(
hidden_size,
num_key_value_heads * head_dim,
vb.pp(k_proj_pp_name),
)?,
linear_no_bias(
hidden_size,
num_key_value_heads * head_dim,
vb.pp(v_proj_pp_name),
)?,
linear_no_bias(
num_attention_heads * head_dim,
hidden_size,
vb.pp(o_proj_pp_name),
)?,
)
};
Ok(Self { Ok(Self {
q_proj, q_proj,
src/models/deepseek_ocr/model.rs
+3 -9
View File
@@ -1,10 +1,7 @@
use anyhow::Result; use anyhow::Result;
use candle_core::{D, IndexOp, Tensor}; use candle_core::{D, IndexOp, Tensor};
use candle_nn::{ use candle_nn::{
Activation, Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, Activation, Conv2d, Embedding, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b, linear_no_bias, ops::{sigmoid, softmax}, rms_norm
linear, linear_no_bias,
ops::{sigmoid, softmax},
rms_norm,
}; };
use candle_transformers::models::segment_anything::LayerNorm2d; use candle_transformers::models::segment_anything::LayerNorm2d;
@@ -79,11 +76,8 @@ impl Attention {
) -> Result<Self> { ) -> Result<Self> {
let head_dim = dim / num_heads; let head_dim = dim / num_heads;
let scaling = 1.0 / (head_dim as f64).sqrt(); let scaling = 1.0 / (head_dim as f64).sqrt();
let qkv = if qkv_bias { let qkv = linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?;
linear(dim, dim * 3, vb.pp("qkv"))?
} else {
linear_no_bias(dim, dim * 3, vb.pp("qkv"))?
};
let proj = linear(dim, dim, vb.pp("proj"))?; let proj = linear(dim, dim, vb.pp("proj"))?;
let mut rel_pos_h = None; let mut rel_pos_h = None;
let mut rel_pos_w = None; let mut rel_pos_w = None;
src/models/hunyuan_ocr/model.rs
+27 -19
View File
@@ -1,8 +1,8 @@
use anyhow::{Result, anyhow}; use anyhow::{Result, anyhow};
use candle_core::{D, IndexOp, Tensor}; use candle_core::{D, IndexOp, Tensor};
use candle_nn::{ use candle_nn::{
Conv2d, Embedding, Init, Linear, Module, RmsNorm, VarBuilder, embedding, linear, Conv2d, Embedding, Init, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b,
linear_no_bias, rms_norm, rms_norm,
}; };
use crate::{ use crate::{
@@ -284,23 +284,31 @@ impl HunYuanVLAttention {
) -> Result<Self> { ) -> Result<Self> {
let num_kv_groups = num_attention_heads / num_key_value_heads; let num_kv_groups = num_attention_heads / num_key_value_heads;
let scaling = 1f64 / f64::sqrt(head_dim as f64); let scaling = 1f64 / f64::sqrt(head_dim as f64);
let (q_proj, k_proj, v_proj, o_proj) = if attention_bias { let q_proj = linear_b(
let q_proj = linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; hidden_size,
let k_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; num_attention_heads * head_dim,
let v_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; attention_bias,
let o_proj = linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?; vb.pp("q_proj"),
(q_proj, k_proj, v_proj, o_proj) )?;
} else { let k_proj = linear_b(
let q_proj = hidden_size,
linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; num_key_value_heads * head_dim,
let k_proj = attention_bias,
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; vb.pp("k_proj"),
let v_proj = )?;
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; let v_proj = linear_b(
let o_proj = hidden_size,
linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?; num_key_value_heads * head_dim,
(q_proj, k_proj, v_proj, o_proj) attention_bias,
}; vb.pp("v_proj"),
)?;
let o_proj = linear_b(
num_attention_heads * head_dim,
hidden_size,
attention_bias,
vb.pp("o_proj"),
)?;
let query_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("query_layernorm"))?; let query_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("query_layernorm"))?;
let key_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("key_layernorm"))?; let key_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("key_layernorm"))?;
Ok(Self { Ok(Self {
src/models/qwen3/model.rs
+26 -18
View File
@@ -1,7 +1,8 @@
use anyhow::Result; use anyhow::Result;
use candle_core::Tensor; use candle_core::Tensor;
use candle_nn::{ use candle_nn::{
Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_no_bias, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, linear_b, linear_no_bias,
rms_norm,
}; };
use crate::{ use crate::{
@@ -36,23 +37,30 @@ impl Qwen3Attention {
let num_key_value_heads = config.num_key_value_heads; let num_key_value_heads = config.num_key_value_heads;
let num_kv_groups = num_attention_heads / num_key_value_heads; let num_kv_groups = num_attention_heads / num_key_value_heads;
let scaling = 1f64 / f64::sqrt(head_dim as f64); let scaling = 1f64 / f64::sqrt(head_dim as f64);
let (q_proj, k_proj, v_proj, o_proj) = if config.attention_bias { let q_proj = linear_b(
let q_proj = linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; hidden_size,
let k_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; num_attention_heads * head_dim,
let v_proj = linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; config.attention_bias,
let o_proj = linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?; vb.pp("q_proj"),
(q_proj, k_proj, v_proj, o_proj) )?;
} else { let k_proj = linear_b(
let q_proj = hidden_size,
linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?; num_key_value_heads * head_dim,
let k_proj = config.attention_bias,
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?; vb.pp("k_proj"),
let v_proj = )?;
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?; let v_proj = linear_b(
let o_proj = hidden_size,
linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?; num_key_value_heads * head_dim,
(q_proj, k_proj, v_proj, o_proj) config.attention_bias,
}; vb.pp("v_proj"),
)?;
let o_proj = linear_b(
num_attention_heads * head_dim,
hidden_size,
config.attention_bias,
vb.pp("o_proj"),
)?;
let q_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("q_norm"))?; let q_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("q_norm"))?;
let k_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("k_norm"))?; let k_norm = rms_norm(head_dim, config.rms_norm_eps, vb.pp("k_norm"))?;
Ok(Self { Ok(Self {
src/models/rmbg2_0/model.rs
+3 -6
View File
@@ -2,7 +2,7 @@ use anyhow::{Result, anyhow};
use candle_core::{D, DType, Device, IndexOp, Shape, Tensor}; use candle_core::{D, DType, Device, IndexOp, Shape, Tensor};
use candle_nn::{ use candle_nn::{
Activation, BatchNorm, Conv2d, Init, LayerNorm, Linear, Module, ModuleT, VarBuilder, linear, Activation, BatchNorm, Conv2d, Init, LayerNorm, Linear, Module, ModuleT, VarBuilder, linear,
linear_no_bias, ops::sigmoid, linear_b, linear_no_bias, ops::sigmoid,
}; };
use crate::{ use crate::{
@@ -91,11 +91,8 @@ impl WindowAttention {
) -> Result<Self> { ) -> Result<Self> {
let head_dim = dim / num_heads; let head_dim = dim / num_heads;
let scaling = 1.0 / (head_dim as f64).sqrt(); let scaling = 1.0 / (head_dim as f64).sqrt();
let qkv = if qkv_bias { let qkv = linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?;
linear(dim, dim * 3, vb.pp("qkv"))?
} else {
linear_no_bias(dim, dim * 3, vb.pp("qkv"))?
};
let proj = linear(dim, dim, vb.pp("proj"))?; let proj = linear(dim, dim, vb.pp("proj"))?;
let relative_position_bias_table = vb.get_with_hints( let relative_position_bias_table = vb.get_with_hints(
((2 * window_size.0 - 1) * (2 * window_size.1 - 1), num_heads), ((2 * window_size.0 - 1) * (2 * window_size.1 - 1), num_heads),