カスタマイズ

ノーコードチューニングの動作をカスタマイズする方法について説明します。 独自のプリセット・評価器を作成することで、任意のフレームワークやメトリクスに対応できます。

カスタムプリセットの作成

独自のフレームワーク向けプリセットを作成するには、TuningPresetを継承し、レジストリに登録します。

from aibooster.intelligence.zenith_tune.presets.base import (
    IntParameter,
    FloatParameter,
    CategoricalParameter,
    BoolParameter,
    TuningPreset,
)
from aibooster.intelligence.zenith_tune.presets.registry import PresetRegistry
from aibooster.intelligence.zenith_tune.command import CommandBuilder
from aibooster.intelligence.zenith_tune.strategies.random import RandomStrategy
from aibooster.intelligence.zenith_tune.evaluators.regex import RegexEvaluator


@PresetRegistry.register("my_framework")
class MyFrameworkPreset(TuningPreset):
    def get_search_space(self):
        return {
            "batch_size": IntParameter(low=8, high=256, step=8, default=32),
            "learning_rate": FloatParameter(low=1e-5, high=1e-2, log=True),
            "use_amp": BoolParameter(default=True),
            "optimizer": CategoricalParameter(choices=["adam", "sgd", "adamw"]),
        }

    def apply_parameters(self, command, params):
        cmd = CommandBuilder(command)
        cmd.update(f"--batch-size {params['batch_size']}")
        cmd.update(f"--lr {params['learning_rate']}")
        cmd.update(f"--optimizer {params['optimizer']}")
        if params["use_amp"]:
            cmd.append("--fp16")
        else:
            cmd.remove("--fp16")
        return cmd.get_command()

    def get_recommended_strategy(self):
        return RandomStrategy()

    def get_recommended_evaluator(self):
        return RegexEvaluator(
            pattern=r"loss:\s*([\d.]+)",
            selector="last",
        )

TuningPresetのメソッド

以下のメソッドを実装する必要があります:

メソッド	説明
get_search_space()	探索するパラメータとその範囲を定義する
apply_parameters(command, params)	探索されたパラメータをコマンドに適用し、変更後のコマンド文字列を返す
get_recommended_strategy()	デフォルト戦略を返す
get_recommended_evaluator()	デフォルトの評価器を返す

以下はオプションでオーバーライドできます:

メソッド	説明
prune(params)	無効なパラメータ組み合わせを事前にスキップする。Trueで試行をスキップ

--argsとコンストラクタの対応

--args key1=val1,key2=val2で渡された値は、プリセットの__init__キーワード引数にマッピングされます。 コンストラクタで必須引数を定義することで、CLIから--args経由でユーザーに値を要求できます。

@PresetRegistry.register("my_framework")
class MyFrameworkPreset(TuningPreset):
    def __init__(self, *, n_gpus: int, lr_min: float = 1e-5, **kwargs):
        self._n_gpus = int(n_gpus)
        self._lr_min = float(lr_min)

zenithtune optimize --preset my_framework --args n_gpus=8,lr_min=1e-4 -- bash train.sh

パラメータ型

型	説明	主なフィールド
IntParameter	整数	low, high, step, default
FloatParameter	浮動小数点	low, high, log（対数スケール）, default
CategoricalParameter	カテゴリカル	choices, default
BoolParameter	ブール	default

low・high・choices以外のフィールドはすべて省略可能です。

CommandBuilder

apply_parametersでコマンドにパラメータを適用する際に使用するユーティリティです。

メソッド	動作	例
update(option)	既存オプションを上書き、なければ末尾に追加	cmd.update("--lr 0.01")
append(option)	既存オプションに値を追加（重複許可）	cmd.append("--fp16")
remove(option)	オプションを削除	cmd.remove("--fp16")
get_command()	変更後のコマンド文字列を取得	cmd.get_command()

normalize=Trueで初期化すると、ハイフンとアンダースコアを同一視してマッチングします（例: --batch-sizeと--batch_size）。

カスタムStrategyの作成

独自の探索戦略を作成するには、TuningStrategyを継承し、StrategyRegistryに登録します。

from aibooster.intelligence.zenith_tune.strategies.base import TuningStrategy
from aibooster.intelligence.zenith_tune.strategies.registry import StrategyRegistry


@StrategyRegistry.register("my_strategy")
class MyStrategy(TuningStrategy):
    def optimize(self, eval_fn, search_space, direction):
        while True:
            params = ...  # search_spaceからパラメータを選択
            eval_fn(params)  # 試行を実行（予算超過時にTrialExhaustedが送出される）

optimizeメソッド内でeval_fn(params)を繰り返し呼ぶことで探索が進行します。試行予算（--n-trials）を超えるとeval_fnからTrialExhausted例外が送出され、探索が終了します。

段階的BlackBox探索の作成

探索空間を2段階に分けて探索する戦略を作成するには、StagedBlackboxStrategyを継承します。 Stage 1で粗い構造パラメータを探索し、その最良結果を固定してStage 2で残りのパラメータを最適化する、という2段階探索のフレームワークが提供されます。各ステージではOptuna TPEが使用されます。

megatron-staged-blackboxはこのクラスの具体的な実装例です。

from aibooster.intelligence.zenith_tune.strategies.staged_blackbox import (
    StagedBlackboxStrategy,
)
from aibooster.intelligence.zenith_tune.strategies.registry import StrategyRegistry


@StrategyRegistry.register("my_staged")
class MyStagedStrategy(StagedBlackboxStrategy):
    def __init__(self, *, stage1_trials: int = 30, **kwargs):
        super().__init__(stage1_trials=stage1_trials, **kwargs)

    def _validate_search_space(self, search_space):
        # search_spaceに必要なパラメータが含まれているか検証
        ...

    def _enumerate_candidates(self, search_space):
        # Stage 1で探索する候補をエンコードした文字列のリストを返す
        return ["candidate_1", "candidate_2", ...]

    def _build_stage1_params(self, candidate):
        # 候補文字列をデコードし、Stage 1のパラメータ辞書を返す
        # Stage 2パラメータはベースライン値に固定
        return {"param_a": ..., "param_b": ..., "param_c": default, "param_d": default}

    def _build_stage2_params(self, trial, search_space, stage1_best):
        # Stage 1の最良パラメータを固定し、Stage 2パラメータをOptunaのtrialから提案
        params = dict(stage1_best)
        params["param_c"] = trial.suggest_categorical("param_c", ...)
        params["param_d"] = trial.suggest_int("param_d", ...)
        return params

    def _log_stage1_start(self, n_candidates):
        ...

    def _log_stage1_best(self, best_params, value):
        ...

    def _log_stage2_start(self, stage1_best):
        ...

StagedBlackboxStrategyのサブクラスで実装が必要なメソッド:

メソッド	説明
_validate_search_space(search_space)	探索空間が両ステージの要件を満たすか検証する
_enumerate_candidates(search_space)	Stage 1の候補をエンコードした文字列のリストを返す
_build_stage1_params(candidate)	候補文字列をデコードしてStage 1のパラメータ辞書を返す
_build_stage2_params(trial, search_space, stage1_best)	Stage 2のパラメータ辞書を構築する
_log_stage1_start(n_candidates)	Stage 1開始時のログ出力
_log_stage1_best(best_params, value)	Stage 1最良結果のログ出力
_log_stage2_start(stage1_best)	Stage 2開始時のログ出力

Preset専用Strategyとして使用する場合は、Presetのget_recommended_strategyから返します。--argsで渡されたパラメータは_coerce_kwargsによりStrategyのコンストラクタにもマッピングされます。

カスタムEvaluatorの作成

RegexEvaluatorを使う

標準出力から正規表現で値を抽出する場合はRegexEvaluatorが利用できます。プリセットのget_recommended_evaluator内で返します。

from aibooster.intelligence.zenith_tune.evaluators.regex import RegexEvaluator
from aibooster.intelligence.zenith_tune.evaluators.base import Direction

# lossを最小化（デフォルト）
RegexEvaluator(pattern=r"loss:\s*([\d.]+)", selector="last")

# スループットを最大化
RegexEvaluator(
    pattern=r"throughput:\s*([\d.]+)",
    selector="mean",
    direction=Direction.MAXIMIZE,
)

selector: 正規表現が複数マッチしたときの集約方法

selector	説明
first	最初のマッチ
last	最後のマッチ（デフォルト）
min / max	最小値 / 最大値
mean / median	平均値 / 中央値

TuningEvaluatorを継承する

より柔軟な評価ロジックが必要な場合はTuningEvaluatorを継承します。

from aibooster.intelligence.zenith_tune.evaluators.base import Direction, TuningEvaluator

class MyEvaluator(TuningEvaluator):
    direction = Direction.MAXIMIZE  # または Direction.MINIMIZE

    def evaluate(self, stdout: str, metadata: dict) -> float:
        # stdout から値を抽出するロジックを実装
        # metadata には trial_id, command, parameters, duration が含まれる
        ...