diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..f6cb8ad --- /dev/null +++ b/.clang-format @@ -0,0 +1 @@ +BasedOnStyle: Google diff --git a/DEVELOPERS.md b/DEVELOPERS.md index bdc02c0..f670c49 100644 --- a/DEVELOPERS.md +++ b/DEVELOPERS.md @@ -71,19 +71,31 @@ The implementation code is roughly split into 4 layers, from high to low level: 4. Backend (`highway`) - Low-level hardware interface (SIMD in the case of highway) supporting the implementations in (3). +Besides these layers, supporting utilities are: + +- `compression/` - model compression operations. The 8-bit switched floating + point model conversion is here. +- `util/` - command line argument handling and any other utilities. + +## Style and Formatting + +A `.clang-format` configuration is provided with our defaults, please run source +files through `clang-format` (or a formatter that produces equivalent behavior) +before finalizing PR for submission. + ## Compile-Time Flags (Advanced) -There are several compile-time flags to be aware of (note these may or may not +There are several compile-time flags to be aware of (note these may or may not be exposed to the build system): -- `GEMMA_WEIGHT_T` : Sets the level of compression for weights (surfaced as - WEIGHT_TYPE in CMakeLists.txt). Currently this should be set to `SfpStream` - (default, if no flag is specified) for 8-bit SFP, or `hwy::bfloat16_t` to +- `GEMMA_WEIGHT_T` : Sets the level of compression for weights (surfaced as + WEIGHT_TYPE in CMakeLists.txt). Currently this should be set to `SfpStream` + (default, if no flag is specified) for 8-bit SFP, or `hwy::bfloat16_t` to enable for higher-fidelity (but slower) bfloat16 support. This is defined in `gemma.h`. - `GEMMA_MAX_SEQ_LEN` : Sets maximum sequence length to preallocate for the KV Cache. The default is 4096 tokens but can be overridden. This is not exposed - through `CMakeLists.txt` yet. + through `CMakeLists.txt` yet. In the medium term both of these will likely be deprecated in favor of handling options at runtime - allowing for multiple weight compression schemes in a single diff --git a/configs.h b/configs.h index 4be5f75..bf25596 100644 --- a/configs.h +++ b/configs.h @@ -21,7 +21,7 @@ // Allow changing pre-allocated kv cache size as a compiler flag #ifndef GEMMA_MAX_SEQLEN #define GEMMA_MAX_SEQLEN 4096 -#endif // !GEMMA_MAX_SEQLEN +#endif // !GEMMA_MAX_SEQLEN #include @@ -34,10 +34,10 @@ struct ConfigGemma7B { static constexpr int kVocabSize = 256128; static constexpr int kLayers = 28; static constexpr int kModelDim = 3072; - static constexpr int kFFHiddenDim = 16 * 3072 / 2; // = 24576 + static constexpr int kFFHiddenDim = 16 * 3072 / 2; // = 24576 static constexpr int kHeads = 16; - static constexpr int kKVHeads = 16; // standard MHA - static constexpr int kQKVDim = 256; // query size == key size == value size + static constexpr int kKVHeads = 16; // standard MHA + static constexpr int kQKVDim = 256; // query size == key size == value size static constexpr int kTopK = 1; }; @@ -46,13 +46,13 @@ struct ConfigGemma2B { static constexpr int kVocabSize = 256128; static constexpr int kLayers = 18; static constexpr int kModelDim = 2048; - static constexpr int kFFHiddenDim = 16 * 2048 / 2; // = 16384 + static constexpr int kFFHiddenDim = 16 * 2048 / 2; // = 16384 static constexpr int kHeads = 8; - static constexpr int kKVHeads = 8; // TODO(austinvhuang): add MQA support - static constexpr int kQKVDim = 256; // query size == key size == value size + static constexpr int kKVHeads = 8; // TODO(austinvhuang): add MQA support + static constexpr int kQKVDim = 256; // query size == key size == value size static constexpr int kTopK = 1; }; -} // namespace gcpp +} // namespace gcpp -#endif // THIRD_PARTY_GEMMA_CPP_CONFIGS_H_ +#endif // THIRD_PARTY_GEMMA_CPP_CONFIGS_H_ diff --git a/gemma.h b/gemma.h index 1e76a37..1ff98c1 100644 --- a/gemma.h +++ b/gemma.h @@ -26,15 +26,19 @@ // copybara:import_next_line:gemma_cpp #include "compression/compress.h" // SfpStream/NuqStream +// copybara:end // copybara:import_next_line:gemma_cpp -#include "configs.h" // kSeqLen +#include "configs.h" // kSeqLen +// copybara:end // copybara:import_next_line:gemma_cpp -#include "util/args.h" // ArgsBase +#include "util/args.h" // ArgsBase +// copybara:end #include "hwy/aligned_allocator.h" #include "hwy/base.h" // hwy::bfloat16_t #include "hwy/contrib/thread_pool/thread_pool.h" // copybara:import_next_line:sentencepiece #include "src/sentencepiece_processor.h" +// copybara:end namespace gcpp { @@ -42,7 +46,7 @@ namespace gcpp { // float, hwy::bfloat16_t, SfpStream, NuqStream #ifndef GEMMA_WEIGHT_T #define GEMMA_WEIGHT_T SfpStream -#endif // !GEMMA_WEIGHT_T +#endif // !GEMMA_WEIGHT_T using WeightT = GEMMA_WEIGHT_T; using EmbedderInputT = hwy::bfloat16_t; @@ -51,9 +55,9 @@ constexpr bool kSystemPrompt = false; struct KVCache { hwy::AlignedFreeUniquePtr - key_cache; // batch_size * kSeqLen * kLayers * kKVHeads * kQKVDim + key_cache; // batch_size * kSeqLen * kLayers * kKVHeads * kQKVDim hwy::AlignedFreeUniquePtr - value_cache; // batch_size * kSeqLen * kLayers * kKVHeads * kQKVDim + value_cache; // batch_size * kSeqLen * kLayers * kKVHeads * kQKVDim }; // Model variants: see configs.h for details. @@ -61,9 +65,9 @@ enum class Model { GEMMA_2B, GEMMA_7B }; enum class ModelTraining { GEMMA_IT, GEMMA_PT }; struct LoaderArgs : public ArgsBase { - LoaderArgs(int argc, char *argv[]) { InitAndParse(argc, argv); } + LoaderArgs(int argc, char* argv[]) { InitAndParse(argc, argv); } - static std::string ToLower(const std::string &text) { + static std::string ToLower(const std::string& text) { std::string result = text; std::transform(begin(result), end(result), begin(result), [](unsigned char c) { return std::tolower(c); }); @@ -89,7 +93,7 @@ struct LoaderArgs : public ArgsBase { } // Returns error string or nullptr if OK. - const char *Validate() const { + const char* Validate() const { const std::string model_type_lc = ToLower(model_type); if (model_type_lc != "2b-pt" && model_type_lc != "7b-pt" && model_type_lc != "2b-it" && model_type_lc != "7b-it") { @@ -111,11 +115,12 @@ struct LoaderArgs : public ArgsBase { } Path tokenizer; - Path model; // uncompressed weights OR - Path cache; // compressed weights + Path model; // uncompressed weights OR + Path cache; // compressed weights std::string model_type; - template void ForEach(const Visitor &visitor) { + template + void ForEach(const Visitor& visitor) { visitor(tokenizer, "tokenizer", Path(), "Path name of tokenizer model file. (required)"); visitor( @@ -138,10 +143,10 @@ struct LoaderArgs : public ArgsBase { struct GemmaInterface; struct Gemma { - Gemma(const LoaderArgs &args, hwy::ThreadPool &pool); - ~Gemma(); // must be defined after GemmaInterface's dtor is defined. + Gemma(const LoaderArgs& args, hwy::ThreadPool& pool); + ~Gemma(); // must be defined after GemmaInterface's dtor is defined. - const sentencepiece::SentencePieceProcessor &Tokenizer() const; + const sentencepiece::SentencePieceProcessor& Tokenizer() const; std::unique_ptr impl_; gcpp::ModelTraining model_training; @@ -153,7 +158,7 @@ using StreamFunc = std::function; using AcceptFunc = std::function; struct InferenceArgs : public ArgsBase { - InferenceArgs(int argc, char *argv[]) { InitAndParse(argc, argv); } + InferenceArgs(int argc, char* argv[]) { InitAndParse(argc, argv); } size_t max_tokens; size_t max_generated_tokens; @@ -163,7 +168,7 @@ struct InferenceArgs : public ArgsBase { bool multiturn; // Returns error string or nullptr if OK. - const char *Validate() const { + const char* Validate() const { if (max_tokens > gcpp::kSeqLen) { return "max_tokens is larger than the maximum sequence length (see " "configs.h)."; @@ -175,7 +180,8 @@ struct InferenceArgs : public ArgsBase { return nullptr; } - template void ForEach(const Visitor &visitor) { + template + void ForEach(const Visitor& visitor) { visitor(max_tokens, "max_tokens", size_t{3072}, "Maximum number of tokens in prompt + generation."); visitor(max_generated_tokens, "max_generated_tokens", size_t{2048}, @@ -186,19 +192,20 @@ struct InferenceArgs : public ArgsBase { "Make top-k sampling deterministic", 2); visitor(multiturn, "multiturn", false, "Multiturn mode (if 0, this clears the KV cache after every " - "interaction without quitting)\n Default = 0 (conversation resets every turn)"); + "interaction without quitting)\n Default = 0 (conversation " + "resets every turn)"); } }; -void GenerateGemma(Gemma &gemma, const InferenceArgs &args, - const std::vector &prompt, size_t start_pos, - hwy::ThreadPool &pool, hwy::ThreadPool &inner_pool, - const StreamFunc &stream_token, - const AcceptFunc &accept_token, std::mt19937 &g, +void GenerateGemma(Gemma& gemma, const InferenceArgs& args, + const std::vector& prompt, size_t start_pos, + hwy::ThreadPool& pool, hwy::ThreadPool& inner_pool, + const StreamFunc& stream_token, + const AcceptFunc& accept_token, std::mt19937& g, int verbosity); constexpr int EOS_ID = 1; -} // namespace gcpp +} // namespace gcpp -#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_H_ +#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_H_