From 5223e85f64bbb5625cc483623faa01afe9a288b4 Mon Sep 17 00:00:00 2001 From: Jennifer Lee Date: Tue, 13 Jan 2026 19:24:10 -0800 Subject: [PATCH] Tune xgboost hyperparams (#334) Summary: - Add `ml_selector_trainer_utils` to expose train test split and ml model training function - Add `ml_selector_trainer_tuner` to tune hyperparameters that are used in xgboost training - Update readme with details about tuning process Differential Revision: D90044668 --- cli/args/TrainArgs.h | 9 + tools/ml_selector/BUCK | 48 +- tools/ml_selector/CMakeLists.txt | 4 + tools/ml_selector/README.md | 57 ++ tools/ml_selector/ml_features.cpp | 22 + tools/ml_selector/ml_features.h | 10 +- tools/ml_selector/ml_selector_trainer.cpp | 531 ++---------------- .../ml_selector/ml_selector_trainer_utils.cpp | 467 +++++++++++++++ tools/ml_selector/ml_selector_trainer_utils.h | 104 ++++ tools/ml_selector/ml_selector_tuner.cpp | 398 +++++++++++++ tools/ml_selector/ml_selector_tuner.h | 122 ++++ tools/ml_selector/tests/BUCK | 2 + .../tests/test_mlSelectorTrainer.cpp | 51 ++ tools/training/train_params.h | 1 + 14 files changed, 1324 insertions(+), 502 deletions(-) create mode 100644 tools/ml_selector/ml_selector_trainer_utils.cpp create mode 100644 tools/ml_selector/ml_selector_trainer_utils.h create mode 100644 tools/ml_selector/ml_selector_tuner.cpp create mode 100644 tools/ml_selector/ml_selector_tuner.h diff --git a/cli/args/TrainArgs.h b/cli/args/TrainArgs.h index adb576e0c..8e08bd45d 100644 --- a/cli/args/TrainArgs.h +++ b/cli/args/TrainArgs.h @@ -107,6 +107,12 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { 0, false, "Enables pareto frontier training. This will output a directory containing all compressors in the pareto frontier."); + parser.addCommandFlag( + cmd(), + kTuneHyperparams, + 0, + false, + "Enable hyperparameter tuning for ML selector training."); } explicit TrainArgs(const arg::ParsedArgs& parsed) @@ -181,6 +187,8 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { parsed.cmdHasFlag(cmd(), kNoAceSuccessors); trainParams.noClustering = parsed.cmdHasFlag(cmd(), kNoClustering); + trainParams.tuneHyperparams = + parsed.cmdHasFlag(cmd(), kTuneHyperparams); trainParams.compressorGenFunc = custom_parsers::createCompressorFromSerialized; } @@ -223,6 +231,7 @@ class TrainArgs : public GlobalArgs, public ProfileArgs { inline static const std::string kMaxFileSizeMb = "max-file-size-mb"; inline static const std::string kMaxTotalSizeMb = "max-total-size-mb"; inline static const std::string kParetoFrontier = "pareto-frontier"; + inline static const std::string kTuneHyperparams = "tune-hyperparams"; }; } // namespace openzl::cli diff --git a/tools/ml_selector/BUCK b/tools/ml_selector/BUCK index ce862d884..238f103c0 100644 --- a/tools/ml_selector/BUCK +++ b/tools/ml_selector/BUCK @@ -1,5 +1,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. +# @noautodeps + load("../../defs.bzl", "zs_cxxlibrary", "zs_library") oncall("data_compression") @@ -22,14 +24,29 @@ zs_cxxlibrary( ], deps = [ "../..:zstronglib", - "../../cpp:openzl_cpp", - "../../tests/datagen:datagen", # @manual - "../training:train_common", "../training/graph_mutation:graph_mutation", "../training/sample_collection:sample_collection", - "../training/utils:training_utils", ":ml_features", ":ml_selector_graph", + ":ml_selector_trainer_utils", + ":ml_selector_tuner", + ], + exported_deps = [ + "../training:train_common", + "../training/utils:training_utils", + ], +) + +zs_cxxlibrary( + name = "ml_selector_trainer_utils", + srcs = ["ml_selector_trainer_utils.cpp"], + headers = ["ml_selector_trainer_utils.h"], + compiler_flags = [ + "-Wno-unused-exception-parameter", # For XGBoost warnings + ], + deps = [ + "..:logger", + "../../cpp:openzl_cpp", ], exported_external_deps = [ ("xgboost", None, "xgboost-cpp"), @@ -44,7 +61,28 @@ zs_cxxlibrary( deps = [ "..:logger", "../..:zstronglib", - "../../cpp:openzl_cpp", + ], + exported_deps = [ "../training/utils:training_utils", ], ) + +zs_cxxlibrary( + name = "ml_selector_tuner", + srcs = ["ml_selector_tuner.cpp"], + headers = ["ml_selector_tuner.h"], + compiler_flags = [ + "-Wno-unused-exception-parameter", # For XGBoost warnings + ], + deps = [ + "..:io", + "../..:zstronglib", + ":ml_features", + ":ml_selector_graph", + ":ml_selector_trainer_utils", + ], + external_deps = [ + ("xgboost", None, "xgboost-cpp"), + ("xgboost", None, "dmlc"), + ], +) diff --git a/tools/ml_selector/CMakeLists.txt b/tools/ml_selector/CMakeLists.txt index 416b42bbb..9279388d6 100644 --- a/tools/ml_selector/CMakeLists.txt +++ b/tools/ml_selector/CMakeLists.txt @@ -67,6 +67,10 @@ add_library(ml_selector_graph ml_selector_graph.h ml_selector_trainer.cpp ml_selector_trainer.h + ml_selector_trainer_utils.cpp + ml_selector_trainer_utils.h + ml_selector_tuner.cpp + ml_selector_tuner.h ml_features.cpp ml_features.h ) diff --git a/tools/ml_selector/README.md b/tools/ml_selector/README.md index e3b1388ea..9102c93dc 100644 --- a/tools/ml_selector/README.md +++ b/tools/ml_selector/README.md @@ -76,3 +76,60 @@ When you compress data with a trained model: 1. Features are extracted from the input 2. The `gbtModel` uses those features to predict which successor to use 3. Compress data with selected successor + +# ML Selector Tuner + +Tune XGBoost hyperparameters to optimize for compression size and speed. + +## How to Use + +```bash +buck2 run @//mode/opt tools/ml_selector:ml_selector_tuner -- [population_size] [survival_rate] [max_iterations] [convergence_threshold] [mutation_rate] [compression_weight] +``` + +### Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `population_size` | 20 | Number of hyperparameter sets per generation | +| `survival_rate` | 0.25 | Fraction of population that survives to next generation | +| `max_iterations` | 10 | Maximum number of generations to run | +| `convergence_threshold` | 2 | Generations without improvement before stopping | +| `mutation_rate` | 0.25 | Probability of random modification to a hyperparameter | +| `compression_weight` | 0.75 | Weight for compression size vs. time (0.0–1.0) | + +## How It Works + +The tuner uses a genetic algorithm to search for optimal XGBoost hyperparameters: + +- **Initialization**: Generates an initial population by dividing each hyperparameter range into equal intervals and randomly sampling exactly once from each interval. + +- **Evaluation**: Scores each candidate by training an XGBoost model and compressing test inputs, computing a weighted score of compression size and time + +- **Selection**: Ranks candidates by score and selects the top performers (based on `survival_rate`) as parents for the next generation + +- **Crossover**: Creates child configurations by randomly inheriting each hyperparameter from one of two randomly selected parents + +- **Mutation**: Applies random noise to genes with probability `mutation_rate`, clamping to valid bounds + +- **Convergence**: Stops when max iterations are reached or the best score hasn't improved by more than 0.0005% for `convergence_threshold` consecutive generations + +- **Output**: Returns the best hyperparameter configuration and compares it against default XGBoost parameters + +### Tuned Hyperparameters + +The following XGBoost hyperparameters are searched: + +| Parameter | Range | Description | +|-----------|-------|-------------| +| `learning_rate` | 0.001–1.0 | How much model adjusts weights in response to estimated error during training | +| `min_child_weight` | 0–30 | Minimum sum of instance weight in a child | +| `subsample` | 0.1–1.0 | Fraction of samples used per tree | +| `colsample_bynode` | 0.1–1.0 | Fraction of features used per split | +| `max_depth` | 0–20 | Maximum tree depth | +| `max_leaves` | 0–60 | Maximum number of leaves per tree | +| `reg_alpha` | 0–10 | L1 regularization | +| `gamma` | 0–20 | Minimum loss reduction for split | +| `num_boost_round` | 5–60 | Number of boosting rounds | +| `max_delta_step` | 0–10 | Maximum delta step for weight estimation | +| `scale_pos_weight` | 0.5–20 | Balance of positive/negative weights (for imbalanced datasets) | diff --git a/tools/ml_selector/ml_features.cpp b/tools/ml_selector/ml_features.cpp index 3837715cf..a9f466a10 100644 --- a/tools/ml_selector/ml_features.cpp +++ b/tools/ml_selector/ml_features.cpp @@ -2,6 +2,7 @@ #include "tools/ml_selector/ml_features.h" #include #include +#include #include #include #include "openzl/zl_reflection.h" @@ -26,6 +27,27 @@ static size_t compress(CCtx& cctx, Compressor& compressor, const Input& input) return cctx.compressOne(input).size(); } +ChoiceFunction makeWeightedChoiceFunc(float weight) +{ + return [weight](std::vector& targets) { + std::vector result; + for (size_t i = 0; i < targets.size(); i++) { + std::priority_queue< + std::pair, + std::vector>, + std::greater>> + pq; + for (const auto& it : targets[i]) { + float score = weight * it.second.at("size") + + (1 - weight) * it.second.at("ctime"); + pq.emplace(score, it.first); + } + result.push_back((float)pq.top().second); + } + return result; + }; +} + std::vector minSizeChoiceFunc(std::vector& targets) { std::vector result; diff --git a/tools/ml_selector/ml_features.h b/tools/ml_selector/ml_features.h index f0f2a2f69..a233e3ab3 100644 --- a/tools/ml_selector/ml_features.h +++ b/tools/ml_selector/ml_features.h @@ -1,5 +1,6 @@ // Copyright (c) Meta Platforms, Inc. and affiliates. #pragma once +#include #include #include #include @@ -19,7 +20,8 @@ using FeatureMap = VECTOR(LabeledFeature); * @returns vector containing index of "best" possible successors. Note that * index is cast as float due to XGBoost requirement. */ -using ChoiceFunction = std::vector (*)(std::vector& targets); +using ChoiceFunction = + std::function(std::vector&)>; /** * Container for processed ML training data with features and labels. @@ -38,6 +40,12 @@ struct ProcessedMLTrainingSamples { std::vector featurePtrNames; }; +/** + * Make weighted choice function that chooses index successor with best 'score' + * which is calculated by weight * cSize + (1 - weight) * cTime. + */ +ChoiceFunction makeWeightedChoiceFunc(float weight); + /** * Default choice function that chooses index successor with minimum compression * size as "best". diff --git a/tools/ml_selector/ml_selector_trainer.cpp b/tools/ml_selector/ml_selector_trainer.cpp index 9263104d0..2fd0a20f0 100644 --- a/tools/ml_selector/ml_selector_trainer.cpp +++ b/tools/ml_selector/ml_selector_trainer.cpp @@ -1,514 +1,21 @@ // Copyright (c) Meta Platforms, Inc. and affiliates. #include "tools/ml_selector/ml_selector_trainer.h" #include -#include -#include #include #include #include "ml_features.h" #include "openzl/zl_reflection.h" #include "src/openzl/compress/selectors/ml/features.h" -#include "tools/logger/Logger.h" #include "tools/ml_selector/ml_selector_graph.h" +#include "tools/ml_selector/ml_selector_trainer_utils.h" +#include "tools/ml_selector/ml_selector_tuner.h" #include "tools/training/graph_mutation/graph_mutation_utils.h" #include "tools/training/sample_collection/training_sample_collector.h" -// Suppress warnings for XGBoost headers -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wcast-qual" -#pragma GCC diagnostic ignored "-Wfloat-equal" -#pragma GCC diagnostic ignored "-Wcast-align" -#include -#include -#pragma GCC diagnostic pop -using namespace openzl::tools::logger; - namespace openzl::training { const std::string ML_SELECTOR_GRAPH_NAME = "mlSelector"; -namespace { -/** - * Macro provided by xgboost cpp to guard all calls - */ -#define safe_xgboost(call) \ - { \ - int err = (call); \ - if (err != 0) { \ - throw std::runtime_error( \ - std::string(__FILE__) + ":" + std::to_string(__LINE__) \ - + ": error in " + #call + ":" + XGBGetLastError()); \ - } \ - } - -/** - * Contains both the 2D feature matrices and their flattened versions. The 2D - * format is for train test split, while the flattened format is required by - * XGBoost's DMatrix API. - * - * Naming follows scikit-learn conventions: - * - X: Feature matrix - * - y: Label vector - */ -struct FeatureData { - std::vector> X; - std::vector y; - std::vector XFlat; - - FeatureData( - const std::vector>& X_, - const std::vector& y_) - : X(X_), y(y_) - { - for (const auto& x : X) { - XFlat.insert(XFlat.end(), x.begin(), x.end()); - } - } -}; - -/** - * Holds the train/test split data for XGBoost model training. - * - */ -struct TestTrainData { - FeatureData train; - FeatureData test; -}; - -/** - * Owns the memory for a GBTPredictor. - */ -struct GBTPredictorWrapper { - /// Node arrays for each tree (indexed by tree) - std::vector>> core_nodes_; - /// Tree arrays for each forest (indexed by forest) - std::vector>> core_trees_; - /// All forests in the model - std::unique_ptr> core_forests_; - std::unique_ptr core_predictor_; -}; - -} // namespace - -/** - * Extracts a numeric value from an XGBoost JSON node. - * - * Handles both JsonNumber (float) and JsonInteger, converting - * the value to a float. Throws exception if unable to extract value. - * - * @param json The XGBoost JSON node to extract the value from - * @return The numeric value as a float, or -1.0f if not a numeric type - */ -static float extractNumericVal(const xgboost::Json& json) -{ - if (xgboost::IsA(json)) { - return xgboost::get(json); - } else if (xgboost::IsA(json)) { - return static_cast( - xgboost::get(json)); - } - throw Exception( - "Expected JsonNumber. Unable to extract numeric value from XGBoost JSON node"); -} - -/** - * Recursively computes the maximum node ID in an XGBoost decision tree. - * - * Traverses the JSON format for XGBoost model. Used to - * pre-allocate a contiguous vector for storing parsed tree nodes. - * - * @param json The root JSON node of the XGBoost tree - * @return The maximum node ID found in the tree - */ -static int findMaxNodeId(const xgboost::Json& json) -{ - auto& objMap = xgboost::get(json); - int maxId = extractNumericVal(json["nodeid"]); - - if (objMap.find("children") != objMap.end()) { - auto& children = - xgboost::get(json["children"]); - for (const auto& child : children) { - maxId = std::max(maxId, findMaxNodeId(child)); - } - } - return maxId; -} - -/** - * Deleter functor to manage XGBoost BoosterHandle resource. Automatically calls - * XGBoosterFree when unique_pointer goes out of scope - */ -struct BoosterHandleDeleter { - using pointer = BoosterHandle; - void operator()(BoosterHandle handle) const - { - if (handle != nullptr) { - XGBoosterFree(handle); - } - } -}; - -using BoosterUniquePtr = std::unique_ptr; - -/** - * Deleter functor to manage XGBoost DMatrixHandle resource. Automatically calls - * XGDMatrixFree when unique_pointer goes out of scope - */ -struct DMatrixHandleDeleter { - using pointer = DMatrixHandle; - void operator()(DMatrixHandle handle) const - { - if (handle != nullptr) { - XGDMatrixFree(handle); - } - } -}; - -using DMatrixUniquePtr = std::unique_ptr; - -/** - * Recursively parses an XGBoost decision tree node from JSON into a - * GBTPredictor_Node structure. - * - * The parsed node is stored directly into the pre-allocated nodes vector - * at the position corresponding to its nodeid. - * - * @param nodeJson The JSON representation of the tree node - * @param nodes Output vector where parsed nodes are stored by their ID - * - * @throws Exception if nodeid exceeds the bounds of the nodes vector - */ -static void parseXGBoostNode( - const xgboost::Json& nodeJson, - std::vector& nodes) -{ - GBTPredictor_Node node{}; - - auto& objMap = xgboost::get(nodeJson); - - size_t nodeid = extractNumericVal(nodeJson["nodeid"]); - if (objMap.find("leaf") != objMap.end()) { - // Leaf node - float leafValue = extractNumericVal(nodeJson["leaf"]); - node.value = leafValue; - node.featureIdx = -1; - - Logger::log_c(VERBOSE1, "Leaf node: %f\n", leafValue); - } else { - // Internal node - // Parse feature index from "fN" format - std::string splitFeature = - xgboost::get(nodeJson["split"]); - float threshold = extractNumericVal(nodeJson["split_condition"]); - int yesChild = extractNumericVal(nodeJson["yes"]); - int noChild = extractNumericVal(nodeJson["no"]); - int missingChild = extractNumericVal(nodeJson["missing"]); - - node.featureIdx = - std::stoi(splitFeature.substr(1)); // Skip the 'f' prefix - node.leftChildIdx = yesChild; - node.rightChildIdx = noChild; - node.missingChildIdx = missingChild; - node.value = threshold; - - Logger::log_c(VERBOSE1, "Internal Node %d: ,", nodeid); - - Logger::log_c( - VERBOSE1, "split[%s < %f] \n", splitFeature.c_str(), threshold); - - Logger::log_c( - VERBOSE1, - ", Yes node %d, no node %d, missing node %d\n", - yesChild, - noChild, - missingChild); - - // Recursively process children - if (objMap.find("children") != objMap.end()) { - auto& children = xgboost::get( - nodeJson["children"]); - for (const auto& child : children) { - parseXGBoostNode(child, nodes); - } - } - } - - if (nodeid >= nodes.size()) { - throw Exception("Node id is out of range"); - } - nodes[nodeid] = node; -} - -/** - * Splits feature and label data into training and test sets. Used for - * evaluating model performance during XGBoost boosting rounds. - * - * @param features Feature vectors to split - * @param labels Corresponding labels to split - * @param testSize Fraction of data to reserve for testing - * @param shuffle Whether to shuffle (in place) data before splitting - * @param randomState Seed for shuffling - * - * @return TrainTestData containing XTrain, XTest, yTrain, and yTest splits - * - * @throws Exception if features and labels have different sizes - */ -static TestTrainData trainTestSplit( - std::vector>& features, - std::vector& labels, - float testSize = 0.2, - bool shuffle = true, - unsigned int randomState = 40) -{ - if (features.size() != labels.size()) { - throw Exception("Features and labels must have same size"); - } - - std::vector indices(features.size()); - std::iota(indices.begin(), indices.end(), 0); - - if (shuffle) { - std::mt19937 rng(randomState); - std::shuffle(indices.begin(), indices.end(), rng); - - std::vector> shuffledFeatures; - std::vector shuffledLabels; - shuffledFeatures.reserve(features.size()); - shuffledLabels.reserve(labels.size()); - - for (size_t idx : indices) { - shuffledFeatures.push_back(features[idx]); - shuffledLabels.push_back(labels[idx]); - } - features = shuffledFeatures; - labels = shuffledLabels; - } - - size_t trainSize = static_cast(features.size() * (1.0f - testSize)); - - std::vector> XTrain( - features.begin(), features.begin() + trainSize); - std::vector> XTest( - features.begin() + trainSize, features.end()); - std::vector yTrain(labels.begin(), labels.begin() + trainSize); - std::vector yTest(labels.begin() + trainSize, labels.end()); - - return { .train = FeatureData(XTrain, yTrain), - .test = FeatureData(XTest, yTest) }; -} - -/** - * Transforms a trained XGBoost model into GBTPredictor. - * - * @param xgBoostDump JSON strings for each tree - * @param numClasses Number of classification classes (2 for binary) - * - * @return GBTPredictorWrapper containing the converted model with ownership - * of all internal structures - * @throws Exception if the XGBoost dump is unexpected size - */ -static GBTPredictorWrapper createGBTModelFromXGBoost( - std::vector& xgBoostDump, - size_t numClasses) -{ - GBTPredictorWrapper pred; - // For binary classification, XGBoost creates 1 forest - // For multiclass, it creates one forest per class - if (numClasses == 1) { - throw Exception("Only 1 class found in XGBoost dump, expected 2"); - } - size_t numForests = numClasses == 2 ? 1 : numClasses; - size_t numRounds = xgBoostDump.size() / numForests; - - /** - * For each boosting round, XGBoost creates a tree per class (unless its - * binary classification then it's just one tree). So we have something like - * this: - * - * r0: tree_label_0, tree_label_1, .... tree_label_n - * r1: tree_label_0, tree_label_1, .... tree_label_n - * ... - * - * So for forest i, we use (jth_round * numForests) + i to get all trees in - * that forest - */ - for (size_t forestIdx = 0; forestIdx < numForests; forestIdx++) { - std::vector trees; - for (size_t round = 0; round < numRounds; round++) { - size_t treeIdx = round * numForests + forestIdx; - if (treeIdx >= xgBoostDump.size()) { - throw Exception("XGBoost Dump unexpected size"); - } - const std::string& treeStr = xgBoostDump[treeIdx]; - xgboost::Json treeJson = xgboost::Json::Load( - xgboost::StringView(treeStr.data(), treeStr.size())); - - std::vector nodes; - int maxNodeId = findMaxNodeId(treeJson); - nodes.resize(maxNodeId + 1); - - parseXGBoostNode(treeJson, nodes); - - pred.core_nodes_.push_back( - std::make_unique>( - std::move(nodes))); - - trees.push_back( - { .numNodes = pred.core_nodes_.back()->size(), - .nodes = pred.core_nodes_.back()->data() }); - } - - pred.core_trees_.push_back( - std::make_unique>( - std::move(trees))); - - if (!pred.core_forests_) { - pred.core_forests_ = - std::make_unique>(); - } - - pred.core_forests_->push_back( - { .numTrees = pred.core_trees_.back()->size(), - .trees = pred.core_trees_.back()->data() }); - } - - pred.core_predictor_ = std::make_unique( - GBTPredictor{ .numForests = pred.core_forests_->size(), - .forests = pred.core_forests_->data() }); - - return pred; -} - -/** - * Trains an XGBoost gradient boosted tree model and converts it to the - * GBTPredictor format for inference. - * - * @param data Train/test split data - * @param num_classes Number of classification classes - * - * @return GBTPredictorWrapper containing the trained model - * - */ -static GBTPredictorWrapper trainXGBoostModel( - TestTrainData& data, - size_t num_classes) -{ - if (data.train.X.empty() || data.train.y.empty()) { - throw Exception("Training data cannot be empty"); - } - - if (data.test.X.empty() || data.test.y.empty()) { - throw Exception("Test data cannot be empty"); - } - - std::vector xgBoostDump; - - // equivalent to n_estimators in python XGBoost - constexpr size_t DEFAULT_XGBOOST_ROUNDS = 30; - - // Create raw handles first, then wrap in unique_ptr - DMatrixHandle trainHandleRaw = nullptr; - DMatrixHandle testHandleRaw = nullptr; - BoosterHandle boosterHandleRaw = nullptr; - - // Convert training and test data to DMatrixHandle needed by XGBoost - safe_xgboost(XGDMatrixCreateFromMat( - data.train.XFlat.data(), - data.train.X.size(), - data.train.X.front().size(), - -1, - &trainHandleRaw)); - DMatrixUniquePtr train(trainHandleRaw); - - safe_xgboost(XGDMatrixSetFloatInfo( - train.get(), "label", data.train.y.data(), data.train.y.size())); - - safe_xgboost(XGDMatrixCreateFromMat( - data.test.XFlat.data(), - data.test.X.size(), - data.test.X.front().size(), - -1, - &testHandleRaw)); - DMatrixUniquePtr test(testHandleRaw); - - // Set model parameters - safe_xgboost(XGDMatrixSetFloatInfo( - test.get(), "label", data.test.y.data(), data.test.y.size())); - - DMatrixHandle trainHandle = train.get(); - safe_xgboost(XGBoosterCreate(&trainHandle, 1, &boosterHandleRaw)); - BoosterUniquePtr booster(boosterHandleRaw); - - BoosterHandle boosterHandle = booster.get(); - safe_xgboost(XGBoosterSetParam(boosterHandle, "booster", "gbtree")); - safe_xgboost(XGBoosterSetParam(boosterHandle, "learning_rate", "0.1")); - // equivalent to n_jobs in python XGBoost - safe_xgboost(XGBoosterSetParam(boosterHandle, "nthread", "1")); - safe_xgboost(XGBoosterSetParam(boosterHandle, "min_child_weight", "0.0")); - safe_xgboost(XGBoosterSetParam(boosterHandle, "subsample", "0.7")); - safe_xgboost(XGBoosterSetParam(boosterHandle, "colsample_bynode", "0.8")); - - // Set the objective function based on whether multiclass or binary - if (num_classes > 2) { - safe_xgboost(XGBoosterSetParam( - boosterHandle, "objective", "multi:softprob")); - safe_xgboost(XGBoosterSetParam( - boosterHandle, - "num_class", - std::to_string(num_classes).c_str())); - } else { - safe_xgboost(XGBoosterSetParam( - boosterHandle, "objective", "binary:logistic")); - // Explicitly set base_score to 0.5 (this just means that data has 50/50 - // chance to be in class 1 or class 2 as a start point). This is to - // avoid auto-computation error, which happens if training data is - // heavily imbalanced and only labeled to one class. - safe_xgboost(XGBoosterSetParam(boosterHandle, "base_score", "0.5")); - } - - const int eval_dmats_size = 2; - DMatrixHandle eval_dmats[eval_dmats_size] = { trainHandle, test.get() }; - Logger::log_c( - VERBOSE1, - "Training XGBoost model with %d boosting rounds\n", - DEFAULT_XGBOOST_ROUNDS); - // Start training the model - for (size_t i = 0; i < DEFAULT_XGBOOST_ROUNDS; ++i) { - // Update the model performance for each iteration - safe_xgboost( - XGBoosterUpdateOneIter(boosterHandle, (int)i, trainHandle)); - - const char* eval_names[eval_dmats_size] = { "train", "test" }; - const char* eval_result = NULL; - safe_xgboost(XGBoosterEvalOneIter( - boosterHandle, - (int)i, - eval_dmats, - eval_names, - eval_dmats_size, - &eval_result)); - Logger::log_c(VERBOSE1, "%s\n", eval_result); - } - - bst_ulong len = 0; - - const char** dump = nullptr; - - safe_xgboost(XGBoosterDumpModelEx( - boosterHandle, - "", // fmap (empty string for no feature map) - 0, // No stats - "json", - &len, - &dump)); - - xgBoostDump = std::vector(dump, dump + len); - - // unique_ptr will automatically free the handles when they go out of scope - return createGBTModelFromXGBoost(xgBoostDump, num_classes); -} - std::shared_ptr trainMLSelectorGraph( const std::vector& inputs, Compressor& compressor, @@ -555,8 +62,40 @@ std::shared_ptr trainMLSelectorGraph( TestTrainData splitData = trainTestSplit( trainingSample.features, trainingSample.numericLabels); + Hyperparams params = defaultXGBoostHyperParams; + if (trainParams.tuneHyperparams) { + training::TrainingContext trainingContext = { + .featureNames = trainingSample.featureNames, + .featurePtrNames = trainingSample.featurePtrNames, + .splitData = splitData, + .successorGraphs = successorGraphs, + .compressor = &compressor, + }; + + TuningConfig config; + std::vector initialPop = + generateInitialTuningPop(defaultParamRanges, config); + + auto [tunedResult, defaultResult] = training::tuneHyperparams( + trainingContext, + mlSelectorInputs, + defaultParamRanges, + initialPop, + config); + + auto tunedScore = config.compressionWeight * tunedResult.size + + (1 - config.compressionWeight) * tunedResult.ctime; + + auto defaultScore = config.compressionWeight * defaultResult.size + + (1 - config.compressionWeight) * defaultResult.ctime; + + if (tunedScore < defaultScore) { + params = tunedResult.params; + } + } + GBTPredictorWrapper gbtPred = - trainXGBoostModel(splitData, successors.nbGraphIDs); + trainXGBoostModel(splitData, successors.nbGraphIDs, params); GBTModel coreModel = { .predictor = gbtPred.core_predictor_.get(), diff --git a/tools/ml_selector/ml_selector_trainer_utils.cpp b/tools/ml_selector/ml_selector_trainer_utils.cpp new file mode 100644 index 000000000..653a9bbcd --- /dev/null +++ b/tools/ml_selector/ml_selector_trainer_utils.cpp @@ -0,0 +1,467 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#include "tools/ml_selector/ml_selector_trainer_utils.h" +#include +#include +#include +#include +#include +#include "openzl/cpp/Exception.hpp" +#include "tools/logger/Logger.h" + +// Suppress warnings for XGBoost headers +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wcast-qual" +#pragma GCC diagnostic ignored "-Wfloat-equal" +#pragma GCC diagnostic ignored "-Wcast-align" +#include +#include +#pragma GCC diagnostic pop +using namespace openzl::tools::logger; +using openzl::Exception; + +namespace openzl::training { + +namespace { +/** + * Macro provided by xgboost cpp to guard all calls + */ +#define safe_xgboost(call) \ + { \ + int err = (call); \ + if (err != 0) { \ + throw std::runtime_error( \ + std::string(__FILE__) + ":" + std::to_string(__LINE__) \ + + ": error in " + #call + ":" + XGBGetLastError()); \ + } \ + } + +} // namespace + +/** + * Extracts a numeric value from an XGBoost JSON node. + * + * Handles both JsonNumber (float) and JsonInteger, converting + * the value to a float. Throws exception if unable to extract value. + * + * @param json The XGBoost JSON node to extract the value from + * @return The numeric value as a float, or -1.0f if not a numeric type + */ +static float extractNumericVal(const xgboost::Json& json) +{ + if (xgboost::IsA(json)) { + return xgboost::get(json); + } else if (xgboost::IsA(json)) { + return static_cast( + xgboost::get(json)); + } + throw Exception( + "Expected JsonNumber. Unable to extract numeric value from XGBoost JSON node"); +} + +/** + * Recursively computes the maximum node ID in an XGBoost decision tree. + * + * Traverses the JSON format for XGBoost model. Used to + * pre-allocate a contiguous vector for storing parsed tree nodes. + * + * @param json The root JSON node of the XGBoost tree + * @return The maximum node ID found in the tree + */ +static int findMaxNodeId(const xgboost::Json& json) +{ + auto& objMap = xgboost::get(json); + int maxId = extractNumericVal(json["nodeid"]); + + if (objMap.find("children") != objMap.end()) { + auto& children = + xgboost::get(json["children"]); + for (const auto& child : children) { + maxId = std::max(maxId, findMaxNodeId(child)); + } + } + return maxId; +} + +/** + * Deleter functor to manage XGBoost BoosterHandle resource. Automatically calls + * XGBoosterFree when unique_pointer goes out of scope + */ +struct BoosterHandleDeleter { + using pointer = BoosterHandle; + void operator()(BoosterHandle handle) const + { + if (handle != nullptr) { + XGBoosterFree(handle); + } + } +}; + +using BoosterUniquePtr = std::unique_ptr; + +/** + * Deleter functor to manage XGBoost DMatrixHandle resource. Automatically calls + * XGDMatrixFree when unique_pointer goes out of scope + */ +struct DMatrixHandleDeleter { + using pointer = DMatrixHandle; + void operator()(DMatrixHandle handle) const + { + if (handle != nullptr) { + XGDMatrixFree(handle); + } + } +}; + +using DMatrixUniquePtr = std::unique_ptr; + +/** + * Recursively parses an XGBoost decision tree node from JSON into a + * GBTPredictor_Node structure. + * + * The parsed node is stored directly into the pre-allocated nodes vector + * at the position corresponding to its nodeid. + * + * @param nodeJson The JSON representation of the tree node + * @param nodes Output vector where parsed nodes are stored by their ID + * + * @throws Exception if nodeid exceeds the bounds of the nodes vector + */ +static void parseXGBoostNode( + const xgboost::Json& nodeJson, + std::vector& nodes) +{ + GBTPredictor_Node node{}; + + auto& objMap = xgboost::get(nodeJson); + + size_t nodeid = extractNumericVal(nodeJson["nodeid"]); + if (objMap.find("leaf") != objMap.end()) { + // Leaf node + float leafValue = extractNumericVal(nodeJson["leaf"]); + node.value = leafValue; + node.featureIdx = -1; + + Logger::log_c(EVERYTHING, "Leaf node: %f", leafValue); + } else { + // Internal node + // Parse feature index from "fN" format + std::string splitFeature = + xgboost::get(nodeJson["split"]); + float threshold = extractNumericVal(nodeJson["split_condition"]); + int yesChild = extractNumericVal(nodeJson["yes"]); + int noChild = extractNumericVal(nodeJson["no"]); + int missingChild = extractNumericVal(nodeJson["missing"]); + + node.featureIdx = + std::stoi(splitFeature.substr(1)); // Skip the 'f' prefix + node.leftChildIdx = yesChild; + node.rightChildIdx = noChild; + node.missingChildIdx = missingChild; + node.value = threshold; + + Logger::log_c(EVERYTHING, "Internal Node %d: ,", nodeid); + + Logger::log_c( + EVERYTHING, "split[%s < %f] ", splitFeature.c_str(), threshold); + + Logger::log_c( + EVERYTHING, + ", Yes node %d, no node %d, missing node %d", + yesChild, + noChild, + missingChild); + + // Recursively process children + if (objMap.find("children") != objMap.end()) { + auto& children = xgboost::get( + nodeJson["children"]); + for (const auto& child : children) { + parseXGBoostNode(child, nodes); + } + } + } + + if (nodeid >= nodes.size()) { + throw Exception("Node id is out of range"); + } + nodes[nodeid] = node; +} + +/** + * Splits feature and label data into training and test sets. Used for + * evaluating model performance during XGBoost boosting rounds. + * + * @param features Feature vectors to split + * @param labels Corresponding labels to split + * @param testSize Fraction of data to reserve for testing + * @param shuffle Whether to shuffle (in place) data before splitting + * @param randomState Seed for shuffling + * + * @return TrainTestData containing XTrain, XTest, yTrain, and yTest splits + * + * @throws Exception if features and labels have different sizes + */ +TestTrainData trainTestSplit( + std::vector>& features, + std::vector& labels, + float testSize, + bool shuffle, + unsigned int randomState) +{ + if (features.size() != labels.size()) { + throw Exception("Features and labels must have same size"); + } + + std::vector indices(features.size()); + std::iota(indices.begin(), indices.end(), 0); + + if (shuffle) { + std::mt19937 rng(randomState); + std::shuffle(indices.begin(), indices.end(), rng); + + std::vector> shuffledFeatures; + std::vector shuffledLabels; + shuffledFeatures.reserve(features.size()); + shuffledLabels.reserve(labels.size()); + + for (size_t idx : indices) { + shuffledFeatures.push_back(features[idx]); + shuffledLabels.push_back(labels[idx]); + } + features = shuffledFeatures; + labels = shuffledLabels; + } + + size_t trainSize = static_cast(features.size() * (1.0f - testSize)); + + std::vector> XTrain( + features.begin(), features.begin() + trainSize); + std::vector> XTest( + features.begin() + trainSize, features.end()); + std::vector yTrain(labels.begin(), labels.begin() + trainSize); + std::vector yTest(labels.begin() + trainSize, labels.end()); + + return { .train = FeatureData(XTrain, yTrain), + .test = FeatureData(XTest, yTest) }; +} + +/** + * Transforms a trained XGBoost model into GBTPredictor. + * + * @param xgBoostDump JSON strings for each tree + * @param numClasses Number of classification classes (2 for binary) + * + * @return GBTPredictorWrapper containing the converted model with ownership + * of all internal structures + * @throws Exception if the XGBoost dump is unexpected size + */ +static GBTPredictorWrapper createGBTModelFromXGBoost( + std::vector& xgBoostDump, + size_t numClasses) +{ + GBTPredictorWrapper pred; + // For binary classification, XGBoost creates 1 forest + // For multiclass, it creates one forest per class + if (numClasses == 1) { + throw Exception("Only 1 class found in XGBoost dump, expected 2"); + } + size_t numForests = numClasses == 2 ? 1 : numClasses; + size_t numRounds = xgBoostDump.size() / numForests; + + /** + * For each boosting round, XGBoost creates a tree per class (unless its + * binary classification then it's just one tree). So we have something like + * this: + * + * r0: tree_label_0, tree_label_1, .... tree_label_n + * r1: tree_label_0, tree_label_1, .... tree_label_n + * ... + * + * So for forest i, we use (jth_round * numForests) + i to get all trees in + * that forest + */ + for (size_t forestIdx = 0; forestIdx < numForests; forestIdx++) { + std::vector trees; + for (size_t round = 0; round < numRounds; round++) { + size_t treeIdx = round * numForests + forestIdx; + if (treeIdx >= xgBoostDump.size()) { + throw Exception("XGBoost Dump unexpected size"); + } + const std::string& treeStr = xgBoostDump[treeIdx]; + xgboost::Json treeJson = xgboost::Json::Load( + xgboost::StringView(treeStr.data(), treeStr.size())); + + std::vector nodes; + int maxNodeId = findMaxNodeId(treeJson); + nodes.resize(maxNodeId + 1); + + parseXGBoostNode(treeJson, nodes); + + pred.core_nodes_.push_back( + std::make_unique>( + std::move(nodes))); + + trees.push_back( + { .numNodes = pred.core_nodes_.back()->size(), + .nodes = pred.core_nodes_.back()->data() }); + } + + pred.core_trees_.push_back( + std::make_unique>( + std::move(trees))); + + if (!pred.core_forests_) { + pred.core_forests_ = + std::make_unique>(); + } + + pred.core_forests_->push_back( + { .numTrees = pred.core_trees_.back()->size(), + .trees = pred.core_trees_.back()->data() }); + } + + pred.core_predictor_ = std::make_unique( + GBTPredictor{ .numForests = pred.core_forests_->size(), + .forests = pred.core_forests_->data() }); + + return pred; +} + +/** + * Trains an XGBoost gradient boosted tree model and converts it to the + * GBTPredictor format for inference. + * + * @param data Train/test split data + * @param num_classes Number of classification classes + * + * @return GBTPredictorWrapper containing the trained model + * + */ +GBTPredictorWrapper trainXGBoostModel( + TestTrainData& data, + size_t num_classes, + const Hyperparams& hyperparams) +{ + if (data.train.X.empty() || data.train.y.empty()) { + throw Exception("Training data cannot be empty"); + } + + if (data.test.X.empty() || data.test.y.empty()) { + throw Exception("Test data cannot be empty"); + } + + std::vector xgBoostDump; + + // equivalent to n_estimators in python XGBoost + constexpr size_t DEFAULT_XGBOOST_ROUNDS = 30; + + // Create raw handles first, then wrap in unique_ptr + DMatrixHandle trainHandleRaw = nullptr; + DMatrixHandle testHandleRaw = nullptr; + BoosterHandle boosterHandleRaw = nullptr; + + // Convert training and test data to DMatrixHandle needed by XGBoost + safe_xgboost(XGDMatrixCreateFromMat( + data.train.XFlat.data(), + data.train.X.size(), + data.train.X.front().size(), + -1, + &trainHandleRaw)); + DMatrixUniquePtr train(trainHandleRaw); + + safe_xgboost(XGDMatrixSetFloatInfo( + train.get(), "label", data.train.y.data(), data.train.y.size())); + + safe_xgboost(XGDMatrixCreateFromMat( + data.test.XFlat.data(), + data.test.X.size(), + data.test.X.front().size(), + -1, + &testHandleRaw)); + DMatrixUniquePtr test(testHandleRaw); + + // Set model parameters + safe_xgboost(XGDMatrixSetFloatInfo( + test.get(), "label", data.test.y.data(), data.test.y.size())); + + DMatrixHandle trainHandle = train.get(); + safe_xgboost(XGBoosterCreate(&trainHandle, 1, &boosterHandleRaw)); + BoosterUniquePtr booster(boosterHandleRaw); + + BoosterHandle boosterHandle = booster.get(); + safe_xgboost(XGBoosterSetParam(boosterHandle, "booster", "gbtree")); + // equivalent to n_jobs in python XGBoost + safe_xgboost(XGBoosterSetParam(boosterHandle, "nthread", "1")); + + for (auto& [param, val] : hyperparams) { + if (param == "num_boost_round") { + continue; + } + safe_xgboost( + XGBoosterSetParam(boosterHandle, param.c_str(), val.c_str())); + } + + // Set the objective function based on whether multiclass or binary + if (num_classes > 2) { + safe_xgboost(XGBoosterSetParam( + boosterHandle, "objective", "multi:softprob")); + safe_xgboost(XGBoosterSetParam( + boosterHandle, + "num_class", + std::to_string(num_classes).c_str())); + } else { + safe_xgboost(XGBoosterSetParam( + boosterHandle, "objective", "binary:logistic")); + // Explicitly set base_score to 0.5 (this just means that data has 50/50 + // chance to be in class 1 or class 2 as a start point). This is to + // avoid auto-computation error, which happens if training data is + // heavily imbalanced and only labeled to one class. + safe_xgboost(XGBoosterSetParam(boosterHandle, "base_score", "0.5")); + } + + const int eval_dmats_size = 2; + DMatrixHandle eval_dmats[eval_dmats_size] = { trainHandle, test.get() }; + size_t num_rounds = hyperparams.find("num_boost_round") != hyperparams.end() + ? std::stoi(hyperparams.at("num_boost_round")) + : DEFAULT_XGBOOST_ROUNDS; + Logger::log_c( + VERBOSE1, + "Training XGBoost model with %d boosting rounds", + num_rounds); + // Start training the model + + for (size_t i = 0; i < num_rounds; ++i) { + // Update the model performance for each iteration + safe_xgboost( + XGBoosterUpdateOneIter(boosterHandle, (int)i, trainHandle)); + + const char* eval_names[eval_dmats_size] = { "train", "test" }; + const char* eval_result = NULL; + safe_xgboost(XGBoosterEvalOneIter( + boosterHandle, + (int)i, + eval_dmats, + eval_names, + eval_dmats_size, + &eval_result)); + Logger::log_c(VERBOSE3, "Evaluation Results %s", eval_result); + } + + bst_ulong len = 0; + + const char** dump = nullptr; + + safe_xgboost(XGBoosterDumpModelEx( + boosterHandle, + "", // fmap (empty string for no feature map) + 0, // No stats + "json", + &len, + &dump)); + + xgBoostDump = std::vector(dump, dump + len); + + // unique_ptr will automatically free the handles when they go out of scope + return createGBTModelFromXGBoost(xgBoostDump, num_classes); +} +} // namespace openzl::training diff --git a/tools/ml_selector/ml_selector_trainer_utils.h b/tools/ml_selector/ml_selector_trainer_utils.h new file mode 100644 index 000000000..f917b443e --- /dev/null +++ b/tools/ml_selector/ml_selector_trainer_utils.h @@ -0,0 +1,104 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#pragma once +#include +#include +#include +#include +#include "src/openzl/compress/selectors/ml/gbt.h" + +namespace openzl::training { +using Hyperparams = std::map; + +/** + * Default XGBoost hyperparameters. + */ +inline const Hyperparams defaultXGBoostHyperParams = { + { "learning_rate", "0.1" }, + { "min_child_weight", "0.0" }, + { "subsample", "0.7" }, + { "colsample_bynode", "0.8" }, +}; + +/** + * Owns the memory for a GBTPredictor. + */ +struct GBTPredictorWrapper { + /// Node arrays for each tree (indexed by tree) + std::vector>> core_nodes_; + /// Tree arrays for each forest (indexed by forest) + std::vector>> core_trees_; + /// All forests in the model + std::unique_ptr> core_forests_; + std::unique_ptr core_predictor_; +}; + +/** + * Contains both the 2D feature matrices and their flattened versions. The 2D + * format is for train test split, while the flattened format is required by + * XGBoost's DMatrix API. + * + * Naming follows scikit-learn conventions: + * - X: Feature matrix + * - y: Label vector + */ +struct FeatureData { + std::vector> X; + std::vector y; + std::vector XFlat; + + FeatureData( + const std::vector>& X_, + const std::vector& y_) + : X{ X_ }, y{ y_ } + { + for (const auto& x : X) { + XFlat.insert(XFlat.end(), x.begin(), x.end()); + } + } +}; + +/** + * Holds the train/test split data for XGBoost model training. + * + */ +struct TestTrainData { + FeatureData train; + FeatureData test; +}; + +/** + * Splits feature and label data into training and test sets. Used for + * evaluating model performance during XGBoost boosting rounds. + * + * @param features Feature vectors to split + * @param labels Corresponding labels to split + * @param testSize Fraction of data to reserve for testing + * @param shuffle Whether to shuffle (in place) data before splitting + * @param randomState Seed for shuffling + * + * @return TrainTestData containing XTrain, XTest, yTrain, and yTest splits + * + * @throws Exception if features and labels have different sizes + */ +TestTrainData trainTestSplit( + std::vector>& features, + std::vector& labels, + float testSize = 0.2, + bool shuffle = true, + unsigned int randomState = 40); + +/** + * Trains an XGBoost gradient boosted tree model and converts it to the + * GBTPredictor format for inference. + * + * @param data Train/test split data + * @param num_classes Number of classification classes + * + * @return GBTPredictorWrapper containing the trained model + */ +GBTPredictorWrapper trainXGBoostModel( + TestTrainData& data, + size_t num_classes, + const Hyperparams& hyperparams = defaultXGBoostHyperParams); + +} // namespace openzl::training diff --git a/tools/ml_selector/ml_selector_tuner.cpp b/tools/ml_selector/ml_selector_tuner.cpp new file mode 100644 index 000000000..065cb90cd --- /dev/null +++ b/tools/ml_selector/ml_selector_tuner.cpp @@ -0,0 +1,398 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#include "tools/ml_selector/ml_selector_tuner.h" +#include +#include +#include +#include +#include +#include +#include "openzl/zl_compressor.h" +#include "tools/logger/Logger.h" +#include "tools/ml_selector/ml_selector_graph.h" +using namespace openzl::tools::logger; + +namespace openzl::training { +namespace { // Anonymous namespace + +struct EvaluationResultComparator { + float cSizeWeight; + + explicit EvaluationResultComparator(float weight = 0.75f) + : cSizeWeight{ weight } + { + } + + bool operator()(const EvaluationResult& a, const EvaluationResult& b) const + { + float scoreA = calculateHyperparamScore(a.size, a.ctime, cSizeWeight); + float scoreB = calculateHyperparamScore(b.size, b.ctime, cSizeWeight); + return scoreA > scoreB; + } +}; + +using SortedPopulation = std::priority_queue< + EvaluationResult, + std::vector, + EvaluationResultComparator>; + +const std::set kIntegerParams = { "max_depth", + "num_boost_round", + "max_leaves" }; + +} // namespace + +static std::pair evaluateWithContext( + TrainingContext& ctx, + const std::vector& inputs, + const Hyperparams& params) +{ + GBTPredictorWrapper gbtPred = trainXGBoostModel( + ctx.splitData, ctx.successorGraphs.size(), params); + + GBTModel coreModel = { + .predictor = gbtPred.core_predictor_.get(), + .featureGenerator = FeatureGen_integer, + .nbSuccessors = ctx.successorGraphs.size(), + .nbFeatures = ctx.featurePtrNames.size(), + .featureLabels = ctx.featurePtrNames.data(), + }; + ZL_MLSelectorConfig config = { .model = ZL_GBT, + .runtimeConfig = &coreModel }; + + // Register new ML selector graph with trained model + auto mlSelectorGraphId = unwrap(ZL_MLSelector_registerGraph( + ctx.compressor->get(), + &config, + ctx.successorGraphs.data(), + ctx.successorGraphs.size())); + + ZL_GraphID staticGraph = ZL_Compressor_registerStaticGraph_fromNode1o( + ctx.compressor->get(), + ZL_NODE_CONVERT_SERIAL_TO_NUM_LE64, + mlSelectorGraphId); + + ctx.compressor->selectStartingGraph(staticGraph); + + CCtx cctx = refCCtxForTraining(*ctx.compressor); + + auto const timerStart = std::chrono::steady_clock::now(); + size_t totalSize = 0; + for (size_t i = 0; i < inputs.size(); i++) { + size_t compressedSize = cctx.compressOne(inputs[i]->front()).size(); + totalSize += compressedSize; + } + std::chrono::duration const timeElapsedMS = + (std::chrono::steady_clock::now() - timerStart); + return std::make_pair(totalSize, timeElapsedMS.count()); +} + +static SortedPopulation evaluatePopulation( + TrainingContext& ctx, + const std::vector& inputs, + const std::vector& currPop, + float cSizeWeight) +{ + EvaluationResultComparator comparator(cSizeWeight); + SortedPopulation compressionSizes(comparator); + + float minScore = std::numeric_limits::max(); + float maxScore = std::numeric_limits::lowest(); + float minSize = 0, minCtime = 0; + float maxSize = 0, maxCtime = 0; + for (size_t i = 0; i < currPop.size(); i++) { + auto [size, ctime] = evaluateWithContext(ctx, inputs, currPop[i]); + compressionSizes.push({ size, ctime, currPop[i] }); + float score = calculateHyperparamScore(size, ctime, cSizeWeight); + + if (score < minScore) { + minScore = score; + minSize = size; + minCtime = ctime; + } + if (score > maxScore) { + maxScore = score; + maxSize = size; + maxCtime = ctime; + } + } + + Logger::log_c( + VERBOSE2, + " Worst result in this generation with compression size %0.2f and ctime %0.2f", + maxSize, + maxCtime); + + Logger::log_c( + VERBOSE2, + " Best result in this generation with compression size %0.2f and ctime %0.2f", + minSize, + minCtime); + + return compressionSizes; +} + +static std::vector crossover( + const std::vector& survivingPop, + size_t childPopSize, + const TuningConfig& config) +{ + std::uniform_int_distribution parentDist( + 0, survivingPop.size() - 1); + std::uniform_int_distribution coinFlip(0, 1); + + std::vector childPop; + for (size_t i = 0; i < childPopSize; i++) { + size_t parent1 = parentDist(config.rng); + size_t parent2; + do { + parent2 = parentDist(config.rng); + } while (parent2 == parent1 && survivingPop.size() > 1); + + Hyperparams child; + for (const auto& [gene, _] : survivingPop[parent1]) { + child[gene] = coinFlip(config.rng) == 0 + ? survivingPop[parent1].at(gene) + : survivingPop[parent2].at(gene); + } + childPop.push_back(child); + } + return childPop; +} + +static void mutate( + std::vector& pop, + const TuningConfig& config, + const std::map& paramRanges) +{ + std::uniform_real_distribution probDist(0.0f, 1.0f); + std::uniform_real_distribution noiseDist(-0.2f, 0.2f); + + for (auto& individual : pop) { + for (auto& [gene, value] : individual) { + if (probDist(config.rng) < config.mutationRate) { + const auto& range = paramRanges.at(gene); + float currentVal = std::stof(value); + float mutatedVal = std::clamp( + currentVal + + noiseDist(config.rng) + * (range.max - range.min), + range.min, + range.max); + + if (kIntegerParams.count(gene)) { + value = std::to_string( + static_cast(std::round(mutatedVal))); + } else { + value = std::to_string(mutatedVal); + } + } + } + } +} + +float calculateHyperparamScore(float size, float ctime, float weight) +{ + return weight * size + (1 - weight) * ctime; +} + +std::vector generateInitialTuningPop( + const std::map& ranges, + const TuningConfig& config) +{ + std::vector population; + population.reserve(config.populationSize); + + if (config.populationSize == 0) { + std::cerr << "Error: Population size must be greater than 0" + << std::endl; + return population; + } + + // For each parameter, create stratified intervals + std::map> stratifiedValues; + + for (const auto& [name, range] : ranges) { + std::vector values; + float stepSize = (range.max - range.min) / config.populationSize; + + float curr = range.min; + for (size_t i = 0; i < config.populationSize; i++) { + std::uniform_real_distribution dist(curr, curr + stepSize); + float value = dist(config.rng); + if (kIntegerParams.count(name)) { + values.push_back((int)value); + } else { + values.push_back(value); + } + curr += stepSize; + } + // Shuffle to randomize order + std::shuffle(values.begin(), values.end(), config.rng); + stratifiedValues[name] = std::move(values); + } + + for (size_t i = 0; i < config.populationSize; i++) { + Hyperparams params; + for (const auto& [name, values] : stratifiedValues) { + if (kIntegerParams.count(name)) { + params[name] = + std::to_string(static_cast(std::round(values[i]))); + } else { + params[name] = std::to_string(values[i]); + } + } + population.push_back(params); + } + return population; +} + +static int getLabelDistribution(const TrainingContext& ctx) +{ + // Print label distribution across classes + std::map labelCounts; + for (float label : ctx.splitData.train.y) { + labelCounts[static_cast(label)]++; + } + for (float label : ctx.splitData.test.y) { + labelCounts[static_cast(label)]++; + } + + size_t totalSamples = + ctx.splitData.train.y.size() + ctx.splitData.test.y.size(); + Logger::log_c( + VERBOSE2, "Label distribution across %zu samples:", totalSamples); + for (const auto& [classIdx, count] : labelCounts) { + float percentage = 100.0f * count / totalSamples; + Logger::log_c( + VERBOSE2, + " Class %d: %zu samples (%.1f%%)", + classIdx, + count, + percentage); + } + return labelCounts.size(); +} + +static void printTuningResults( + EvaluationResult bestPop, + const std::vector& inputs, + float defaultSize, + float defaultCtime) +{ + Logger::log(VERBOSE1, "Best Hyperparameter"); + for (const auto& [name, value] : bestPop.params) { + Logger::log_c(VERBOSE1, " %s: %s", name.c_str(), value.c_str()); + } + size_t uncompressedSize = 0; + for (auto& input : inputs) { + uncompressedSize += + (*input).front() + .contentSize(); // Gets first Input's contentSize + } + + Logger::log_c( + VERBOSE2, + "Compressed size changed from %0.2f to %0.2f", + defaultSize, + bestPop.size); + + Logger::log_c( + VERBOSE2, + "Compressed ratio changed from %0.2f to %0.2f", + uncompressedSize / defaultSize, + uncompressedSize / bestPop.size); + + Logger::log_c( + VERBOSE2, + "Compression time changed from %0.2f to %0.2f", + defaultCtime, + bestPop.ctime); +} + +std::pair tuneHyperparams( + TrainingContext& ctx, + const std::vector& inputs, + const std::map& paramRanges, + const std::vector& initialPop, + const TuningConfig& config) +{ + if (getLabelDistribution(ctx) < 2) { + Logger::log( + VERBOSE1, + "Error: Not enough classes for tuning. Need at least 2 classes."); + return std::make_pair( + EvaluationResult{ 0, 0, defaultXGBoostHyperParams }, + EvaluationResult{ 0, 0, defaultXGBoostHyperParams }); + } + + float cSizeWeight = config.compressionWeight; + size_t iterationsLeft = config.maxIterations; + size_t convergeVal = config.convergenceThreshold; + float prevBestScore = -1; + + size_t survivingPopSize = + static_cast(config.populationSize * config.survivalRate); + std::vector survivingPop; + std::vector currentPop = initialPop; + + while (iterationsLeft > 0 && convergeVal > 0) { + // Evaluate current population + Logger::log_c( + VERBOSE2, + "Evaluating current population %d", + config.maxIterations - iterationsLeft); + + auto compressionSizes = + evaluatePopulation(ctx, inputs, currentPop, cSizeWeight); + + // Check for convergence + const auto& best = compressionSizes.top(); + float score = + calculateHyperparamScore(best.size, best.ctime, cSizeWeight); + if (prevBestScore < 0) { + prevBestScore = score; + } else { + if (std::abs(score - prevBestScore) < 0.0005f * prevBestScore) { + // If there is no improvement then decrease the convergence + convergeVal--; + prevBestScore = score; + } else { + convergeVal = config.convergenceThreshold; + } + } + + // Select parents for next generation + survivingPop.clear(); + for (size_t i = 0; i < survivingPopSize; i++) { + survivingPop.push_back(compressionSizes.top().params); + compressionSizes.pop(); + } + + // Crossover + auto childPop = crossover( + survivingPop, config.populationSize - survivingPopSize, config); + + survivingPop.insert( + survivingPop.end(), childPop.begin(), childPop.end()); + + // Mutation + mutate(survivingPop, config, paramRanges); + + iterationsLeft--; + currentPop = survivingPop; + } + Logger::log(VERBOSE2, "Stopping tuning... Evaluating final population"); + auto finalPop = evaluatePopulation(ctx, inputs, survivingPop, cSizeWeight); + auto [defaultSize, defaultCtime] = + evaluateWithContext(ctx, inputs, defaultXGBoostHyperParams); + auto bestPop = finalPop.top(); + + printTuningResults(bestPop, inputs, defaultSize, defaultCtime); + + return std::make_pair( + bestPop, + EvaluationResult{ + defaultSize, defaultCtime, defaultXGBoostHyperParams }); +} + +} // namespace openzl::training diff --git a/tools/ml_selector/ml_selector_tuner.h b/tools/ml_selector/ml_selector_tuner.h new file mode 100644 index 000000000..11fbf7561 --- /dev/null +++ b/tools/ml_selector/ml_selector_tuner.h @@ -0,0 +1,122 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +#pragma once +#include +#include +#include +#include +#include "openzl/cpp/Compressor.hpp" +#include "tools/ml_selector/ml_selector_trainer_utils.h" +#include "tools/training/utils/utils.h" + +namespace openzl::training { + +/** + * Defines the search range for a single hyperparameter. + */ +struct ParamRange { + float min; + float max; +}; + +inline const std::map defaultParamRanges = { + { "learning_rate", { 0.001f, 1.0f } }, + { "min_child_weight", { 0.0f, 30.0f } }, + { "subsample", { 0.1f, 1.0f } }, + { "colsample_bynode", { 0.1f, 1.0f } }, +}; + +/** + * Defines the training context for ML selector training. + * + * @param featureNames names of each feature for the ML model + * @param splitData train/test split data + * @param successorGraphs list of successor graphs to train on + * @param compressor base compressor to train + */ +struct TrainingContext { + std::vector featureNames; + std::vector featurePtrNames; + TestTrainData splitData; + std::vector successorGraphs; + Compressor* compressor; +}; + +/** + * Tuning config for ML selector hyperparameter tuning. + * + * @param populationSize number of hyperparameter sets per iteration + * @param survivalRate fraction of population to survive to next iteration + * @param maxIterations max number of iterations to run + * @param convergenceThreshold num of iterations without improvement before + * stopping + * @param mutationRate probability of random modification to hyperparameter + * @param compressionWeight weight of compression size versus compression time + * @param rng random number generator for stochastic operations + */ +struct TuningConfig { + size_t populationSize = 20; + float survivalRate = 0.25f; + size_t maxIterations = 10; + size_t convergenceThreshold = 2; + float mutationRate = 0.25f; + float compressionWeight = 1.0f; + mutable std::mt19937 rng{ std::random_device{}() }; +}; + +/** + * Compression size and time for given set of hyperparameters. + */ +struct EvaluationResult { + float size; + float ctime; + Hyperparams params; +}; + +/** + * Generates an initial population of hyperparameters using stratified sampling. + * + * Each parameter range is divided into `populationSize` equal segments, and + * each population member samples one random value from each segment. This + * ensures the initial population is evenly distributed across the entire + * parameter space. + */ +std::vector generateInitialTuningPop( + const std::map& paramRanges, + const TuningConfig& config); + +/** + * Calculates a weighted fitness score for a set of hyperparameters. + * @param size Compressed size + * @param ctime Compression time + * @param weight Weight for compression size vs time (0.0-1.0). + * Higher values prioritize size over speed. + * @returns Weighted fitness score (lower is better) + */ +float calculateHyperparamScore(float size, float ctime, float weight); + +/** + * Genetic algo hyperparameter tuning for ML selector. + * + * Algorithm: + * 1. Start with random population of hyperparams + * 2. For each iteration, evaluate each hyperparam through compression sizes + * 3. Select the top N hyperparams that survive to next iteration + * 4. Crossover and create new hyperparams by combining hyperparams from top N + * 5. Mutate hyperparams by randomly changing a hyperparam value + * 6. Repeat until convergence or max iterations reached + * + * @param ctx Training context + * @param inputs Inputs for training and evaluating ML selector + * @param paramRanges Map of param name to {min, max} range to search + * @param initialPop Initial population of hyperparams to start with + * @param config Tuning configuration + * @returns Best hyperparameters found based on compression size + */ +std::pair tuneHyperparams( + TrainingContext& ctx, + const std::vector& inputs, + const std::map& paramRanges, + const std::vector& initialPop, + const TuningConfig& config); + +} // namespace openzl::training diff --git a/tools/ml_selector/tests/BUCK b/tools/ml_selector/tests/BUCK index 8e407f7e1..fc7ff9bde 100644 --- a/tools/ml_selector/tests/BUCK +++ b/tools/ml_selector/tests/BUCK @@ -11,6 +11,8 @@ zs_unittest( "..:ml_features", "..:ml_selector_graph", "..:ml_selector_trainer", + "..:ml_selector_trainer_utils", + "..:ml_selector_tuner", "../../../tests:ml_selector_utils", "../../../tests:utils", "../../../tests/datagen:datagen", diff --git a/tools/ml_selector/tests/test_mlSelectorTrainer.cpp b/tools/ml_selector/tests/test_mlSelectorTrainer.cpp index c8efa72bd..b8a259c32 100644 --- a/tools/ml_selector/tests/test_mlSelectorTrainer.cpp +++ b/tools/ml_selector/tests/test_mlSelectorTrainer.cpp @@ -9,6 +9,7 @@ #include "tools/ml_selector/ml_features.h" #include "tools/ml_selector/ml_selector_graph.h" #include "tools/ml_selector/ml_selector_trainer.h" +#include "tools/ml_selector/ml_selector_tuner.h" #include "tools/training/train.h" #include "tools/training/train_params.h" #include "tools/training/utils/utils.h" @@ -229,6 +230,19 @@ class TestMLSelectorTrainer : public testing::Test { return trainData; } + std::pair compressAndTime( + Compressor& compressor, + const std::vector& input) + { + auto compressBound = ZL_compressBound(input.size() * sizeof(uint64_t)); + std::string cBuffer = std::string(compressBound, '\0'); + auto start = std::chrono::steady_clock::now(); + auto compressedSize = compress(compressor, cBuffer, input); + auto end = std::chrono::steady_clock::now(); + std::chrono::duration defaultElapsed = end - start; + + return std::make_pair(compressedSize, defaultElapsed.count()); + } size_t compress( Compressor& compressor, std::string& dst, @@ -509,5 +523,42 @@ TEST_F(TestMLSelectorTrainer, TestAmbiguousData) testSelection(testData, compressor_, mlCompressor); } +TEST_F(TestMLSelectorTrainer, TestHyperparamTuner) +{ + multiInputs_.clear(); + generateTrainData(20); + EXPECT_EQ(multiInputs_.size(), 60); + multiSuccessors_ = registerSuccessors(compressor_, true); + setUpCompressor(trainedCompressor_, true); + auto serializedCompressor = openzl::training::trainMLSelectorGraph( + multiInputs_, trainedCompressor_, trainParams_); + + // Set up a fresh compressor for training with hyperparam tuning + Compressor tunedTrainedCompressor; + setUpCompressor(tunedTrainedCompressor, true); + trainParams_.tuneHyperparams = true; + auto tunedSerializedCompressor = openzl::training::trainMLSelectorGraph( + multiInputs_, tunedTrainedCompressor, trainParams_); + + // Deserialized compressor with default hyperparams and compress test data + Compressor mlCompressor = deserializeCompressor(serializedCompressor); + auto [defaultSize, defaultTime] = + compressAndTime(mlCompressor, testData_[0]); + + // Deserialized compressor with tuned hyperparams and compress test data + Compressor tunedMlCompressor = + deserializeCompressor(tunedSerializedCompressor); + auto [tunedSize, tunedTime] = + compressAndTime(tunedMlCompressor, testData_[0]); + + training::TuningConfig config; + auto tunedScore = config.compressionWeight * tunedSize + + (1 - config.compressionWeight) * tunedTime; + + auto defaultScore = config.compressionWeight * defaultSize + + (1 - config.compressionWeight) * defaultTime; + EXPECT_LE(tunedScore, defaultScore); +} + } // namespace } // namespace openzl::tests diff --git a/tools/training/train_params.h b/tools/training/train_params.h index a544215c0..28bdcfd0b 100644 --- a/tools/training/train_params.h +++ b/tools/training/train_params.h @@ -29,6 +29,7 @@ struct TrainParams { poly::optional maxFileSizeMb; poly::optional maxTotalSizeMb; bool paretoFrontier{ false }; + bool tuneHyperparams{ false }; }; } // namespace openzl::training