diff --git a/csrc/config/compressed_tensors_config.cpp b/csrc/config/compressed_tensors_config.cpp new file mode 100644 index 000000000..d612e1951 --- /dev/null +++ b/csrc/config/compressed_tensors_config.cpp @@ -0,0 +1,345 @@ +#include "compressed_tensors_config.hpp" +#include "module_target_matcher.hpp" + +#include +#include +#include +#include +#include + +namespace infinilm::config { +namespace { + +using json = nlohmann::json; + +[[noreturn]] void config_error(const std::string &path, const std::string &message) { + throw std::invalid_argument(path + ": " + message); +} + +std::string lowercase(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + return value; +} + +std::string uppercase(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), + [](unsigned char c) { return static_cast(std::toupper(c)); }); + return value; +} + +std::string require_string(const json &value, const std::string &path) { + if (!value.is_string()) { + config_error(path, "expected a string"); + } + return value.get(); +} + +int require_int(const json &value, const std::string &path) { + if (!value.is_number_integer()) { + config_error(path, "expected an integer"); + } + const auto parsed = value.get(); + if (parsed < std::numeric_limits::min() + || parsed > std::numeric_limits::max()) { + config_error(path, "integer is outside the supported range"); + } + return static_cast(parsed); +} + +std::vector parse_string_list(const json &value, const std::string &path) { + if (!value.is_array()) { + config_error(path, "expected an array of strings"); + } + + std::vector result; + result.reserve(value.size()); + for (size_t i = 0; i < value.size(); ++i) { + result.push_back(require_string( + value.at(i), path + "[" + std::to_string(i) + "]")); + } + return result; +} + +QuantizationValueType parse_value_type(const json &value, const std::string &path) { + const auto type = lowercase(require_string(value, path)); + if (type == "int") { + return QuantizationValueType::INT; + } + if (type == "float") { + return QuantizationValueType::FLOAT; + } + config_error(path, "expected \"int\" or \"float\""); +} + +QuantizationStrategy parse_strategy(const json &value, const std::string &path) { + const auto strategy = lowercase(require_string(value, path)); + if (strategy == "tensor") { + return QuantizationStrategy::TENSOR; + } + if (strategy == "channel") { + return QuantizationStrategy::CHANNEL; + } + if (strategy == "group") { + return QuantizationStrategy::GROUP; + } + if (strategy == "block") { + return QuantizationStrategy::BLOCK; + } + if (strategy == "token") { + return QuantizationStrategy::TOKEN; + } + if (strategy == "tensor_group") { + return QuantizationStrategy::TENSOR_GROUP; + } + if (strategy == "attn_head") { + return QuantizationStrategy::ATTN_HEAD; + } + config_error(path, "unknown quantization strategy \"" + strategy + "\""); +} + +QuantizationDynamicMode parse_dynamic(const json &value, const std::string &path) { + if (value.is_boolean()) { + return value.get() + ? QuantizationDynamicMode::DYNAMIC + : QuantizationDynamicMode::STATIC; + } + if (value.is_string() && lowercase(value.get()) == "local") { + return QuantizationDynamicMode::LOCAL; + } + config_error(path, "expected false, true, or \"local\""); +} + +QuantizationArgs parse_quantization_args(const json &value, const std::string &path) { + if (!value.is_object()) { + config_error(path, "expected an object"); + } + + QuantizationArgs args; + if (value.contains("num_bits") && !value.at("num_bits").is_null()) { + args.num_bits = require_int(value.at("num_bits"), path + ".num_bits"); + if (args.num_bits <= 0) { + config_error(path + ".num_bits", "must be positive"); + } + } + if (value.contains("type") && !value.at("type").is_null()) { + args.type = parse_value_type(value.at("type"), path + ".type"); + } + if (value.contains("symmetric") && !value.at("symmetric").is_null()) { + if (!value.at("symmetric").is_boolean()) { + config_error(path + ".symmetric", "expected a boolean"); + } + args.symmetric = value.at("symmetric").get(); + } + if (value.contains("strategy") && !value.at("strategy").is_null()) { + args.strategy = parse_strategy(value.at("strategy"), path + ".strategy"); + } + if (value.contains("dynamic") && !value.at("dynamic").is_null()) { + args.dynamic = parse_dynamic(value.at("dynamic"), path + ".dynamic"); + } + + // Fields outside InfiniLM's `W8A8` execution scope remain available through + // `CompressedTensorsConfig::raw_config`. The scheme resolver will reject an + // unsupported combination instead of silently selecting a kernel. + return args; +} + +QuantizationGroup make_w8a8_preset( + const std::string &name, + const json &targets, + const std::string &path) { + QuantizationGroup group; + group.name = name; + group.targets = parse_string_list(targets, path); + + QuantizationArgs weights; + weights.type = QuantizationValueType::INT; + weights.strategy = QuantizationStrategy::CHANNEL; + group.weights = weights; + + QuantizationArgs inputs; + inputs.type = QuantizationValueType::INT; + inputs.strategy = QuantizationStrategy::TOKEN; + inputs.dynamic = QuantizationDynamicMode::DYNAMIC; + group.input_activations = inputs; + return group; +} + +QuantizationGroup parse_group( + const std::string &name, + const json &value, + const std::string &path) { + if (value.is_array()) { + const auto preset = uppercase(name); + if (preset == "W8A8" || preset == "INT8") { + return make_w8a8_preset(name, value, path); + } + if (preset == "UNQUANTIZED") { + QuantizationGroup group; + group.name = name; + group.targets = parse_string_list(value, path); + return group; + } + config_error( + path, + "unsupported preset \"" + name + "\"; use an explicit group object"); + } + if (!value.is_object()) { + config_error(path, "expected a group object or preset target list"); + } + if (!value.contains("targets")) { + config_error(path + ".targets", "missing required field"); + } + + QuantizationGroup group; + group.name = name; + group.targets = parse_string_list(value.at("targets"), path + ".targets"); + + if (value.contains("weights") && !value.at("weights").is_null()) { + group.weights = parse_quantization_args(value.at("weights"), path + ".weights"); + } + if (value.contains("input_activations") + && !value.at("input_activations").is_null()) { + group.input_activations = parse_quantization_args( + value.at("input_activations"), path + ".input_activations"); + } + if (value.contains("output_activations") + && !value.at("output_activations").is_null()) { + group.output_activations = parse_quantization_args( + value.at("output_activations"), path + ".output_activations"); + } + if (value.contains("format") && !value.at("format").is_null()) { + group.format = require_string(value.at("format"), path + ".format"); + } + return group; +} + +QuantizationStatus parse_status(const json &value, const std::string &path) { + const auto status = lowercase(require_string(value, path)); + if (status == "initialized") { + return QuantizationStatus::INITIALIZED; + } + if (status == "calibration") { + return QuantizationStatus::CALIBRATION; + } + if (status == "frozen") { + return QuantizationStatus::FROZEN; + } + if (status == "compressed") { + return QuantizationStatus::COMPRESSED; + } + if (status == "decompressed") { + return QuantizationStatus::DECOMPRESSED; + } + config_error(path, "unknown quantization status \"" + status + "\""); +} + +} // namespace + +CompressedTensorsConfig CompressedTensorsConfig::from_json(const json &config) { + constexpr std::string_view root = "quantization_config"; + if (!config.is_object()) { + config_error(std::string(root), "expected an object"); + } + + CompressedTensorsConfig parsed; + parsed.raw_config = config; + + if (config.contains("quant_method") && !config.at("quant_method").is_null()) { + parsed.quant_method = require_string( + config.at("quant_method"), std::string(root) + ".quant_method"); + } + if (parsed.quant_method != "compressed-tensors") { + config_error( + std::string(root) + ".quant_method", + "expected \"compressed-tensors\""); + } + if (config.contains("format") && !config.at("format").is_null()) { + parsed.format = require_string( + config.at("format"), std::string(root) + ".format"); + } + if (config.contains("quantization_status") + && !config.at("quantization_status").is_null()) { + parsed.quantization_status = parse_status( + config.at("quantization_status"), + std::string(root) + ".quantization_status"); + } + if (config.contains("ignore") && !config.at("ignore").is_null()) { + parsed.ignore = parse_string_list( + config.at("ignore"), std::string(root) + ".ignore"); + } + if (config.contains("kv_cache_scheme") + && !config.at("kv_cache_scheme").is_null()) { + parsed.kv_cache_scheme = parse_quantization_args( + config.at("kv_cache_scheme"), + std::string(root) + ".kv_cache_scheme"); + } + if (config.contains("global_compression_ratio") + && !config.at("global_compression_ratio").is_null()) { + const auto &ratio = config.at("global_compression_ratio"); + if (!ratio.is_number()) { + config_error( + std::string(root) + ".global_compression_ratio", + "expected a number"); + } + parsed.global_compression_ratio = ratio.get(); + } + + if (!config.contains("config_groups")) { + config_error(std::string(root) + ".config_groups", "missing required field"); + } + const auto &groups = config.at("config_groups"); + if (!groups.is_object()) { + config_error(std::string(root) + ".config_groups", "expected an object"); + } + parsed.config_groups.reserve(groups.size()); + for (const auto &[name, group] : groups.items()) { + parsed.config_groups.push_back(parse_group( + name, + group, + std::string(root) + ".config_groups." + name)); + } + return parsed; +} + +bool CompressedTensorsConfig::is_ignored( + std::string_view module_name, + std::string_view module_type) const { + return ModuleTargetMatcher::match_any(ignore, module_name, module_type) + != ModuleTargetMatcher::MatchKind::NONE; +} + +const QuantizationGroup *CompressedTensorsConfig::resolve_group( + std::string_view module_name, + std::string_view module_type) const { + if (is_ignored(module_name, module_type)) { + return nullptr; + } + + const QuantizationGroup *resolved_group = nullptr; + const QuantizationGroup *conflicting_group = nullptr; + auto best_match = ModuleTargetMatcher::MatchKind::NONE; + for (const auto &group : config_groups) { + const auto current_match = ModuleTargetMatcher::match_any( + group.targets, module_name, module_type); + if (current_match > best_match) { + resolved_group = &group; + conflicting_group = nullptr; + best_match = current_match; + } else if (current_match != ModuleTargetMatcher::MatchKind::NONE + && current_match == best_match) { + conflicting_group = &group; + } + } + + if (conflicting_group != nullptr) { + throw std::invalid_argument( + "ambiguous `compressed-tensors` groups for module `" + + std::string(module_name) + "`: `" + resolved_group->name + + "` and `" + conflicting_group->name + + "` match with equal specificity"); + } + return resolved_group; +} + +} // namespace infinilm::config diff --git a/csrc/config/compressed_tensors_config.hpp b/csrc/config/compressed_tensors_config.hpp new file mode 100644 index 000000000..45359b243 --- /dev/null +++ b/csrc/config/compressed_tensors_config.hpp @@ -0,0 +1,88 @@ +#pragma once + +#include "nlohmann/json.hpp" + +#include +#include +#include +#include + +namespace infinilm::config { + +// Values mirror the public `compressed-tensors` configuration vocabulary. They +// describe checkpoint metadata, not a concrete InfiniLM kernel implementation. +enum class QuantizationValueType { + INT, + FLOAT, +}; + +enum class QuantizationStrategy { + TENSOR, + CHANNEL, + GROUP, + BLOCK, + TOKEN, + TENSOR_GROUP, + ATTN_HEAD, +}; + +// The `compressed-tensors` format accepts `false` (static), `true` (fully +// dynamic), and `"local"` (only local quantization parameters are dynamic). +enum class QuantizationDynamicMode { + STATIC, + DYNAMIC, + LOCAL, +}; + +enum class QuantizationStatus { + INITIALIZED, + CALIBRATION, + FROZEN, + COMPRESSED, + DECOMPRESSED, +}; + +struct QuantizationArgs { + int num_bits = 8; + QuantizationValueType type = QuantizationValueType::INT; + bool symmetric = true; + QuantizationStrategy strategy = QuantizationStrategy::TENSOR; + QuantizationDynamicMode dynamic = QuantizationDynamicMode::STATIC; +}; + +struct QuantizationGroup { + std::string name; + std::vector targets; + std::optional weights; + std::optional input_activations; + std::optional output_activations; + std::optional format; +}; + +struct CompressedTensorsConfig { + static CompressedTensorsConfig from_json(const nlohmann::json &config); + + bool is_ignored( + std::string_view module_name, + std::string_view module_type) const; + + // Returns `nullptr` when the module is ignored or no group matches. Throws + // `std::invalid_argument` if distinct groups match with equal specificity. + const QuantizationGroup *resolve_group( + std::string_view module_name, + std::string_view module_type) const; + + std::string quant_method = "compressed-tensors"; + std::string format = "fakequant"; + QuantizationStatus quantization_status = QuantizationStatus::INITIALIZED; + std::vector config_groups; + std::vector ignore; + std::optional kv_cache_scheme; + std::optional global_compression_ratio; + + // Keep the source object for diagnostics and forward-compatible inspection + // of metadata that does not affect the currently supported inference path. + nlohmann::json raw_config; +}; + +} // namespace infinilm::config diff --git a/csrc/config/model_config.cpp b/csrc/config/model_config.cpp index f0095558e..f2e1766b5 100644 --- a/csrc/config/model_config.cpp +++ b/csrc/config/model_config.cpp @@ -1,8 +1,27 @@ #include "model_config.hpp" +namespace { +nlohmann::json extract_quantization_config(const nlohmann::json &config) { + auto quantization_config = config.find("quantization_config"); + if (quantization_config != config.end()) { + return *quantization_config; + } + + auto text_config = config.find("text_config"); + if (text_config != config.end() && text_config->is_object()) { + quantization_config = text_config->find("quantization_config"); + if (quantization_config != text_config->end()) { + return *quantization_config; + } + } + + return nullptr; +} +} // namespace + namespace infinilm::config { ModelConfig::ModelConfig(const nlohmann::json &json) : config_json(json) { - this->quant_config = QuantConfig(config_json["quantization_config"]); + this->quant_config = QuantConfig(extract_quantization_config(config_json)); }; ModelConfig::ModelConfig(const std::string &path) { @@ -13,7 +32,7 @@ ModelConfig::ModelConfig(const std::string &path) { } else { throw std::runtime_error("Could not open config file: " + path); } - this->quant_config = QuantConfig(config_json["quantization_config"]); + this->quant_config = QuantConfig(extract_quantization_config(config_json)); } infinilm::quantization::QuantScheme diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index dc0e89287..f751e9890 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -88,6 +88,12 @@ class ModelConfig { return quant_config.get_quantization_method(); } + std::shared_ptr get_quantization_method( + std::string_view module_name, + std::string_view module_type) const { + return quant_config.get_quantization_method(module_name, module_type); + } + infinicore::DataType get_dtype() const; infinilm::quantization::QuantScheme get_quant_scheme() const; diff --git a/csrc/config/module_target_matcher.cpp b/csrc/config/module_target_matcher.cpp new file mode 100644 index 000000000..7babc3995 --- /dev/null +++ b/csrc/config/module_target_matcher.cpp @@ -0,0 +1,70 @@ +#include "module_target_matcher.hpp" + +#include +#include + +namespace infinilm::config { +namespace { + +constexpr std::string_view REGEX_PREFIX = "re:"; + +bool is_regex_target(std::string_view target) { + return target.size() >= REGEX_PREFIX.size() + && target.substr(0, REGEX_PREFIX.size()) == REGEX_PREFIX; +} + +} // namespace + +ModuleTargetMatcher::MatchKind ModuleTargetMatcher::match( + std::string_view target, + std::string_view module_name, + std::string_view module_type) { + if (target == module_name) { + return MatchKind::EXACT_NAME; + } + + if (is_regex_target(target)) { + const std::string pattern(target.substr(REGEX_PREFIX.size())); + try { + const std::regex expression( + pattern, + std::regex_constants::ECMAScript | std::regex_constants::optimize); + if (std::regex_search( + module_name.begin(), + module_name.end(), + expression, + std::regex_constants::match_continuous)) { + return MatchKind::REGEX; + } + } catch (const std::regex_error &error) { + throw std::invalid_argument( + "invalid module target regex `" + std::string(target) + + "`: " + error.what()); + } + return MatchKind::NONE; + } + + if (target == module_type) { + return MatchKind::MODULE_TYPE; + } + return MatchKind::NONE; +} + +ModuleTargetMatcher::MatchKind ModuleTargetMatcher::match_any( + const std::vector &targets, + std::string_view module_name, + std::string_view module_type) { + MatchKind best_match = MatchKind::NONE; + for (const auto &target : targets) { + const auto current_match = match(target, module_name, module_type); + if (current_match > best_match) { + best_match = current_match; + if (best_match == MatchKind::EXACT_NAME) { + break; + } + } + } + return best_match; +} + +} // namespace infinilm::config diff --git a/csrc/config/module_target_matcher.hpp b/csrc/config/module_target_matcher.hpp new file mode 100644 index 000000000..790a347f1 --- /dev/null +++ b/csrc/config/module_target_matcher.hpp @@ -0,0 +1,30 @@ +#pragma once + +#include +#include +#include + +namespace infinilm::config { + +class ModuleTargetMatcher { +public: + // Match kinds are ordered from least to most specific. + enum class MatchKind { + NONE, + MODULE_TYPE, + REGEX, + EXACT_NAME, + }; + + static MatchKind match( + std::string_view target, + std::string_view module_name, + std::string_view module_type); + + static MatchKind match_any( + const std::vector &targets, + std::string_view module_name, + std::string_view module_type); +}; + +} // namespace infinilm::config diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index e58966d89..d11e5ac8d 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -1,8 +1,45 @@ #include "quant_config.hpp" +#include + namespace infinilm::config { +namespace { + +bool is_unquantized(const QuantizationGroup &group) { + return !group.weights.has_value() + && !group.input_activations.has_value() + && !group.output_activations.has_value(); +} + +bool is_supported_w8a8(const QuantizationGroup &group) { + if (!group.weights.has_value() + || !group.input_activations.has_value() + || group.output_activations.has_value()) { + return false; + } + + const auto &weights = *group.weights; + const auto &inputs = *group.input_activations; + return weights.num_bits == 8 + && weights.type == QuantizationValueType::INT + && weights.symmetric + && weights.strategy == QuantizationStrategy::CHANNEL + && weights.dynamic == QuantizationDynamicMode::STATIC + && inputs.num_bits == 8 + && inputs.type == QuantizationValueType::INT + && inputs.symmetric + && inputs.strategy == QuantizationStrategy::TOKEN + && inputs.dynamic == QuantizationDynamicMode::DYNAMIC; +} + +} // namespace + QuantConfig::QuantConfig(const nlohmann::json &json) : quantization_config(json) { - this->quantization_method = get_quantization_method(); + if (!quantization_config.is_null() + && quantization_config.value("quant_method", "") == "compressed-tensors") { + compressed_tensors_config_ = CompressedTensorsConfig::from_json(quantization_config); + } + quantization_method = get_quantization_method(); } std::shared_ptr @@ -22,11 +59,42 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "quark") { return std::make_shared(quantization_config); - } else { - return std::make_shared(quantization_config); } // Add other schemes as needed return std::make_shared(quantization_config); // Default case if no matching scheme } + +std::shared_ptr +QuantConfig::get_quantization_method( + std::string_view module_name, + std::string_view module_type) const { + if (!compressed_tensors_config_.has_value()) { + return get_quantization_method(); + } + + const auto &config = *compressed_tensors_config_; + const auto *group = config.resolve_group(module_name, module_type); + if (group == nullptr || is_unquantized(*group)) { + return std::make_shared( + quantization_config); + } + + const auto &format = group->format.value_or(config.format); + if (format != "int-quantized") { + throw std::invalid_argument( + "unsupported `compressed-tensors` format `" + format + + "` for module `" + std::string(module_name) + "`"); + } + if (!is_supported_w8a8(*group)) { + throw std::invalid_argument( + "unsupported `compressed-tensors` scheme in group `" + group->name + + "` for module `" + std::string(module_name) + + "`; expected static symmetric INT8 channel weights and dynamic " + "symmetric INT8 per-token input activations"); + } + + return std::make_shared( + quantization_config); +} } // namespace infinilm::config diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index fb0b8abf3..737ea7c15 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -1,9 +1,11 @@ #pragma once -#include "../utils.hpp" #include "../layers/quantization/quantization.hpp" +#include "../utils.hpp" +#include "compressed_tensors_config.hpp" #include "nlohmann/json.hpp" #include #include +#include namespace infinilm::config { @@ -15,6 +17,12 @@ class QuantConfig { QuantConfig(const nlohmann::json &json); std::shared_ptr get_quantization_method() const; + std::shared_ptr + get_quantization_method(std::string_view module_name, std::string_view module_type) const; + + const std::optional &get_compressed_tensors_config() const { + return compressed_tensors_config_; + } infinilm::quantization::QuantScheme get_quant_scheme() const { if (quantization_method != nullptr) { @@ -58,6 +66,7 @@ class QuantConfig { private: nlohmann::json quantization_config; std::shared_ptr quantization_method; + std::optional compressed_tensors_config_; infinilm::quantization::KVQuantAlgo kv_quant_scheme = infinilm::quantization::KVQuantAlgo::NONE; std::optional kv_cache_dtype_ = std::nullopt; diff --git a/csrc/layers/attention/attention.cpp b/csrc/layers/attention/attention.cpp index 16506ef02..5c101baff 100644 --- a/csrc/layers/attention/attention.cpp +++ b/csrc/layers/attention/attention.cpp @@ -2,6 +2,9 @@ #include "../../utils.hpp" #include "../rotary_embedding/rotary_embedding.hpp" +#include +#include + namespace infinilm::layers::attention { Attention::Attention(std::shared_ptr model_config, @@ -26,14 +29,29 @@ Attention::Attention(std::shared_ptr model_config num_attention_heads_ = total_num_heads / tp_size; num_key_value_heads_ = total_num_kv_heads < tp_size ? 1 : total_num_kv_heads / tp_size; - auto quantization_method = model_config->get_quantization_method(); + const std::string module_prefix = "model.layers." + std::to_string(layer_idx) + ".self_attn."; + auto q_quantization = model_config->get_quantization_method( + module_prefix + "q_proj", "Linear"); + auto k_quantization = model_config->get_quantization_method( + module_prefix + "k_proj", "Linear"); + auto v_quantization = model_config->get_quantization_method( + module_prefix + "v_proj", "Linear"); + if (q_quantization->get_quant_scheme() != k_quantization->get_quant_scheme() + || q_quantization->get_quant_scheme() + != v_quantization->get_quant_scheme()) { + throw std::invalid_argument( + "fused QKV projections require the same quantization scheme in `" + + module_prefix + "{q_proj,k_proj,v_proj}`"); + } + auto o_quantization = model_config->get_quantization_method( + module_prefix + "o_proj", "Linear"); auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; qkv_proj_ = std::make_shared( hidden_size_, head_dim_, total_num_heads, total_num_kv_heads, "q_proj", "k_proj", "v_proj", register_fn, - quantization_method, use_bias, dtype, device, rank_info); + q_quantization, use_bias, dtype, device, rank_info); o_proj_ = this->register_module( - "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method, + "o_proj", total_num_heads * head_dim_, hidden_size_, o_quantization, use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm); rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); diff --git a/csrc/layers/mlp/mlp.cpp b/csrc/layers/mlp/mlp.cpp index f7604c505..4aa446e33 100644 --- a/csrc/layers/mlp/mlp.cpp +++ b/csrc/layers/mlp/mlp.cpp @@ -2,6 +2,9 @@ #include "../../global_state/global_state.hpp" #include "infinicore/ops.hpp" +#include +#include + namespace infinilm::layers::mlp { MLP::MLP(std::shared_ptr model_config, @@ -26,6 +29,42 @@ MLP::MLP(std::shared_ptr model_config, use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm); } +MLP::MLP(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device) { + + const auto &dtype{model_config->get_dtype()}; + hidden_size_ = model_config->get("hidden_size"); + intermediate_size_ = model_config->get("intermediate_size"); + use_bias_ = model_config->get_or("mlp_bias", false); + + const engine::distributed::RankInfo &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); + int tp_rank = rank_info.tp_rank; + int tp_size = rank_info.tp_size; + + const std::string module_prefix = "model.layers." + std::to_string(layer_idx) + ".mlp."; + auto gate_quantization = model_config->get_quantization_method( + module_prefix + "gate_proj", "Linear"); + auto up_quantization = model_config->get_quantization_method( + module_prefix + "up_proj", "Linear"); + if (gate_quantization->get_quant_scheme() + != up_quantization->get_quant_scheme()) { + throw std::invalid_argument( + "fused gate/up projections require the same quantization scheme in `" + + module_prefix + "{gate_proj,up_proj}`"); + } + auto down_quantization = model_config->get_quantization_method( + module_prefix + "down_proj", "Linear"); + + auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; + gate_up_proj_ = std::make_shared( + hidden_size_, intermediate_size_, "gate_proj", "up_proj", register_fn, + gate_quantization, use_bias_, dtype, device, rank_info); + down_proj_ = this->register_module( + "down_proj", intermediate_size_, hidden_size_, down_quantization, + use_bias_, dtype, device, tp_rank, tp_size, rank_info.comm); +} + infinicore::Tensor MLP::forward(const infinicore::Tensor &hidden_states) const { // 1. Project to gate and up auto hidden_states_mutable = hidden_states; diff --git a/csrc/layers/mlp/mlp.hpp b/csrc/layers/mlp/mlp.hpp index abd81bd88..227480add 100644 --- a/csrc/layers/mlp/mlp.hpp +++ b/csrc/layers/mlp/mlp.hpp @@ -28,6 +28,10 @@ class MLP : public infinicore::nn::Module { MLP(std::shared_ptr model_config, const infinicore::Device &device); + MLP(std::shared_ptr model_config, + size_t layer_idx, + const infinicore::Device &device); + /** * @brief Forward pass: compute MLP output * diff --git a/csrc/layers/quantization/compressed_tensors.cpp b/csrc/layers/quantization/compressed_tensors.cpp index ff5617a1e..5a456aa61 100644 --- a/csrc/layers/quantization/compressed_tensors.cpp +++ b/csrc/layers/quantization/compressed_tensors.cpp @@ -18,12 +18,12 @@ std::vector CompressedTensors::get_param_layout( descs.push_back({"weight", {out_features, in_features}, infinicore::DataType::I8, split_dim, tp_rank, tp_size}); // weight_scale is per-output-channel [out_features, 1]; always split on // dim0 (output dimension) for ColumnParallel, and don't split for RowParallel. - int scale_split_dim = (split_dim == 0) ? 0 : -1; - int scale_tp_size = (split_dim == 0) ? tp_size : 1; - int scale_tp_rank = (split_dim == 0) ? tp_rank : 0; - descs.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, scale_split_dim, scale_tp_rank, scale_tp_size}); + int output_split_dim = (split_dim == 0) ? 0 : -1; + int output_tp_size = (split_dim == 0) ? tp_size : 1; + int output_tp_rank = (split_dim == 0) ? tp_rank : 0; + descs.push_back({"weight_scale", {out_features, 1}, infinicore::DataType::F32, output_split_dim, output_tp_rank, output_tp_size}); if (bias) { - descs.push_back({"bias", {out_features}, dtype, -1, 0, 1}); + descs.push_back({"bias", {out_features}, dtype, output_split_dim, output_tp_rank, output_tp_size}); } return descs; } diff --git a/csrc/models/qwen3/qwen3_attention.cpp b/csrc/models/qwen3/qwen3_attention.cpp index 7d9beb043..12f82023f 100644 --- a/csrc/models/qwen3/qwen3_attention.cpp +++ b/csrc/models/qwen3/qwen3_attention.cpp @@ -3,6 +3,9 @@ #include "../../layers/attention/attention.hpp" #include "../../utils.hpp" +#include +#include + namespace infinilm::models::qwen3 { Qwen3Attention::Qwen3Attention(std::shared_ptr model_config, @@ -28,14 +31,29 @@ Qwen3Attention::Qwen3Attention(std::shared_ptr mo ? 1 : total_num_kv_heads / tp_size; - auto quantization_method = model_config->get_quantization_method(); + const std::string module_prefix = "model.layers." + std::to_string(layer_idx) + ".self_attn."; + auto q_quantization = model_config->get_quantization_method( + module_prefix + "q_proj", "Linear"); + auto k_quantization = model_config->get_quantization_method( + module_prefix + "k_proj", "Linear"); + auto v_quantization = model_config->get_quantization_method( + module_prefix + "v_proj", "Linear"); + if (q_quantization->get_quant_scheme() != k_quantization->get_quant_scheme() + || q_quantization->get_quant_scheme() + != v_quantization->get_quant_scheme()) { + throw std::invalid_argument( + "fused QKV projections require the same quantization scheme in `" + + module_prefix + "{q_proj,k_proj,v_proj}`"); + } + auto o_quantization = model_config->get_quantization_method( + module_prefix + "o_proj", "Linear"); auto register_fn = [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }; qkv_proj_ = std::make_shared( hidden_size_, head_dim_, total_num_heads, total_num_kv_heads, "q_proj", "k_proj", "v_proj", register_fn, - quantization_method, use_bias, dtype, device, rank_info); + q_quantization, use_bias, dtype, device, rank_info); o_proj_ = this->register_module( - "o_proj", total_num_heads * head_dim_, hidden_size_, quantization_method, + "o_proj", total_num_heads * head_dim_, hidden_size_, o_quantization, use_output_bias, dtype, device, tp_rank, tp_size, rank_info.comm); rotary_emb_ = infinilm::layers::rotary_embedding::get_rope(model_config, device); diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index b210b0c56..1f023b8e3 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -129,7 +129,9 @@ def load_state_dict( for k in f.keys(): tensor = f.get_tensor(k) preserve_fp32 = k.endswith(preserve_fp32_suffixes) - if tensor.is_floating_point() and not preserve_fp32: + if tensor.is_floating_point() and k.endswith(".weight_scale"): + tensor = tensor.to(device=device, dtype=torch.float32) + elif tensor.is_floating_point() and not preserve_fp32: tensor = tensor.to(device=device, dtype=dtype) else: tensor = tensor.to(device=device) diff --git a/test/config/compressed_tensors_config_test.cpp b/test/config/compressed_tensors_config_test.cpp new file mode 100644 index 000000000..1418518cf --- /dev/null +++ b/test/config/compressed_tensors_config_test.cpp @@ -0,0 +1,318 @@ +#include "csrc/config/compressed_tensors_config.hpp" + +#include +#include +#include + +namespace { + +using infinilm::config::CompressedTensorsConfig; +using infinilm::config::QuantizationDynamicMode; +using infinilm::config::QuantizationStatus; +using infinilm::config::QuantizationStrategy; +using infinilm::config::QuantizationValueType; +using nlohmann::json; + +void expect(bool condition, const std::string &message) { + if (!condition) { + throw std::runtime_error(message); + } +} + +void expect_invalid_config( + const json &config, + const std::string &expected_message) { + try { + CompressedTensorsConfig::from_json(config); + } catch (const std::invalid_argument &error) { + expect( + std::string(error.what()).find(expected_message) != std::string::npos, + "unexpected validation error: " + std::string(error.what())); + return; + } + throw std::runtime_error("expected invalid compressed-tensors config"); +} + +void test_explicit_w8a8_group() { + const auto config = json::parse(R"json({ + "quant_method": "compressed-tensors", + "format": "int-quantized", + "quantization_status": "compressed", + "ignore": ["lm_head", "re:model\\.layers\\.0\\..*"], + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": { + "num_bits": 8, + "type": "int", + "symmetric": true, + "strategy": "channel", + "dynamic": false + }, + "input_activations": { + "num_bits": 8, + "type": "int", + "symmetric": true, + "strategy": "token", + "dynamic": true + } + } + } + })json"); + + const auto parsed = CompressedTensorsConfig::from_json(config); + expect(parsed.format == "int-quantized", "format was not parsed"); + expect( + parsed.quantization_status == QuantizationStatus::COMPRESSED, + "quantization_status was not parsed"); + expect(parsed.ignore.size() == 2, "ignore list was not parsed"); + expect(parsed.config_groups.size() == 1, "config group was not parsed"); + + const auto &group = parsed.config_groups.front(); + expect(group.name == "group_0", "group name was not preserved"); + expect(group.targets == std::vector{"Linear"}, "targets were not parsed"); + expect(group.weights.has_value(), "weights were not parsed"); + expect(group.input_activations.has_value(), "input activations were not parsed"); + expect(group.weights->num_bits == 8, "weight bit width was not parsed"); + expect(group.weights->type == QuantizationValueType::INT, "weight type was not parsed"); + expect(group.weights->strategy == QuantizationStrategy::CHANNEL, "weight strategy was not parsed"); + expect( + group.weights->dynamic == QuantizationDynamicMode::STATIC, + "weight dynamic mode was not parsed"); + expect( + group.input_activations->strategy == QuantizationStrategy::TOKEN, + "activation strategy was not parsed"); + expect( + group.input_activations->dynamic == QuantizationDynamicMode::DYNAMIC, + "activation dynamic mode was not parsed"); + expect(parsed.raw_config == config, "raw config was not preserved"); +} + +void test_w8a8_preset() { + const json config = { + {"quant_method", "compressed-tensors"}, + {"config_groups", {{"W8A8", {"Linear"}}}}, + }; + + const auto group = CompressedTensorsConfig::from_json(config).config_groups.front(); + expect(group.weights.has_value(), "W8A8 preset did not create weights"); + expect(group.input_activations.has_value(), "W8A8 preset did not create activations"); + expect(group.weights->num_bits == 8, "W8A8 preset did not use 8-bit weights"); + expect(group.weights->strategy == QuantizationStrategy::CHANNEL, "W8A8 weight strategy is incorrect"); + expect(group.input_activations->strategy == QuantizationStrategy::TOKEN, "W8A8 activation strategy is incorrect"); + expect(group.input_activations->dynamic == QuantizationDynamicMode::DYNAMIC, "W8A8 activations are not dynamic"); +} + +void test_defaults_and_optional_fields() { + const auto config = json::parse(R"json({ + "quant_method": "compressed-tensors", + "global_compression_ratio": 0.5, + "kv_cache_scheme": { + "num_bits": 8, + "type": "int", + "strategy": "channel", + "dynamic": "local" + }, + "config_groups": { + "INT8": ["Linear"], + "UNQUANTIZED": ["lm_head"], + "optional": { + "targets": ["re:model\\.layers\\..*"], + "format": "fakequant", + "output_activations": { + "type": "float", + "strategy": "tensor_group", + "dynamic": "local" + } + } + } + })json"); + + const auto parsed = CompressedTensorsConfig::from_json(config); + expect(parsed.format == "fakequant", "default format is incorrect"); + expect( + parsed.quantization_status == QuantizationStatus::INITIALIZED, + "default quantization status is incorrect"); + expect(parsed.ignore.empty(), "default ignore list is not empty"); + expect(parsed.global_compression_ratio == 0.5, "compression ratio was not parsed"); + expect(parsed.kv_cache_scheme.has_value(), "KV-cache scheme was not parsed"); + expect( + parsed.kv_cache_scheme->dynamic == QuantizationDynamicMode::LOCAL, + "local dynamic mode was not parsed"); + + const auto find_group = [&parsed](const std::string &name) -> const auto & { + for (const auto &group : parsed.config_groups) { + if (group.name == name) { + return group; + } + } + throw std::runtime_error("missing config group: " + name); + }; + + const auto &int8 = find_group("INT8"); + expect(int8.weights.has_value(), "INT8 preset did not create weights"); + expect(int8.input_activations.has_value(), "INT8 preset did not create activations"); + + const auto &unquantized = find_group("UNQUANTIZED"); + expect(!unquantized.weights.has_value(), "UNQUANTIZED preset created weights"); + expect(!unquantized.input_activations.has_value(), "UNQUANTIZED preset created activations"); + + const auto &optional = find_group("optional"); + expect(optional.format == "fakequant", "group format was not parsed"); + expect(optional.output_activations.has_value(), "output activations were not parsed"); + expect( + optional.output_activations->type == QuantizationValueType::FLOAT, + "floating-point value type was not parsed"); + expect( + optional.output_activations->strategy == QuantizationStrategy::TENSOR_GROUP, + "tensor-group strategy was not parsed"); +} + +void test_group_resolution() { + const auto config = json::parse(R"json({ + "quant_method": "compressed-tensors", + "ignore": [ + "model.layers.3.self_attn.o_proj", + "re:model\\.layers\\.4\\." + ], + "config_groups": { + "type_group": { + "targets": ["Linear"] + }, + "regex_group": { + "targets": ["re:model\\.layers\\.3\\."] + }, + "exact_group": { + "targets": ["model.layers.3.self_attn.q_proj"] + } + } + })json"); + const auto parsed = CompressedTensorsConfig::from_json(config); + + const auto *exact = parsed.resolve_group( + "model.layers.3.self_attn.q_proj", "Linear"); + expect(exact != nullptr, "exact group was not resolved"); + expect(exact->name == "exact_group", "exact group did not take priority"); + + const auto *regex = parsed.resolve_group( + "model.layers.3.self_attn.v_proj", "Linear"); + expect(regex != nullptr, "regex group was not resolved"); + expect(regex->name == "regex_group", "regex group did not take priority over type"); + + const auto *type = parsed.resolve_group( + "model.layers.2.self_attn.q_proj", "Linear"); + expect(type != nullptr, "module-type group was not resolved"); + expect(type->name == "type_group", "incorrect module-type group was resolved"); + + expect( + parsed.resolve_group("model.embed_tokens", "Embedding") == nullptr, + "unmatched module resolved to a group"); + expect( + parsed.is_ignored("model.layers.3.self_attn.o_proj", "Linear"), + "exact ignore rule did not match"); + expect( + parsed.resolve_group("model.layers.3.self_attn.o_proj", "Linear") + == nullptr, + "ignored module resolved to a group"); + expect( + parsed.resolve_group("model.layers.4.self_attn.q_proj", "Linear") + == nullptr, + "regex-ignored module resolved to a group"); +} + +void test_group_resolution_checks_ambiguity_at_highest_specificity() { + const auto lower_priority_tie = json::parse(R"json({ + "quant_method": "compressed-tensors", + "config_groups": { + "type_a": {"targets": ["Linear"]}, + "type_b": {"targets": ["Linear"]}, + "z_exact": {"targets": ["model.layers.0.self_attn.q_proj"]} + } + })json"); + + const auto parsed = CompressedTensorsConfig::from_json(lower_priority_tie); + const auto *resolved = parsed.resolve_group( + "model.layers.0.self_attn.q_proj", "Linear"); + expect(resolved != nullptr, "exact group was not resolved"); + expect( + resolved->name == "z_exact", + "a lower-specificity tie incorrectly overrode the exact match"); + + const auto highest_priority_tie = json::parse(R"json({ + "quant_method": "compressed-tensors", + "config_groups": { + "exact_a": {"targets": ["model.layers.0.self_attn.q_proj"]}, + "exact_b": {"targets": ["model.layers.0.self_attn.q_proj"]} + } + })json"); + + try { + CompressedTensorsConfig::from_json(highest_priority_tie).resolve_group("model.layers.0.self_attn.q_proj", "Linear"); + } catch (const std::invalid_argument &error) { + expect( + std::string(error.what()).find("equal specificity") != std::string::npos, + "unexpected group-resolution error: " + std::string(error.what())); + return; + } + throw std::runtime_error("expected equally specific groups to be ambiguous"); +} + +void test_validation_errors_include_paths() { + const auto invalid_dynamic = json::parse(R"json({ + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "input_activations": {"dynamic": "global"} + } + } + })json"); + const auto invalid_num_bits = json::parse(R"json({ + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 0} + } + } + })json"); + + expect_invalid_config(json::array(), "quantization_config: expected an object"); + expect_invalid_config( + {{"quant_method", "compressed-tensors"}}, + "quantization_config.config_groups"); + expect_invalid_config( + {{"quant_method", "gptq"}, {"config_groups", json::object()}}, + "quantization_config.quant_method"); + expect_invalid_config( + {{"quant_method", "compressed-tensors"}, {"config_groups", {{"group_0", json::object()}}}}, + "quantization_config.config_groups.group_0.targets"); + expect_invalid_config( + invalid_dynamic, + "quantization_config.config_groups.group_0.input_activations.dynamic"); + expect_invalid_config( + invalid_num_bits, + "quantization_config.config_groups.group_0.weights.num_bits"); + expect_invalid_config( + {{"quant_method", "compressed-tensors"}, {"config_groups", {{"FP8", {"Linear"}}}}}, + "unsupported preset \"FP8\""); +} + +} // namespace + +int main() { + try { + test_explicit_w8a8_group(); + test_w8a8_preset(); + test_defaults_and_optional_fields(); + test_group_resolution(); + test_group_resolution_checks_ambiguity_at_highest_specificity(); + test_validation_errors_include_paths(); + } catch (const std::exception &error) { + std::cerr << "compressed_tensors_config_test failed: " << error.what() << '\n'; + return 1; + } + + std::cout << "compressed_tensors_config_test passed\n"; + return 0; +} diff --git a/test/config/module_target_matcher_test.cpp b/test/config/module_target_matcher_test.cpp new file mode 100644 index 000000000..159b9763c --- /dev/null +++ b/test/config/module_target_matcher_test.cpp @@ -0,0 +1,114 @@ +#include "csrc/config/module_target_matcher.hpp" + +#include +#include +#include +#include +#include + +namespace { + +using MatchKind = infinilm::config::ModuleTargetMatcher::MatchKind; +using infinilm::config::ModuleTargetMatcher; + +constexpr std::string_view MODULE_NAME = "model.layers.3.self_attn.q_proj"; +constexpr std::string_view MODULE_TYPE = "Linear"; + +void expect(bool condition, const std::string &message) { + if (!condition) { + throw std::runtime_error(message); + } +} + +void test_exact_name_matching() { + expect( + ModuleTargetMatcher::match(MODULE_NAME, MODULE_NAME, MODULE_TYPE) + == MatchKind::EXACT_NAME, + "full module name did not match exactly"); + expect( + ModuleTargetMatcher::match("q_proj", MODULE_NAME, MODULE_TYPE) + == MatchKind::NONE, + "module name suffix was treated as an exact match"); +} + +void test_module_type_matching() { + expect( + ModuleTargetMatcher::match("Linear", MODULE_NAME, MODULE_TYPE) + == MatchKind::MODULE_TYPE, + "module type did not match"); + expect( + ModuleTargetMatcher::match("Embedding", MODULE_NAME, MODULE_TYPE) + == MatchKind::NONE, + "unrelated module type matched"); +} + +void test_regex_matching_starts_at_module_name() { + expect( + ModuleTargetMatcher::match( + R"(re:model\.layers\.\d+\.self_attn\..*_proj$)", + MODULE_NAME, + MODULE_TYPE) + == MatchKind::REGEX, + "regular expression did not match the module name"); + expect( + ModuleTargetMatcher::match("re:q_proj$", MODULE_NAME, MODULE_TYPE) + == MatchKind::NONE, + "regular expression searched beyond the start of the module name"); + expect( + ModuleTargetMatcher::match("re:.*q_proj$", MODULE_NAME, MODULE_TYPE) + == MatchKind::REGEX, + "explicit suffix regular expression did not match"); + expect( + ModuleTargetMatcher::match("re:model\\.layers", MODULE_NAME, MODULE_TYPE) + == MatchKind::REGEX, + "regular expression was incorrectly required to match the full name"); +} + +void test_match_any_returns_the_most_specific_match() { + expect( + ModuleTargetMatcher::match_any( + {"Linear", "re:model\\.layers", std::string(MODULE_NAME)}, + MODULE_NAME, + MODULE_TYPE) + == MatchKind::EXACT_NAME, + "exact name did not take priority"); + expect( + ModuleTargetMatcher::match_any( + {"Linear", "re:model\\.layers"}, MODULE_NAME, MODULE_TYPE) + == MatchKind::REGEX, + "regular expression did not take priority over module type"); + expect( + ModuleTargetMatcher::match_any({}, MODULE_NAME, MODULE_TYPE) + == MatchKind::NONE, + "empty target list matched"); +} + +void test_invalid_regex_is_rejected() { + try { + ModuleTargetMatcher::match("re:[", MODULE_NAME, MODULE_TYPE); + } catch (const std::invalid_argument &error) { + expect( + std::string(error.what()).find("re:[") != std::string::npos, + "invalid-regex error did not identify the target"); + return; + } + throw std::runtime_error("invalid regular expression was accepted"); +} + +} // namespace + +int main() { + try { + test_exact_name_matching(); + test_module_type_matching(); + test_regex_matching_starts_at_module_name(); + test_match_any_returns_the_most_specific_match(); + test_invalid_regex_is_rejected(); + } catch (const std::exception &error) { + std::cerr << "module_target_matcher_test failed: " << error.what() << '\n'; + return 1; + } + + std::cout << "module_target_matcher_test passed\n"; + return 0; +} diff --git a/test/config/quant_config_test.cpp b/test/config/quant_config_test.cpp new file mode 100644 index 000000000..08b728af6 --- /dev/null +++ b/test/config/quant_config_test.cpp @@ -0,0 +1,188 @@ +#include "csrc/config/quant_config.hpp" + +#include +#include +#include +#include + +namespace { + +using infinilm::config::QuantConfig; +using infinilm::config::QuantizationStatus; +using infinilm::quantization::QuantScheme; +using nlohmann::json; + +void expect(bool condition, const std::string &message) { + if (!condition) { + throw std::runtime_error(message); + } +} + +void expect_scheme( + const QuantConfig &config, + std::string_view module_name, + std::string_view module_type, + QuantScheme expected, + const std::string &message) { + expect( + config.get_quantization_method(module_name, module_type) + ->get_quant_scheme() + == expected, + message); +} + +void expect_invalid_method( + const json &config, + const std::string &expected_message) { + try { + QuantConfig(config).get_quantization_method( + "model.layers.0.mlp.up_proj", "Linear"); + } catch (const std::invalid_argument &error) { + expect( + std::string(error.what()).find(expected_message) != std::string::npos, + "unexpected quantization-method error: " + std::string(error.what())); + return; + } + throw std::runtime_error("expected unsupported quantization method"); +} + +json real_qwen3_w8a8_config() { + return json::parse(R"json({ + "quant_method": "compressed-tensors", + "format": "int-quantized", + "quantization_status": "compressed", + "ignore": ["lm_head"], + "config_groups": { + "group_0": { + "format": "int-quantized", + "input_activations": { + "actorder": null, + "block_structure": null, + "dynamic": true, + "group_size": null, + "num_bits": 8, + "observer": null, + "observer_kwargs": {}, + "strategy": "token", + "symmetric": true, + "type": "int" + }, + "output_activations": null, + "targets": ["Linear"], + "weights": { + "actorder": null, + "block_structure": null, + "dynamic": false, + "group_size": null, + "num_bits": 8, + "observer": "mse", + "observer_kwargs": {}, + "strategy": "channel", + "symmetric": true, + "type": "int" + } + } + }, + "global_compression_ratio": null, + "kv_cache_scheme": null, + "sparsity_config": {}, + "transform_config": {}, + "version": "0.13.0" + })json"); +} + +void test_real_qwen3_w8a8_config() { + const auto config = real_qwen3_w8a8_config(); + const QuantConfig quant_config(config); + + const auto &parsed = quant_config.get_compressed_tensors_config(); + expect(parsed.has_value(), "real Qwen3 config was not parsed"); + expect( + parsed->quantization_status == QuantizationStatus::COMPRESSED, + "real Qwen3 quantization status was not parsed"); + expect( + parsed->raw_config.at("version") == "0.13.0", + "real Qwen3 compressed-tensors version was not preserved"); + + for (const auto *module_name : { + "model.layers.0.self_attn.q_proj", + "model.layers.0.self_attn.k_proj", + "model.layers.0.self_attn.v_proj", + "model.layers.0.self_attn.o_proj", + "model.layers.0.mlp.gate_proj", + "model.layers.0.mlp.up_proj", + "model.layers.0.mlp.down_proj", + }) { + expect_scheme( + quant_config, + module_name, + "Linear", + QuantScheme::COMPRESSED_TENSOR_W8A8I8, + std::string("real Qwen3 config did not quantize ") + module_name); + } + expect_scheme( + quant_config, + "lm_head", + "Linear", + QuantScheme::NONE, + "real Qwen3 config did not preserve its lm_head ignore rule"); +} + +void test_unquantized_and_unmatched_modules() { + const QuantConfig quant_config(json{ + {"quant_method", "compressed-tensors"}, + {"config_groups", {{"UNQUANTIZED", {"lm_head"}}}}, + }); + + expect_scheme( + quant_config, + "lm_head", + "Linear", + QuantScheme::NONE, + "UNQUANTIZED group selected a quantized scheme"); + expect_scheme( + quant_config, + "model.embed_tokens", + "Embedding", + QuantScheme::NONE, + "unmatched module was quantized"); +} + +void test_unsupported_schemes_are_rejected() { + const json fakequant_config = { + {"quant_method", "compressed-tensors"}, + {"format", "fakequant"}, + {"config_groups", {{"W8A8", {"Linear"}}}}, + }; + auto w4a8_config = real_qwen3_w8a8_config(); + w4a8_config["config_groups"]["group_0"]["weights"]["num_bits"] = 4; + + expect_invalid_method(fakequant_config, "format `fakequant`"); + expect_invalid_method(w4a8_config, "group `group_0`"); +} + +void test_existing_methods_are_preserved() { + const QuantConfig quant_config(json{{"quant_method", "awq"}}); + expect( + quant_config.get_quantization_method("model.layers.0.mlp.up_proj", "Linear") + ->get_quant_scheme() + == QuantScheme::AWQ_W4A16, + "module-aware lookup changed the existing AWQ selection"); +} + +} // namespace + +int main() { + try { + test_real_qwen3_w8a8_config(); + test_unquantized_and_unmatched_modules(); + test_unsupported_schemes_are_rejected(); + test_existing_methods_are_preserved(); + } catch (const std::exception &error) { + std::cerr << "quant_config_test failed: " << error.what() << '\n'; + return 1; + } + + std::cout << "quant_config_test passed\n"; + return 0; +} diff --git a/test/quantization/compressed_tensors_test.cpp b/test/quantization/compressed_tensors_test.cpp new file mode 100644 index 000000000..1e156ce5a --- /dev/null +++ b/test/quantization/compressed_tensors_test.cpp @@ -0,0 +1,98 @@ +#include "csrc/layers/quantization/compressed_tensors.hpp" + +#include +#include +#include +#include + +namespace { + +using infinilm::quantization::CompressedTensors; +using infinilm::quantization::ParamDescriptor; + +void expect(bool condition, const std::string &message) { + if (!condition) { + throw std::runtime_error(message); + } +} + +const ParamDescriptor &find_descriptor( + const std::vector &layout, + const std::string &name) { + for (const auto &descriptor : layout) { + if (descriptor.name == name) { + return descriptor; + } + } + throw std::runtime_error("missing parameter descriptor: " + name); +} + +void expect_partition( + const ParamDescriptor &descriptor, + int split_dim, + int tp_rank, + int tp_size) { + expect( + descriptor.split_dim == split_dim, + descriptor.name + " has an incorrect split dimension"); + expect( + descriptor.tp_rank == tp_rank, + descriptor.name + " has an incorrect TP rank"); + expect( + descriptor.tp_size == tp_size, + descriptor.name + " has an incorrect TP size"); +} + +void test_replicated_layout() { + const CompressedTensors quantization(nlohmann::json::object()); + const auto layout = quantization.get_param_layout( + 128, 256, -1, 0, 1, -1, infinicore::DataType::F32, true); + + expect(layout.size() == 3, "replicated layout has an incorrect parameter count"); + const auto &weight = find_descriptor(layout, "weight"); + const auto &scale = find_descriptor(layout, "weight_scale"); + const auto &bias = find_descriptor(layout, "bias"); + expect(weight.shape == std::vector{256, 128}, "weight shape is incorrect"); + expect(weight.dtype == infinicore::DataType::I8, "weight dtype is not INT8"); + expect(scale.shape == std::vector{256, 1}, "weight-scale shape is incorrect"); + expect(scale.dtype == infinicore::DataType::F32, "weight-scale dtype is not FP32"); + expect_partition(weight, -1, 0, 1); + expect_partition(scale, -1, 0, 1); + expect_partition(bias, -1, 0, 1); +} + +void test_column_parallel_layout() { + const CompressedTensors quantization(nlohmann::json::object()); + const auto layout = quantization.get_param_layout( + 128, 256, 0, 1, 2, -1, infinicore::DataType::F32, true); + + expect_partition(find_descriptor(layout, "weight"), 0, 1, 2); + expect_partition(find_descriptor(layout, "weight_scale"), 0, 1, 2); + expect_partition(find_descriptor(layout, "bias"), 0, 1, 2); +} + +void test_row_parallel_layout() { + const CompressedTensors quantization(nlohmann::json::object()); + const auto layout = quantization.get_param_layout( + 128, 256, 1, 1, 2, -1, infinicore::DataType::F32, true); + + expect_partition(find_descriptor(layout, "weight"), 1, 1, 2); + expect_partition(find_descriptor(layout, "weight_scale"), -1, 0, 1); + expect_partition(find_descriptor(layout, "bias"), -1, 0, 1); +} + +} // namespace + +int main() { + try { + test_replicated_layout(); + test_column_parallel_layout(); + test_row_parallel_layout(); + } catch (const std::exception &error) { + std::cerr << "compressed_tensors_test failed: " << error.what() << '\n'; + return 1; + } + + std::cout << "compressed_tensors_test passed\n"; + return 0; +} diff --git a/test/test_modeling_utils.py b/test/test_modeling_utils.py new file mode 100644 index 000000000..4cc885990 --- /dev/null +++ b/test/test_modeling_utils.py @@ -0,0 +1,40 @@ +import os +import tempfile +import unittest + +import torch +from infinilm.modeling_utils import load_state_dict +from safetensors.torch import save_file + + +class LoadStateDictTest(unittest.TestCase): + def test_converts_weight_scale_to_float32(self): + with tempfile.TemporaryDirectory() as temp_dir: + checkpoint_path = os.path.join(temp_dir, "model.safetensors") + save_file( + { + "model.layers.0.self_attn.q_proj.weight_scale": torch.ones( + 4, 1, dtype=torch.bfloat16 + ), + "model.layers.0.input_layernorm.weight": torch.ones( + 4, dtype=torch.bfloat16 + ), + }, + checkpoint_path, + metadata={"format": "pt"}, + ) + + state_dict = load_state_dict(checkpoint_path, dtype=torch.float16) + + self.assertEqual( + state_dict["model.layers.0.self_attn.q_proj.weight_scale"].dtype, + torch.float32, + ) + self.assertEqual( + state_dict["model.layers.0.input_layernorm.weight"].dtype, + torch.float16, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/xmake.lua b/xmake.lua index dd04bbbee..bdbdc3d93 100644 --- a/xmake.lua +++ b/xmake.lua @@ -21,6 +21,8 @@ end add_includedirs("third_party/spdlog/include") add_includedirs("third_party/json/single_include/") +local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") + target("_infinilm") add_packages("pybind11") set_default(false) @@ -28,8 +30,6 @@ target("_infinilm") set_languages("cxx17") set_kind("shared") - local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") - -- add_includedirs("csrc", { public = false }) -- add_includedirs("csrc/pybind11", { public = false }) add_includedirs(INFINI_ROOT.."/include", { public = true }) @@ -44,3 +44,66 @@ target("_infinilm") set_installdir("python/infinilm") target_end() + +target("compressed_tensors_config_test") + set_default(false) + set_kind("binary") + set_languages("cxx17") + + add_includedirs(".") + add_files("test/config/compressed_tensors_config_test.cpp") + add_files("csrc/config/compressed_tensors_config.cpp") + add_files("csrc/config/module_target_matcher.cpp") +target_end() + +target("quant_config_test") + set_default(false) + set_kind("binary") + set_languages("cxx17") + + add_includedirs(".") + add_includedirs(INFINI_ROOT.."/include") + add_linkdirs(INFINI_ROOT.."/lib") + add_links("infinicore_cpp_api", "infiniop", "infinirt", "infiniccl") + add_runenvs("LD_LIBRARY_PATH", INFINI_ROOT.."/lib", {pathenv = true}) + add_files("test/config/quant_config_test.cpp") + add_files("csrc/config/compressed_tensors_config.cpp") + add_files("csrc/config/module_target_matcher.cpp") + add_files("csrc/config/quant_config.cpp") + add_files("csrc/global_state/global_state.cpp") + add_files("csrc/layers/quantization/base_quantization.cpp") + add_files("csrc/layers/quantization/none_quantization.cpp") + add_files("csrc/layers/quantization/compressed_tensors.cpp") + add_files("csrc/layers/quantization/awq.cpp") + add_files("csrc/layers/quantization/awq_marlin.cpp") + add_files("csrc/layers/quantization/gptq.cpp") + add_files("csrc/layers/quantization/gptq_marlin.cpp") + add_files("csrc/layers/quantization/gptq_qy.cpp") + add_files("csrc/layers/quantization/marlin_utils.cpp") + add_files("csrc/layers/quantization/mxfp4.cpp") +target_end() + +target("compressed_tensors_test") + set_default(false) + set_kind("binary") + set_languages("cxx17") + + add_includedirs(".") + add_includedirs(INFINI_ROOT.."/include") + add_linkdirs(INFINI_ROOT.."/lib") + add_links("infinicore_cpp_api", "infiniop", "infinirt", "infiniccl") + add_runenvs("LD_LIBRARY_PATH", INFINI_ROOT.."/lib", {pathenv = true}) + add_files("test/quantization/compressed_tensors_test.cpp") + add_files("csrc/layers/quantization/base_quantization.cpp") + add_files("csrc/layers/quantization/compressed_tensors.cpp") +target_end() + +target("module_target_matcher_test") + set_default(false) + set_kind("binary") + set_languages("cxx17") + + add_includedirs(".") + add_files("test/config/module_target_matcher_test.cpp") + add_files("csrc/config/module_target_matcher.cpp") +target_end()