diff --git a/tests/test-perplexity-plumbing.cpp b/tests/test-perplexity-plumbing.cpp index 536e540ead46..3815e27d7879 100644 --- a/tests/test-perplexity-plumbing.cpp +++ b/tests/test-perplexity-plumbing.cpp @@ -45,12 +45,18 @@ int main(int argc, char ** argv) { const std::string root = argv[1]; const std::string perplexity = read_file(root + "/tools/perplexity/perplexity.cpp"); + const std::string ppl = slice_between(perplexity, + "static results_perplexity perplexity(llama_context * ctx, const common_params & params, const int32_t n_ctx)", + "static bool decode_helper"); const std::string kl = slice_between(perplexity, "static bool kl_divergence(llama_context * ctx, const common_params & params)", "if (kld.count < 100) return true;"); ok &= expect(perplexity.find("static int ppl_max_logits_rows(int n_vocab, const common_params & params)") != std::string::npos, "perplexity must cap full-vocab logits rows to avoid multi-GiB output buffers"); + ok &= expect(ppl.find("const int max_logits_rows = params.logits_file.empty()") != std::string::npos && + ppl.find(": std::max(1, std::min(n_ctx, params.n_batch))") != std::string::npos, + "logits baseline generation must bypass the ordinary perplexity row cap"); ok &= expect(perplexity.find("logits_stream.write(\"_logits_\", 8)") != std::string::npos, "perplexity must write upstream-compatible v1 logits baselines"); ok &= expect(perplexity.find("kld_logits::") == std::string::npos, @@ -61,14 +67,14 @@ int main(int argc, char ** argv) { "KLD must retain the upstream v1 probability cutoff"); ok &= expect(perplexity.find("if (!kl_divergence(ctx, params))") != std::string::npos, "perplexity KL failures must propagate to a nonzero process exit"); - ok &= expect(kl.find("const int max_logits_rows = ppl_max_logits_rows(n_vocab, params)") != std::string::npos, - "KL divergence must use the bounded logits-row cap"); - ok &= expect(kl.find("const int n_batch = std::max(1, std::min(n_ctx_i, std::min(params.n_batch, max_logits_rows)))") != std::string::npos, - "KL divergence batch size must be bounded by max_logits_rows"); + ok &= expect(kl.find("ppl_max_logits_rows") == std::string::npos, + "KL divergence must not silently cap the requested batch by logits memory"); + ok &= expect(kl.find("const int n_batch = std::max(1, std::min(n_ctx_i, params.n_batch))") != std::string::npos, + "KL divergence batch size must honor the requested batch size"); ok &= expect(kl.find("llama_batch_init(n_batch, 0, 1)") != std::string::npos, "KL divergence batch allocation must match bounded n_batch"); - ok &= expect(kl.find("std::vector log_probs_uint16(size_t(max_logits_rows) * nv)") != std::string::npos, - "KL divergence base-logit buffer must be bounded by max_logits_rows"); + ok &= expect(kl.find("std::vector log_probs_uint16(size_t(n_batch) * nv)") != std::string::npos, + "KL divergence base-logit buffer must match the decode batch size"); ok &= expect(kl.find("const int logits_first = std::max(first, pos_start)") != std::string::npos && kl.find("const int logits_end = std::min(n_ctx_i - 1, pos_start + batch_size)") != std::string::npos, "KL divergence must process only the logits rows produced by the current decode slice"); diff --git a/tools/perplexity/README.md b/tools/perplexity/README.md index f82d34c8a251..9230cc1d586f 100644 --- a/tools/perplexity/README.md +++ b/tools/perplexity/README.md @@ -22,6 +22,17 @@ and finally the `--kl-divergence` argument to indicate that the program should c This is a measure of how similar the FP16 and the quantized logit distributions are with a value of 0 indicating that the distribution are the same. The uncertainty on the mean KL divergence is calculated by assuming the KL divergence per token follows a Gaussian distribution. +KL-divergence baselines and candidates are processed in blocks controlled by +`--batch-size`; `--ubatch-size` controls how those blocks are split internally. +The F32 logits output requires `vocabulary size * batch size * 4` bytes of host +memory. KL-divergence evaluation also keeps a compressed baseline block that +requires roughly half as much memory. Use matching `--batch-size` and +`--ubatch-size` settings when generating a baseline and evaluating a candidate. + +For example, a model with 248320 vocabulary entries and a batch size of 2048 +uses 1940 MiB for the F32 logits block and about 970 MiB for the compressed +baseline block. + In addition to the KL divergence the following statistics are calculated with `--kl-divergence`: * Ratio of mean FP16 PPL and quantized PPL. Uncertainty is estimated on logits, then propagated. The logarithm of this metric is also calculated and printed, it is 0 if the logit distributions are the same. diff --git a/tools/perplexity/perplexity.cpp b/tools/perplexity/perplexity.cpp index e7ddac951d1e..d4e3270a283f 100644 --- a/tools/perplexity/perplexity.cpp +++ b/tools/perplexity/perplexity.cpp @@ -518,7 +518,9 @@ static results_perplexity perplexity(llama_context * ctx, const common_params & const int n_chunk_max = tokens.size() / n_ctx; const int n_vocab = llama_vocab_n_tokens(vocab); - const int max_logits_rows = ppl_max_logits_rows(n_vocab, params); + const int max_logits_rows = params.logits_file.empty() + ? ppl_max_logits_rows(n_vocab, params) + : std::max(1, std::min(n_ctx, params.n_batch)); const int n_seq_ctx = std::max(1, params.n_ctx / n_ctx); const int n_seq = params.logits_file.empty() ? std::min(n_seq_ctx, max_logits_rows) : 1; const int n_batch = std::max(1, std::min(n_ctx, std::min( @@ -1796,8 +1798,7 @@ static bool kl_divergence(llama_context * ctx, const common_params & params) { } const int n_ctx_i = static_cast(n_ctx); - const int max_logits_rows = ppl_max_logits_rows(n_vocab, params); - const int n_batch = std::max(1, std::min(n_ctx_i, std::min(params.n_batch, max_logits_rows))); + const int n_batch = std::max(1, std::min(n_ctx_i, params.n_batch)); const int num_batches = (n_ctx_i + n_batch - 1) / n_batch; const int n_seq = 1; const int nv = 2*((n_vocab + 1)/2) + 4; @@ -1806,7 +1807,7 @@ static bool kl_divergence(llama_context * ctx, const common_params & params) { llama_batch batch = llama_batch_init(n_batch, 0, 1); - std::vector log_probs_uint16(size_t(max_logits_rows) * nv); + std::vector log_probs_uint16(size_t(n_batch) * nv); std::vector kld_values(size_t(n_ctx - 1 - n_ctx/2)*n_chunk); std::vector p_diff_values(size_t(n_ctx - 1 - n_ctx/2)*n_chunk);