happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
gemma.cpp/gemma/weights.h

566 lines
22 KiB
C++
Raw Blame History

// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_H_
#define THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_H_
#include <stddef.h>
#include <array>
#include <complex>
#include <cstdio>
#include <string>
#include <unordered_set>
#include <vector>
#include "compression/compress.h"
#include "compression/shared.h"
#include "gemma/common.h"
#include "gemma/configs.h"
#include "util/allocator.h"
#include "hwy/aligned_allocator.h"
#include "hwy/base.h"
#include "hwy/contrib/thread_pool/thread_pool.h"
namespace gcpp {
// Different tensors need to appear in a ForEachTensor, according to what is
// happening.
enum class ForEachType {
// Under normal circumstances, when not initializing or loading, we can
// include all tensors and ignore the null ones.
kIgnoreNulls,
// If there is a table of contents, we can include all tensors.
kLoadWithToc,
// There is no table of contents, so we have to be careful to only include
// tensors that are actually present.
kLoadNoToc,
// We need to initialize all tensors needed when there is no table of
// contents. This differs from kLoadNoToc in that we need to include any
// tensor that is allocated but not loaded directly from file.
kInitNoToc,
};
template <class TConfig>
struct CompressedLayer {
// Large data is constructed separately.
CompressedLayer()
: attn_vec_einsum_w("att_ein", kModelDim, kHeads * kQKVDim),
qkv_einsum_w("qkv_ein", (kHeads + 2 * kKVHeads) * kQKVDim, kModelDim),
qkv_einsum_w1("qkv1_w", kHeads * kQKVDim, kModelDim),
qkv_einsum_w2("qkv2_w", 2 * kKVHeads * kQKVDim, kModelDim),
attention_output_biases("attn_ob", 1, kAOBiasDim),
griffin({.linear_x_w = {"gr_lin_x_w", kGriffinDim, kGriffinDim},
.linear_x_biases = {"gr_lin_x_b", 1, kGriffinDim},
.linear_y_w = {"gr_lin_y_w", kGriffinDim, kGriffinDim},
.linear_y_biases = {"gr_lin_y_b", 1, kGriffinDim},
.linear_out_w = {"gr_lin_out_w", kGriffinDim, kGriffinDim},
.linear_out_biases = {"gr_lin_out_b", 1, kGriffinDim},
.conv_w = {"gr_conv_w", kConv1dWidth, kGriffinDim},
.conv_biases = {"gr_conv_b", 1, kGriffinDim},
.gate_w = {"gr_gate_w", 2 * kGriffinDim, kGriffinDim / kHeads},
.gate_biases = {"gr_gate_b", 1, kGriffinDim * 2},
.a = {"gr_a", 1, kGriffinDim}}),
// MultiHeadDotProductAttention.
vit({.attn_out_w = {"attn_out_w", kHeads * kQKVDim, kModelDim},
.attn_out_b = {"attn_out_b", 1, kModelDim},
.qkv_einsum_w = {"qkv_ein_w", (kHeads + 2 * kKVHeads) * kQKVDim,
kModelDim},
.qkv_einsum_b = {"qkv_ein_b", (kHeads + 2 * kKVHeads), kQKVDim},
.linear_0_w = {"linear_0_w", kModelDim, kFFHiddenDim},
.linear_0_b = {"linear_0_b", 1, kFFHiddenDim},
.linear_1_w = {"linear_1_w", kFFHiddenDim, kModelDim},
.linear_1_b = {"linear_1_b", 1, kModelDim},
.layer_norm_0_bias = {"ln_0_bias", 1, kModelDim},
.layer_norm_0_scale = {"ln_0_scale", 1, kModelDim},
.layer_norm_1_bias = {"ln_1_bias", 1, kModelDim},
.layer_norm_1_scale = {"ln_1_scale", 1, kModelDim}}),
gating_einsum_w("gating_ein", 2 * kFFHiddenDim, kModelDim),
gating_einsum_w1("gating1_w", kFFHiddenDim, kModelDim),
gating_einsum_w2("gating2_w", kFFHiddenDim, kModelDim),
linear_w("linear_w", kModelDim, kFFHiddenDim),
pre_attention_norm_scale("pre_att_ns", 1, kModelDim),
pre_ffw_norm_scale("pre_ff_ns", 1, kModelDim),
post_attention_norm_scale(
"post_att_ns", 1, kPostNorm == PostNormType::Scale ? kModelDim : 0),
post_ffw_norm_scale("post_ff_ns", 1,
kPostNorm == PostNormType::Scale ? kModelDim : 0),
ffw_gating_biases("ffw_gat_b", 1, kFFBiases ? 2 * kFFHiddenDim : 0),
ffw_output_biases("ffw_out_b", 1, kFFBiases ? kModelDim : 0),
att_weights("att_w", kModelDim, kHeads * kQKVDim)
{}
~CompressedLayer() = default;
using Weight = typename TConfig::Weight;
// If weights are f32, also f32; otherwise at least bf16. Useful for ops that
// do not yet support smaller compressed types, or require at least bf16. When
// weights are f32, we also want such tensors to be f32.
// If weights are complex, this is also complex.
using WeightF32OrBF16 =
hwy::If<hwy::IsSame<Weight, std::complex<double>>(), std::complex<double>,
hwy::If<hwy::IsSame<Weight, double>(), double,
hwy::If<IsF32<Weight>(), float, BF16>>>;
static constexpr size_t kHeads = TConfig::kHeads;
static constexpr size_t kKVHeads = TConfig::kKVHeads;
static constexpr size_t kModelDim = TConfig::kModelDim;
static constexpr size_t kQKVDim = TConfig::kQKVDim;
static constexpr size_t kFFHiddenDim = TConfig::kFFHiddenDim;
static constexpr size_t kAttVecEinsumWSize = kHeads * kQKVDim * kModelDim;
static constexpr size_t kQKVEinsumWSize =
(kHeads + 2 * kKVHeads) * kQKVDim * kModelDim;
static constexpr size_t kQKVEinsumBSize = (kHeads + 2 * kKVHeads) * kQKVDim;
// 2x for (gelu gating vector, gated vector)
static constexpr size_t kGatingEinsumWSize = 2 * kFFHiddenDim * kModelDim;
static constexpr size_t kConv1dWidth = TConfig::kConv1dWidth;
static constexpr bool kFFBiases = TConfig::kFFBiases;
static constexpr PostNormType kPostNorm = TConfig::kPostNorm;
static constexpr size_t kAOBiasDim =
TConfig::kSoftmaxAttnOutputBiases ? kModelDim : 0;
static constexpr size_t kGriffinDim =
TConfig::kGriffinLayers > 0 ? kModelDim : 0;
template <class T>
using ArrayT = MatPtrT<T>;
ArrayT<Weight> attn_vec_einsum_w;
// qkv_einsum_w holds 2 different matrices, which may be separated out.
// On loading, which is used depends on what is in the file.
// At inference, the one with a non-null ptr is used.
ArrayT<Weight> qkv_einsum_w;
ArrayT<Weight> qkv_einsum_w1;
ArrayT<Weight> qkv_einsum_w2;
ArrayT<float> attention_output_biases;
struct {
ArrayT<Weight> linear_x_w;
ArrayT<float> linear_x_biases;
ArrayT<Weight> linear_y_w;
ArrayT<float> linear_y_biases;
ArrayT<Weight> linear_out_w;
ArrayT<float> linear_out_biases;
ArrayT<float> conv_w;
ArrayT<float> conv_biases;
ArrayT<Weight> gate_w;
ArrayT<float> gate_biases;
ArrayT<float> a;
} griffin;
struct {
// MultiHeadDotProductAttention.
ArrayT<WeightF32OrBF16> attn_out_w;
ArrayT<float> attn_out_b;
ArrayT<WeightF32OrBF16> qkv_einsum_w;
ArrayT<float> qkv_einsum_b;
// MlpBlock.
ArrayT<WeightF32OrBF16> linear_0_w;
ArrayT<float> linear_0_b;
ArrayT<WeightF32OrBF16> linear_1_w;
ArrayT<float> linear_1_b;
// LayerNorm.
ArrayT<WeightF32OrBF16> layer_norm_0_bias;
ArrayT<WeightF32OrBF16> layer_norm_0_scale;
ArrayT<WeightF32OrBF16> layer_norm_1_bias;
ArrayT<WeightF32OrBF16> layer_norm_1_scale;
} vit;
// gating_einsum_w holds 2 different matrices, which may be separated out.
// On loading, which is used depends on what is in the file.
// At inference, the one with a non-null ptr is used.
ArrayT<Weight> gating_einsum_w;
ArrayT<Weight> gating_einsum_w1;
ArrayT<Weight> gating_einsum_w2;
ArrayT<Weight> linear_w;
// We don't yet have an RMSNorm that accepts all Weight.
ArrayT<WeightF32OrBF16> pre_attention_norm_scale;
ArrayT<WeightF32OrBF16> pre_ffw_norm_scale;
ArrayT<WeightF32OrBF16> post_attention_norm_scale;
ArrayT<WeightF32OrBF16> post_ffw_norm_scale;
ArrayT<float> ffw_gating_biases;
ArrayT<float> ffw_output_biases;
// Reshaped attention; not loaded from disk via ForEachTensor.
ArrayT<Weight> att_weights;
// Initializes att_weights from attn_vec_einsum_w, hence this must be called
// after loading weights via ForEachTensor.
// TODO: update compression/convert_weights to bake this in.
void Reshape(MatStorage& storage) {
if (attn_vec_einsum_w.data() == nullptr) return;
constexpr size_t kModelDim = TConfig::kModelDim;
constexpr size_t kHeads = TConfig::kHeads;
constexpr size_t kQKVDim = TConfig::kQKVDim;
// Would have to implement a CompressTraits::Copy for NUQ.
static_assert(!hwy::IsSame<Weight, NuqStream>());
// Reshape [kHeads, kModelDim, kQKVDim] to [kModelDim, kHeads * kQKVDim].
storage.Allocate();
att_weights.SetPtr(storage);
for (size_t m = 0; m < kModelDim; ++m) {
Weight* HWY_RESTRICT out_row = att_weights.data() + m * kHeads * kQKVDim;
for (size_t h = 0; h < kHeads; ++h) {
hwy::CopyBytes(
attn_vec_einsum_w.data() + h * kModelDim * kQKVDim + m * kQKVDim,
out_row + h * kQKVDim, kQKVDim * sizeof(Weight));
}
}
att_weights.set_scale(attn_vec_einsum_w.scale());
}
// Used by ForEachTensor for per-layer tensors.
#define GEMMA_CALL_FUNC(member) \
{ \
for (int i = 0; i < ptrs.size(); ++i) { \
tensors[i] = &ptrs[i]->member; \
} \
if (tensors[0]->Ptr() != nullptr || fet != ForEachType::kIgnoreNulls) { \
func(ptrs[0]->member.CacheName(layer_idx, sep, sep_index).c_str(), \
hwy::Span<MatPtr*>(tensors, ptrs.size())); \
} \
}
template <class Func>
static void ForEachTensor(const std::vector<CompressedLayer<TConfig>*>& ptrs,
int layer_idx, ForEachType fet, Func func,
char sep = ' ', int sep_index = -1) {
MatPtr* tensors[ptrs.size()];
auto type = TConfig::kLayerConfig[layer_idx];
if (type == LayerAttentionType::kVit) {
// MHA.
GEMMA_CALL_FUNC(vit.attn_out_w);
GEMMA_CALL_FUNC(vit.attn_out_b);
GEMMA_CALL_FUNC(vit.qkv_einsum_w);
GEMMA_CALL_FUNC(vit.qkv_einsum_b);
// MlpBlock.
GEMMA_CALL_FUNC(vit.linear_0_w);
GEMMA_CALL_FUNC(vit.linear_0_b);
GEMMA_CALL_FUNC(vit.linear_1_w);
GEMMA_CALL_FUNC(vit.linear_1_b);
// LayerNorm.
GEMMA_CALL_FUNC(vit.layer_norm_0_bias);
GEMMA_CALL_FUNC(vit.layer_norm_0_scale);
GEMMA_CALL_FUNC(vit.layer_norm_1_bias);
GEMMA_CALL_FUNC(vit.layer_norm_1_scale);
return;
}
if (type == LayerAttentionType::kGemma) {
if (fet != ForEachType::kLoadNoToc) {
GEMMA_CALL_FUNC(att_weights);
}
if (fet == ForEachType::kInitNoToc || fet == ForEachType::kLoadNoToc ||
fet == ForEachType::kIgnoreNulls) {
GEMMA_CALL_FUNC(attn_vec_einsum_w);
}
GEMMA_CALL_FUNC(qkv_einsum_w);
if (fet == ForEachType::kIgnoreNulls ||
fet == ForEachType::kLoadWithToc) {
// The unwanted ones will be null or not in the toc.
GEMMA_CALL_FUNC(qkv_einsum_w1);
GEMMA_CALL_FUNC(qkv_einsum_w2);
}
} else {
GEMMA_CALL_FUNC(griffin.linear_x_w);
GEMMA_CALL_FUNC(griffin.linear_x_biases);
GEMMA_CALL_FUNC(griffin.linear_y_w);
GEMMA_CALL_FUNC(griffin.linear_y_biases);
GEMMA_CALL_FUNC(griffin.linear_out_w);
GEMMA_CALL_FUNC(griffin.linear_out_biases);
GEMMA_CALL_FUNC(griffin.conv_w);
GEMMA_CALL_FUNC(griffin.conv_biases);
GEMMA_CALL_FUNC(griffin.gate_w);
GEMMA_CALL_FUNC(griffin.gate_biases);
GEMMA_CALL_FUNC(griffin.a);
}
GEMMA_CALL_FUNC(gating_einsum_w);
if (fet == ForEachType::kIgnoreNulls || fet == ForEachType::kLoadWithToc) {
// The unwanted ones will be null or not in the toc.
GEMMA_CALL_FUNC(gating_einsum_w1);
GEMMA_CALL_FUNC(gating_einsum_w2);
}
GEMMA_CALL_FUNC(linear_w);
GEMMA_CALL_FUNC(pre_attention_norm_scale);
GEMMA_CALL_FUNC(pre_ffw_norm_scale);
if (TConfig::kPostNorm == PostNormType::Scale) {
GEMMA_CALL_FUNC(post_attention_norm_scale);
GEMMA_CALL_FUNC(post_ffw_norm_scale);
}
if (TConfig::kFFBiases) {
GEMMA_CALL_FUNC(ffw_gating_biases);
GEMMA_CALL_FUNC(ffw_output_biases);
}
if (TConfig::kSoftmaxAttnOutputBiases &&
type == LayerAttentionType::kGemma) {
GEMMA_CALL_FUNC(attention_output_biases);
}
}
// Sets all the tensors in the layer to zero. Memory must have been allocated.
void ZeroInit(int layer_idx) {
ForEachTensor({this}, layer_idx, ForEachType::kIgnoreNulls,
[](const char*, hwy::Span<MatPtr*> tensors) {
tensors[0]->ZeroInit();
});
}
// Allocates memory for all the tensors in the layer.
// Note that this is slow and only used for a stand-alone layer.
void Allocate() {
layer_storage.clear();
ForEachTensor({this}, /*layer_idx=*/0, ForEachType::kInitNoToc,
[this](const char* name, hwy::Span<MatPtr*> tensors) {
this->layer_storage.emplace_back(*tensors[0]);
layer_storage.back().Allocate();
tensors[0]->SetPtr(layer_storage.back());
});
}
// Storage for all the matrices and vectors. Only used for a stand-alone
// layer. For a model, the CompressedWeights::model_storage is used instead.
std::vector<MatStorage> layer_storage;
};
template <class TConfig>
struct CompressedWeights {
explicit CompressedWeights(hwy::ThreadPool& pool)
: embedder_input_embedding("c_embedding", TConfig::kVocabSize,
TConfig::kModelDim),
final_norm_scale("c_final_norm", 1, TConfig::kModelDim),
vit_encoder_norm_bias("enc_norm_bias", 1,
TConfig::VitConfig::kModelDim),
vit_encoder_norm_scale("enc_norm_scale", 1,
TConfig::VitConfig::kModelDim),
vit_img_embedding_bias("img_emb_bias", 1,
TConfig::VitConfig::kModelDim),
vit_img_embedding_kernel("img_emb_kernel", 14 * 14 * 3,
TConfig::VitConfig::kModelDim),
vit_img_pos_embedding("img_pos_emb", 256,
TConfig::VitConfig::kModelDim),
vit_img_head_bias("img_head_bias", 1, TConfig::kModelDim),
vit_img_head_kernel("img_head_kernel", TConfig::VitConfig::kModelDim,
TConfig::kModelDim),
scale_names({"att_ein", "qkv_ein", "gr_lin_x_w", "gr_lin_y_w",
"gr_lin_out_w", "gr_gate_w", "gating_ein", "linear_w"}) {}
~CompressedWeights() = default;
using Weight = typename TConfig::Weight;
using WeightF32OrBF16 = typename CompressedLayer<TConfig>::WeightF32OrBF16;
using WeightF32OrInputT = hwy::If<hwy::IsSame<WeightF32OrBF16, BF16>(),
EmbedderInputT, WeightF32OrBF16>;
MatPtrT<WeightF32OrInputT> embedder_input_embedding;
MatPtrT<WeightF32OrBF16> final_norm_scale;
// Vit parts.
MatPtrT<WeightF32OrBF16> vit_encoder_norm_bias;
MatPtrT<WeightF32OrBF16> vit_encoder_norm_scale;
MatPtrT<float> vit_img_embedding_bias;
MatPtrT<WeightF32OrBF16> vit_img_embedding_kernel;
MatPtrT<float> vit_img_pos_embedding;
// The head maps from VitConfig::kModelDim (Vit final layer) to
// kModelDim (LLM input).
MatPtrT<float> vit_img_head_bias;
MatPtrT<WeightF32OrBF16> vit_img_head_kernel;
// Storage for all the matrices and vectors.
std::vector<MatStorage> model_storage;
std::unordered_set<std::string> scale_names;
CompressedLayer<TConfig> c_layers[TConfig::kLayers];
CompressedLayer<typename TConfig::VitConfig>
vit_layers[TConfig::VitConfig::kLayers];
// Called by weights.cc after ForEachTensor.
void Reshape(hwy::ThreadPool& pool) {
size_t storage_index = model_storage.size();
for (size_t layer = 0; layer < TConfig::kLayers; ++layer) {
model_storage.emplace_back(GetLayer(layer)->att_weights);
}
pool.Run(0, TConfig::kLayers,
[this, storage_index](uint64_t layer, size_t /*thread*/) {
GetLayer(layer)->Reshape(model_storage[storage_index + layer]);
});
}
void ZeroInit() {
embedder_input_embedding.ZeroInit();
final_norm_scale.ZeroInit();
for (int i = 0; i < TConfig::kLayers; ++i) {
c_layers[i].ZeroInit(i);
}
}
const CompressedLayer<TConfig>* GetLayer(size_t layer) const {
return &c_layers[layer];
}
CompressedLayer<TConfig>* GetLayer(size_t layer) { return &c_layers[layer]; }
const CompressedLayer<typename TConfig::VitConfig>* GetVitLayer(
size_t layer) const {
return &vit_layers[layer];
}
CompressedLayer<typename TConfig::VitConfig>* GetVitLayer(size_t layer) {
return &vit_layers[layer];
}
// Copies the data from other to *this.
void CopyFrom(const CompressedWeights<TConfig>& other) {
ForEachTensor({this, const_cast<CompressedWeights<TConfig>*>(&other)},
ForEachType::kIgnoreNulls,
[](const char*, hwy::Span<MatPtr*> tensors) {
hwy::CopyBytes(tensors[1]->Ptr(), tensors[0]->Ptr(),
tensors[1]->SizeBytes());
});
}
// If scales is empty, computes and returns the scale factors for the tensors,
// otherwise applies the scale factors to the tensors.
void GetOrApplyScales(std::vector<float>& scales) {
int scale_pos = 0;
ForEachTensor(
{this}, ForEachType::kIgnoreNulls,
[&scales, &scale_pos, this](const char*, hwy::Span<MatPtr*> tensors) {
if (this->scale_names.count(tensors[0]->Name())) {
if (scale_pos < scales.size()) {
tensors[0]->set_scale(scales[scale_pos]);
} else {
float scale = ScaleWeights(tensors[0]->data<float>(),
tensors[0]->NumElements());
scales.push_back(scale);
}
++scale_pos;
}
});
HWY_ASSERT(scale_pos == TConfig::kNumTensorScales);
}
template <class Func>
static void ForEachTensor(
const std::vector<CompressedWeights<TConfig>*>& ptrs, ForEachType fet,
Func func) {
std::vector<CompressedLayer<TConfig>*> layers(ptrs.size());
std::vector<CompressedLayer<typename TConfig::VitConfig>*> vit_layers(
ptrs.size());
MatPtr* tensors[ptrs.size()];
// Variables used by GEMMA_CALL_FUNC.
int layer_idx = -1;
char sep = ' ';
int sep_index = -1;
GEMMA_CALL_FUNC(embedder_input_embedding);
GEMMA_CALL_FUNC(final_norm_scale);
if constexpr (TConfig::VitConfig::kLayers > 0) {
// Vit parts.
GEMMA_CALL_FUNC(vit_encoder_norm_bias);
GEMMA_CALL_FUNC(vit_encoder_norm_scale);
GEMMA_CALL_FUNC(vit_img_embedding_bias);
GEMMA_CALL_FUNC(vit_img_embedding_kernel);
GEMMA_CALL_FUNC(vit_img_pos_embedding);
GEMMA_CALL_FUNC(vit_img_head_bias);
GEMMA_CALL_FUNC(vit_img_head_kernel);
}
for (int layer_idx = 0; layer_idx < TConfig::kLayers; ++layer_idx) {
for (int i = 0; i < ptrs.size(); ++i) {
layers[i] = ptrs[i]->GetLayer(layer_idx);
}
CompressedLayer<TConfig>::ForEachTensor(layers, layer_idx, fet, func);
}
// Vit layers. Not supported for compress_weights.
if constexpr (TConfig::VitConfig::kLayers > 0) {
for (int layer_idx = 0; layer_idx < TConfig::VitConfig::kLayers;
++layer_idx) {
auto type = TConfig::VitConfig::kLayerConfig[layer_idx];
HWY_ASSERT(type == LayerAttentionType::kVit);
for (int i = 0; i < ptrs.size(); ++i) {
vit_layers[i] = ptrs[i]->GetVitLayer(layer_idx);
}
CompressedLayer<typename TConfig::VitConfig>::ForEachTensor(
vit_layers, layer_idx, fet, func);
}
}
}
};
#undef GEMMA_CALL_FUNC
// Pair of configs for the compressed and uncompressed weights.
template <class CConfig, class UCConfig>
struct ConfigPair {
using uc = UCConfig;
using c = CConfig;
};
// ----------------------------------------------------------------------------
// Interface
template <typename TConfig>
struct AllocateCompressedWeights {
ByteStorageT operator()(hwy::ThreadPool& pool) const {
using TWeights = CompressedWeights<TConfig>;
ByteStorageT weights_u8 = AllocateSizeof<TWeights>();
TWeights* weights = reinterpret_cast<TWeights*>(weights_u8.get());
new (weights) TWeights(pool);
std::vector<MatPtr*> model_toc;
auto& model_storage = weights->model_storage;
TWeights::ForEachTensor(
{weights}, ForEachType::kInitNoToc,
[&model_toc, &model_storage](const char*, hwy::Span<MatPtr*> tensors) {
model_toc.push_back(tensors[0]);
model_storage.emplace_back(*tensors[0]);
});
// Allocate in parallel using the pool.
pool.Run(0, model_storage.size(),
[&model_toc, &model_storage](uint64_t task, size_t /*thread*/) {
model_storage[task].Allocate();
model_toc[task]->SetPtr(model_storage[task]);
});
return weights_u8;
}
};
template <typename TConfig>
struct ZeroInitCompressedWeights {
void operator()(ByteStorageT& weights_u8, hwy::ThreadPool& pool) const {
CompressedWeights<TConfig>& weights =
*reinterpret_cast<CompressedWeights<TConfig>*>(weights_u8.get());
weights.ZeroInit();
}
};
template <typename TConfig>
struct ReshapeCompressedWeights {
void operator()(ByteStorageT& weights_u8, hwy::ThreadPool& pool) const {
CompressedWeights<TConfig>& weights =
*reinterpret_cast<CompressedWeights<TConfig>*>(weights_u8.get());
weights.Reshape(pool);
}
};
// TODO: also add RandInitCompressedWeights
ByteStorageT LoadCompressedWeights(const Path& weights, Model model_type,
Type weight_type, hwy::ThreadPool& pool);
void LogWeightStats(Model model, Type weight_type, const ByteStorageT& weights);
} // namespace gcpp
#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_WEIGHTS_H_
Reference in New Issue View Git Blame Copy Permalink