Repository: google/gemma.cpp
File: gemma/gemma.cc (C++)

Code before:
            (gen_end - gen_start);
      }
      break;
    }
  }
}

template <class TConfig>
void ComputeCrossEntropy(const ByteStorageT& weights_u8,
                         ByteStorageT& decode_u8,
                         const GemmaTokenizer& tokenizer, size_t max_tokens,
                         const std::vector<int>& prompt, KVCache& kv_cache,
                         hwy::ThreadPool& pool, int verbosity,
                         float& cross_entropy) {
  const WeightsT<TConfig>& weights = GetWeights<TConfig>(weights_u8);
  auto& activations = GetActivations<TConfig, 1>(decode_u8);

  static constexpr size_t kModelDim = TConfig::kModelDim;
  static constexpr size_t kVocabSize = TConfig::kVocabSize;
  std::vector<float> logits(kVocabSize);
  Softmax(activations.logits.data(), kVocabSize);
  float total_entropy = 0.0f;
  for (size_t pos = 0; pos < max_tokens && pos < prompt.size(); ++pos) {
    if (verbosity >= 4) {
      LogTopK(tokenizer, logits.data(), activations.logits.data(), kVocabSize,
              10);
    }
    const int token = prompt[pos];
    const float prob = activations.logits[token];
    if (verbosity >= 3) {
      printf("pos %4zu token %6d = %-12s  %.10e  %14.10f bits\n", pos, token,
             TokenString(tokenizer, token).c_str(), prob,
             -std::log(prob) / std::log(2.0));
    }
    total_entropy -= std::max(std::log(prob), -64.0f);
    if (verbosity >= 2 && pos % 100 == 99) {
      printf("Processed %zu tokens, cross-entropy per token: %f\n", pos + 1,
             total_entropy / std::log(2.0) / (pos + 1));
    }
    Transformer(token, pos, weights, activations, kv_cache, pool,
                /*layers_output=*/nullptr);
    MatVec<kVocabSize, kModelDim>(
        weights.embedder_input_embedding, 0, activations.x.data(),
        activations.even_odd.data(), activations.logits.data(), pool);
    LogitsSoftCap(30.0f, activations.logits.data(), kVocabSize);
    memcpy(logits.data(), activations.logits.data(),
           kVocabSize * sizeof(logits[0]));
    Softmax(activations.logits.data(), kVocabSize);
  }
  cross_entropy = total_entropy / std::log(2.0);
}

Review comment (refactor):
hm, I'm not sure we actually want to encourage f32 weights. For training, wouldn't it make sense to use bf16 weights? Those are considered compressed, though we'd have to build with -DGEMMA_WEIGHT_T=hwy::bfloat16_t. That should be faster, and let us only have a single function here. And maybe we could even get rid of kWeightsAreCompressed?

Note that 'compressed' can also mean f32. It would be nice to get rid of the duplicated non-compressed code now that we have the separate compress_weights binary.

Code after:
    if (token == runtime_config.eos_id) {
      if (runtime_config.verbosity >= 2) {
        const double gen_end = hwy::platform::Now();
        timing_info.gen_tok_sec =
            static_cast<double>(pos_offset - pos_gen_start) /
            (gen_end - gen_start);
      }
      break;
    }
  }
}

template <class TConfig>
void ComputeCrossEntropy(const ByteStorageT& weights_u8,
                         ByteStorageT& decode_u8,
                         const GemmaTokenizer& tokenizer, size_t max_tokens,
                         const std::vector<int>& prompt, KVCache& kv_cache,
                         hwy::ThreadPool& pool, int verbosity,
                         float& cross_entropy) {
  const WeightsT<TConfig>& weights = GetWeights<TConfig>(weights_u8);
  auto& activations = GetActivations<TConfig, 1>(decode_u8);

  static constexpr size_t kModelDim = TConfig::kModelDim;
  static constexpr size_t kVocabSize = TConfig::kVocabSize;
  std::vector<float> logits(kVocabSize);
  Softmax(activations.logits.data(), kVocabSize);
  float total_entropy = 0.0f;
  for (size_t pos = 0; pos < max_tokens && pos < prompt.size(); ++pos) {
    if (verbosity >= 4) {
      LogTopK(tokenizer, logits.data(), activations.logits.data(), kVocabSize,
              10);
    }
    const int token = prompt[pos];
    const float prob = activations.logits[token];
    if (verbosity >= 3) {
      printf("pos %4zu token %6d = %-12s  %.10e  %14.10f bits\n", pos, token,
             TokenString(tokenizer, token).c_str(), prob,
             -std::log(prob) / std::log(2.0));
    }
    total_entropy -= std::max(std::log(prob), -64.0f);
    if (verbosity >= 2 && pos % 100 == 99) {
      printf("Processed %zu tokens, cross-entropy per token: %f\n", pos + 1,
             total_entropy / std::log(2.0) / (pos + 1));
    }
    Transformer(token, pos, weights, activations, kv_cache, pool,
                /*layers_output=*/nullptr);
    MatVec<kVocabSize, kModelDim>(
        weights.embedder_input_embedding, 0, activations.x.data(),
        activations.even_odd.data(), activations.logits.data(), pool);
    LogitsSoftCap(30.0f, activations.logits.data(), kVocabSize);
    memcpy(logits.data(), activations.logits.data(),