Function decoder_layer

src/models/ace_step/text_encoder.cpp:169–229

TextEncoderWeights load_text_encoder_weights(
    ggml_backend_t backend,
    core::BackendType backend_type,
    const AceStepAssets & assets,
    assets::TensorStorageType storage_type) {
    const auto & config = assets.config.text_encoder;
    const auto & source = *assets.text_encoder_weights;
    const int64_t dim = config.head_dim;

    TextEncoderWeights weights;
    weights.store = std::make_shared<core::BackendWeightStore>(
        backend,
        backend_type,
        "ace_step.text_encoder.weights",
        64ull * 1024ull * 1024ull);

    weights.embed_tokens = weights.store->load_tensor(
        source,
        "embed_tokens.weight",
        storage_type,
        {config.vocab_size, config.hidden_size});

    weights.layers.reserve(static_cast<size_t>(config.num_hidden_layers));
    for (int64_t i = 0; i < config.num_hidden_layers; ++i) {
        const std::string prefix = "layers." + std::to_string(i);
        TextEncoderLayerWeights layer;
        layer.input_norm = weights.store->load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size});
        layer.post_norm =
            weights.store->load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size});
        layer.q_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.q_norm.weight", {dim});
        layer.k_norm = weights.store->load_f32_tensor(source, prefix + ".self_attn.k_norm.weight", {dim});
        layer.q_proj = weights.store->load_tensor(
            source,
            prefix + ".self_attn.q_proj.weight",
            storage_type,
            {config.num_attention_heads * dim, config.hidden_size});
        layer.k_proj = weights.store->load_tensor(
            source,
            prefix + ".self_attn.k_proj.weight",
            storage_type,
            {config.num_key_value_heads * dim, config.hidden_size});
        layer.v_proj = weights.store->load_tensor(
            source,
            prefix + ".self_attn.v_proj.weight",
            storage_type,
            {config.num_key_value_heads * dim, config.hidden_size});
        layer.o_proj = weights.store->load_tensor(
            source,
            prefix + ".self_attn.o_proj.weight",
            storage_type,
            {config.hidden_size, config.num_attention_heads * dim});
        layer.gate_proj =
            weights.store->load_tensor(source, prefix + ".mlp.gate_proj.weight", storage_type, {config.intermediate_size, config.hidden_size});
        layer.up_proj =
            weights.store->load_tensor(source, prefix + ".mlp.up_proj.weight", storage_type, {config.intermediate_size, config.hidden_size});
        layer.down_proj =
            weights.store->load_tensor(source, prefix + ".mlp.down_proj.weight", storage_type, {config.hidden_size, config.intermediate_size});
        weights.layers.push_back(std::move(layer));
    }