Tune xgboost hyperparams

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

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"),
+    ],
+)

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
 )

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 -- <input_path> [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) |

tools/ml_selector/ml_features.cpp

@@ -2,6 +2,7 @@
 #include "tools/ml_selector/ml_features.h"
 #include <chrono>
 #include <cstdio>
+#include <queue>
 #include <string>
 #include <vector>
 #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<TargetsMap>& targets) {
+        std::vector<float> result;
+        for (size_t i = 0; i < targets.size(); i++) {
+            std::priority_queue<
+                    std::pair<float, size_t>,
+                    std::vector<std::pair<float, size_t>>,
+                    std::greater<std::pair<float, size_t>>>
+                    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<float> minSizeChoiceFunc(std::vector<TargetsMap>& targets)
 {
     std::vector<float> result;

tools/ml_selector/ml_features.h

@@ -1,5 +1,6 @@
 // Copyright (c) Meta Platforms, Inc. and affiliates.
 #pragma once
+#include <functional>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -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<float> (*)(std::vector<TargetsMap>& targets);
+using ChoiceFunction =
+        std::function<std::vector<float>(std::vector<TargetsMap>&)>;
 
 /**
  * Container for processed ML training data with features and labels.
@@ -38,6 +40,12 @@ struct ProcessedMLTrainingSamples {
     std::vector<const char*> 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".