定义一个模块

class Mlp(nn.Module):
  def __init__(self, hidden: int, inter: int):
    super().__init__()
    self.gate_proj = nn.Linear(hidden, inter, bias=False)
    self.up_proj = nn.Linear(hidden, inter, bias=False)
    self.down_proj = nn.Linear(inter, hidden, bias=False)

  def forward(self, x):
    return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class FfnHalf(nn.Module):
  def __init__(self, hidden: int, inter: int, eps: float):
    super().__init__()
    self.post_attention_layernorm = RmsNorm(hidden, eps)  # weight * x * rsqrt(mean(x²) + eps)
    self.mlp = Mlp(hidden, inter)

  def forward(self, x):
    return x + self.mlp(self.post_attention_layernorm(x))


block = FfnHalf(576, 1536, 1e-5)
prefix = "model.layers.0."
sub = {k[len(prefix):]: v for k, v in model.state_dict().items() if k.startswith(prefix)}
missing, unexpected = block.load_state_dict(sub, strict=False)

use candle_core::{Module, Tensor};
use candle_nn::{linear_no_bias, rms_norm, Linear, RmsNorm, VarBuilder};

struct Mlp {
  gate_proj: Linear,
  up_proj: Linear,
  down_proj: Linear,
}

impl Mlp {
  fn load(vb: VarBuilder, hidden: usize, inter: usize) -> candle_core::Result<Self> {
    Ok(Self {
      gate_proj: linear_no_bias(hidden, inter, vb.pp("gate_proj"))?,
      up_proj: linear_no_bias(hidden, inter, vb.pp("up_proj"))?,
      down_proj: linear_no_bias(inter, hidden, vb.pp("down_proj"))?,
    })
  }
}

impl Module for Mlp {
  fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
    let gate = x.apply(&self.gate_proj)?.silu()?;
    (gate * x.apply(&self.up_proj)?)?.apply(&self.down_proj)
  }
}

struct FfnHalf {
  post_attention_layernorm: RmsNorm,
  mlp: Mlp,
}

impl FfnHalf {
  fn load(vb: VarBuilder, hidden: usize, inter: usize, eps: f64) -> candle_core::Result<Self> {
    Ok(Self {
      post_attention_layernorm: rms_norm(hidden, eps, vb.pp("post_attention_layernorm"))?,
      mlp: Mlp::load(vb.pp("mlp"), hidden, inter)?,
    })
  }
}

impl Module for FfnHalf {
  fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
    x + x.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?
  }
}

let model_dir = "SmolLM2-135M-Instruct";
let path = format!("{model_dir}/model.safetensors");
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[&path], DType::F32, &dev)? };
let block = FfnHalf::load(vb.pp("model.layers.0"), 576, 1536, 1e-5)?;
let y = x.apply(&block)?;

missing []
unexpected ['self_attn.q_proj.weight', 'self_attn.k_proj.weight', 'self_attn.v_proj.weight', 'self_attn.o_proj.weight', 'input_layernorm.weight']
hand block vs transformers modules max diff 0.0

[Cpu] FfnHalf out [1, 5, 576] max_abs_diff=4.196167e-5 (max |ref|=53.606655)
[Metal(MetalDevice(DeviceId(1)))] FfnHalf out [1, 5, 576] max_abs_diff=6.67572e-6 (max |ref|=53.606655)

f32 vs f64: norm 5.494803190231323e-08 mlp 4.567097710150847e-05 max|mlp| 53.762926948947964

[Cpu] norm vs pt32 8.34465e-7 vs pt64 8.34465e-7; mlp vs pt32 4.196167e-5 vs pt64 7.6293945e-5
[Metal(MetalDevice(DeviceId(1)))] norm vs pt32 1.1920929e-7 vs pt64 1.1920929e-7; mlp vs pt32 6.67572e-6 vs pt64 4.9591064e-5

print("block keys:", [(k, list(v.shape)) for k, v in block.state_dict().items()])

let vm = VarMap::new();
let _ = FfnHalf::load(VarBuilder::from_varmap(&vm, DType::F32, &Device::Cpu).pp("model.layers.0"), 576, 1536, 1e-5)?;
let mut names: Vec<(String, Vec<usize>)> =
  vm.data().lock().map_err(|e| anyhow::anyhow!("{e}"))?.iter().map(|(k, v)| (k.clone(), v.dims().to_vec())).collect();
names.sort();

block keys: [('post_attention_layernorm.weight', [576]), ('mlp.gate_proj.weight', [1536, 576]), ('mlp.up_proj.weight', [1536, 576]), ('mlp.down_proj.weight', [576, 1536])]

requested by FfnHalf::load: [("model.layers.0.mlp.down_proj.weight", [576, 1536]), ("model.layers.0.mlp.gate_proj.weight", [1536, 576]), ("model.layers.0.mlp.up_proj.weight", [1536, 576]), ("model.layers.0.post_attention_layernorm.weight", [576])]

Llama::load asked 272 names; state_dict has 273; only in state_dict ["lm_head.weight"]; only in candle []

model.state_dict() keys 273 lm_head.weight in it True same storage as embed True

_IncompatibleKeys(missing_keys=['mlps.gate_proj.weight', 'mlps.up_proj.weight', 'mlps.down_proj.weight'], unexpected_keys=['mlp.gate_proj.weight', 'mlp.up_proj.weight', 'mlp.down_proj.weight'])

match linear_no_bias(576, 1536, vb.pp("mlps").pp("gate_proj")) {
  Ok(_) => println!("loaded"),
  Err(e) => println!("load error: {e}"),
}
match linear_no_bias(1536, 576, vb.pp("mlp").pp("gate_proj")) {
  Ok(_) => println!("loaded"),
  Err(e) => println!("load error: {e}"),
}

load error: cannot find tensor model.layers.0.mlps.gate_proj.weight
load error: shape mismatch for model.layers.0.mlp.gate_proj.weight, expected: [576, 1536], got: [1536, 576]

named_parameters: ['post_attention_layernorm.weight', 'mlp.gate_proj.weight', 'mlp.up_proj.weight', 'mlp.down_proj.weight']

net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
print("state_dict keys:", list(net.state_dict().keys()))

let net = seq()
  .add(linear(4, 8, vb.pp("0"))?)
  .add(Activation::Relu)
  .add(linear(8, 3, vb.pp("2"))?);
let y = net.forward(x)?;

state_dict keys: ['0.weight', '0.bias', '2.weight', '2.bias']
[[0.08920177072286606, -0.008306492120027542, 0.09059013426303864], [0.19197271764278412, -0.08750583976507187, 0.07556053251028061]]

[[0.08920177, -0.008306492, 0.090590134], [0.19197272, -0.08750584, 0.07556053]]
max_abs_diff=0e0
pp("1"): cannot find tensor 1.weight