diff --git a/dist/jni/cp_inference.cpp b/dist/jni/cp_inference.cpp index 2cd435d..371d778 100644 --- a/dist/jni/cp_inference.cpp +++ b/dist/jni/cp_inference.cpp @@ -31,16 +31,33 @@ static std::vector read_floats(const char* path, size_t expect) { fclose(f); return v; } +// Garde-t-on le tenseur en F16 ou on convertit en F32 ? +// Rule : les poids matmul (attn_q/k/v/output, ffn_gate/up/down) restent F16 -> moitié BW +// dans les gros mul_mat (~85% du coût CP). Les norms restent en F32 car ggml_mul exige +// même type des deux opérandes (xn_f32 * w_norm_f32). Économie pratique : embeds ~1024 +// éléments * 5 layers * 6 norms = 30 KB de norms, négligeable. Les matmul valent +// ~5 * (3*1024*1024 attn + 3*1024*MLP ffn) qui dominent. +static bool cp_keep_f16(const std::string& name) { + // Tout ce qui n'est pas "_norm.weight" peut rester F16. Sécurité : on n'autorise + // que les noms connus du CP qwen3 (5L). Les autres tombent en F32 par défaut. + if (name.find("_norm.weight") != std::string::npos) return false; + return true; +} + bool cp_load(CPState& s, const char* gguf_path, const char* heads_path, const char* embs_path) { ggml_context* meta = nullptr; gguf_init_params p; p.no_alloc = true; p.ctx = &meta; gguf_context* g = gguf_init_from_file(gguf_path, p); if (!g) { fprintf(stderr, "cp_load: gguf open fail %s\n", gguf_path); return false; } int64_t n = gguf_get_n_tensors(g); + + // Budget : on alloue à la taille native (F16 reste F16 si autorisé, sinon F32). Borne + // sup = somme des max(nb_f32, nb_native) pour être tranquille. size_t bytes = 0; for (int64_t i = 0; i < n; i++) { ggml_tensor* mt = ggml_get_tensor(meta, gguf_get_tensor_name(g, i)); - bytes += ggml_nelements(mt) * sizeof(float); + size_t nb_f32 = (size_t)ggml_nelements(mt) * sizeof(float); + bytes += nb_f32; // sup } ggml_init_params ip = { bytes + (size_t)n * ggml_tensor_overhead() + (1u << 20), nullptr, false }; s.weights_ctx = ggml_init(ip); @@ -48,19 +65,36 @@ bool cp_load(CPState& s, const char* gguf_path, const char* heads_path, const ch FILE* f = fopen(gguf_path, "rb"); const size_t off = gguf_get_data_offset(g); std::vector tmp; + size_t kept_f16 = 0, conv_f32 = 0; for (int64_t i = 0; i < n; i++) { const char* name = gguf_get_tensor_name(g, i); ggml_tensor* mt = ggml_get_tensor(meta, name); - ggml_tensor* t32 = ggml_new_tensor(s.weights_ctx, GGML_TYPE_F32, ggml_n_dims(mt), mt->ne); - size_t nb = ggml_nbytes(mt); + size_t nb_src = ggml_nbytes(mt); int64_t ne = ggml_nelements(mt); - tmp.resize(nb); + tmp.resize(nb_src); fseek(f, off + gguf_get_tensor_offset(g, i), SEEK_SET); - if (fread(tmp.data(), 1, nb, f) != nb) { fprintf(stderr, "cp_load read fail %s\n", name); fclose(f); return false; } - if (mt->type == GGML_TYPE_F32) memcpy(t32->data, tmp.data(), nb); - else if (mt->type == GGML_TYPE_F16) ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), (float*)t32->data, ne); - else { fprintf(stderr, "cp_load unsupported type %s for %s\n", ggml_type_name(mt->type), name); fclose(f); return false; } - s.tensors[name] = t32; + if (fread(tmp.data(), 1, nb_src, f) != nb_src) { fprintf(stderr, "cp_load read fail %s\n", name); fclose(f); return false; } + + ggml_type dst_type = GGML_TYPE_F32; + if (mt->type == GGML_TYPE_F16 && cp_keep_f16(name)) dst_type = GGML_TYPE_F16; + else if (mt->type == GGML_TYPE_F32) dst_type = GGML_TYPE_F32; + else dst_type = GGML_TYPE_F32; // F16 -> F32 pour les norms + + ggml_tensor* t = ggml_new_tensor(s.weights_ctx, dst_type, ggml_n_dims(mt), mt->ne); + if (dst_type == mt->type) { + memcpy(t->data, tmp.data(), nb_src); + if (dst_type == GGML_TYPE_F16) kept_f16++; + } else if (mt->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F32) { + ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), (float*)t->data, ne); + conv_f32++; + } else if (mt->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { + memcpy(t->data, tmp.data(), nb_src); + } else { + fprintf(stderr, "cp_load unsupported cast %s -> %s for %s\n", + ggml_type_name(mt->type), ggml_type_name(dst_type), name); + fclose(f); return false; + } + s.tensors[name] = t; } fclose(f); gguf_free(g); @@ -68,8 +102,9 @@ bool cp_load(CPState& s, const char* gguf_path, const char* heads_path, const ch const size_t TAB = (size_t)N_CB * N_VOCAB * N_EMBD; s.heads = read_floats(heads_path, TAB); if (s.heads.empty()) return false; s.codec_embs = read_floats(embs_path, TAB); if (s.codec_embs.empty()) return false; - fprintf(stderr, "cp_load: %lld tensors + heads(%.0f MB) + codec_embs(%.0f MB) OK\n", - (long long)n, TAB * 4 / 1048576.0, TAB * 4 / 1048576.0); + fprintf(stderr, "cp_load: %lld tensors (f16 kept=%zu, f32=%zu) + heads(%.0f MB) + codec_embs(%.0f MB) OK\n", + (long long)n, kept_f16, conv_f32 + (size_t)(n - (int64_t)(kept_f16 + conv_f32)), + TAB * 4 / 1048576.0, TAB * 4 / 1048576.0); return true; }