Method Model

src/model.cpp:756–872

}

Model::Model(YALMData& yalm, int context) {
  config = std::make_shared<Config>();
  config->from_yalm(yalm, context);
  std::cout << "loading model with quant: " << quant_to_string(config->weight_quant) << std::endl;

  bool need_weight_scales = config->weight_quant == Quant::F8E5M2;
  bool use_mla = config->use_mla;
  int b0 = config->block_size[0];
  int b1 = config->block_size[1];

  token_embedding_table = check_tensor(
    get_tensor(yalm, "model.embed.weight"),
    config->weight_quant,
    {config->vocab_size, config->dim, 0, 0},
    __LINE__
  );
  if (need_weight_scales) {
    token_embedding_scale = check_tensor(
      get_tensor(yalm, "model.embed.scale"),
      Quant::F32,
      {cdiv(config->vocab_size, b0), cdiv(config->dim, b1), 0, 0},
      __LINE__
    );
  }

  for (int i = 0; i < config->n_layers; ++i) {
    const Tensor* p_rms_att_weight = get_tensor(yalm, fmt::format("model.layers.{}.attn.norm.weight", i));
    const Tensor* p_rms_ffn_weight = get_tensor(yalm, fmt::format("model.layers.{}.mlp.norm.weight", i));
    const Tensor* p_w1 = get_tensor(yalm, fmt::format("model.layers.{}.mlp.w1.weight", i));
    const Tensor* p_s1 = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.mlp.w1.scale", i)) : nullptr;
    const Tensor* p_w2 = get_tensor(yalm, fmt::format("model.layers.{}.mlp.w2.weight", i));
    const Tensor* p_s2 = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.mlp.w2.scale", i)) : nullptr;
    const Tensor* p_w3 = get_tensor(yalm, fmt::format("model.layers.{}.mlp.w3.weight", i));
    const Tensor* p_s3 = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.mlp.w3.scale", i)) : nullptr;
    const Tensor* p_shared_w1 = (i >= config->first_k_dense_replace && config->n_shared_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.shared_mlp.w1.weight", i)) : nullptr;
    const Tensor* p_shared_s1 = (need_weight_scales && i >= config->first_k_dense_replace && config->n_shared_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.shared_mlp.w1.scale", i)) : nullptr;
    const Tensor* p_shared_w2 = (i >= config->first_k_dense_replace && config->n_shared_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.shared_mlp.w2.weight", i)) : nullptr;
    const Tensor* p_shared_s2 = (need_weight_scales && i >= config->first_k_dense_replace && config->n_shared_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.shared_mlp.w2.scale", i)) : nullptr;
    const Tensor* p_shared_w3 = (i >= config->first_k_dense_replace && config->n_shared_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.shared_mlp.w3.weight", i)) : nullptr;
    const Tensor* p_shared_s3 = (need_weight_scales && i >= config->first_k_dense_replace && config->n_shared_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.shared_mlp.w3.scale", i)) : nullptr;
    const Tensor* p_moegate = (i >= config->first_k_dense_replace && config->n_routed_experts > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.moegate.weight", i)) : nullptr;
    const Tensor* p_moegate_bias = (i >= config->first_k_dense_replace && config->n_routed_experts > 0 && config->has_moegate_bias) ? get_tensor(yalm, fmt::format("model.layers.{}.moegate.bias", i)) : nullptr;

    const Tensor* p_rms_q_a_weight = config->q_lora_rank > 0 ? get_tensor(yalm, fmt::format("model.layers.{}.attn.q_a_norm.weight", i)) : nullptr;
    const Tensor* p_rms_kv_a_weight = get_tensor(yalm, fmt::format("model.layers.{}.attn.kv_a_norm.weight", i));
    const Tensor* p_wq_a = config->q_lora_rank > 0 ? get_tensor(yalm, fmt::format("model.layers.{}.attn.wq_a.weight", i)) : nullptr;
    const Tensor* p_sq_a = (need_weight_scales && config->q_lora_rank > 0) ? get_tensor(yalm, fmt::format("model.layers.{}.attn.wq_a.scale", i)) : nullptr;
    const Tensor* p_wkv_a = get_tensor(yalm, fmt::format("model.layers.{}.attn.wkv_a.weight", i));
    const Tensor* p_skv_a = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.attn.wkv_a.scale", i)) : nullptr;
    const Tensor* p_wo = get_tensor(yalm, fmt::format("model.layers.{}.attn.wo.weight", i));
    const Tensor* p_so = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.attn.wo.scale", i)) : nullptr;


    if (use_mla) {
      const Tensor* p_wc = get_tensor(yalm, fmt::format("model.layers.{}.attn.wc.weight", i));
      const Tensor* p_sc = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.attn.wc.scale", i)) : nullptr;
      const Tensor* p_wq_rope_b = get_tensor(yalm, fmt::format("model.layers.{}.attn.wq_rope_b.weight", i));
      const Tensor* p_sq_rope_b = need_weight_scales ? get_tensor(yalm, fmt::format("model.layers.{}.attn.wq_rope_b.scale", i)) : nullptr;