| 115 | }; |
| 116 | |
| 117 | using TorchCudaSamplingPolicy = engine::sampling::TorchCudaSamplingPolicy; |
| 118 | |
| 119 | std::shared_ptr<const Vevo2ARWeights> load_ar_weights( |
| 120 | const Vevo2Assets & assets, |
| 121 | ggml_backend_t backend, |
| 122 | core::BackendType backend_type, |
| 123 | size_t weight_context_bytes, |
| 124 | assets::TensorStorageType storage_type, |
| 125 | const assets::TensorSource & source) { |
| 126 | const auto & config = assets.config.ar; |
| 127 | const int64_t dim = ar_head_dim(config); |
| 128 | auto weights = std::make_shared<Vevo2ARWeights>(); |
| 129 | weights->store = std::make_shared<core::BackendWeightStore>( |
| 130 | backend, |
| 131 | backend_type, |
| 132 | "vevo2.ar.weights", |
| 133 | weight_context_bytes); |
| 134 | weights->token_embedding = weights->store->load_tensor( |
| 135 | source, |
| 136 | "model.embed_tokens.weight", |
| 137 | storage_type, |
| 138 | {config.vocab_size, config.hidden_size}); |
| 139 | weights->layers.reserve(static_cast<size_t>(config.num_hidden_layers)); |
| 140 | for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { |
| 141 | const std::string prefix = "model.layers." + std::to_string(layer); |
| 142 | Vevo2ARWeights::Layer layer_weights; |
| 143 | layer_weights.input_norm = modules::binding::norm_weight_from_source(*weights->store, source, prefix + ".input_layernorm", config.hidden_size); |
| 144 | layer_weights.q_proj = modules::binding::linear_from_source( |
| 145 | *weights->store, |
| 146 | source, |
| 147 | prefix + ".self_attn.q_proj", |
| 148 | storage_type, |
| 149 | config.num_attention_heads * dim, |
| 150 | config.hidden_size, |
| 151 | true); |
| 152 | layer_weights.k_proj = modules::binding::linear_from_source( |
| 153 | *weights->store, |
| 154 | source, |
| 155 | prefix + ".self_attn.k_proj", |
| 156 | storage_type, |
| 157 | config.num_key_value_heads * dim, |
| 158 | config.hidden_size, |
| 159 | true); |
| 160 | layer_weights.v_proj = modules::binding::linear_from_source( |
| 161 | *weights->store, |
| 162 | source, |
| 163 | prefix + ".self_attn.v_proj", |
| 164 | storage_type, |
| 165 | config.num_key_value_heads * dim, |
| 166 | config.hidden_size, |
| 167 | true); |
| 168 | layer_weights.o_proj = modules::binding::linear_from_source( |
| 169 | *weights->store, |
| 170 | source, |
| 171 | prefix + ".self_attn.o_proj", |
| 172 | storage_type, |
| 173 | config.hidden_size, |
| 174 | config.num_attention_heads * dim, |
no test coverage detected