バージョン: v2509

自動プルーニング

ZenithTuneは、チューニング中の試行を自動的に枝刈りする自動プルーナーを提供します。この機能により、実行に時間がかかりすぎたり、GPUメモリを過剰に消費するような試行を検出し、自動的に枝刈りすることができます。

自動プルーナーは、コマンド実行中に特定の条件をモニタリングし、設定した閾値を超えた（もしくは下回った）場合に試行を自動的に中断します。メモリ不足や過剰に時間の要する実行を回避してリソースを効率利用し、有望でない試行を早期に中断してチューニング時間を短縮できます。

プルーナーの種類

TimeoutPruner

指定した時間を超過した試行を枝刈りします。

from zenith_tune.auto_pruners import TimeoutPruner
from zenith_tune.tuners import CommandOutputTuner

# 300秒でタイムアウトするプルーナー
timeout_pruner = TimeoutPruner(timeout_seconds=300.0)

tuner = CommandOutputTuner(
    auto_pruners=[timeout_pruner]
)

AIBooster連携プルーナー

AIBoosterサーバーと連携してGPUメトリクスを監視するプルーナーです。

AIBoosterGPUUtilizationPruner

GPU利用率が閾値を下回った場合に試行を枝刈りします。

from zenith_tune.auto_pruners import AIBoosterGPUUtilizationPruner
from zenith_tune.tuners import CommandOutputTuner

# GPU利用率が5%を下回った場合にプルーニング
gpu_util_pruner = AIBoosterGPUUtilizationPruner(
    aibooster_server_address="http://localhost:16697",
    threshold=5.0,  # 5% GPU利用率
)

tuner = CommandOutputTuner(
    auto_pruners=[gpu_util_pruner]
)

AIBoosterGPUMemoryUsedPruner

GPU メモリ使用量が閾値を上回った場合に試行を枝刈りします。

from zenith_tune.auto_pruners import AIBoosterGPUMemoryUsedPruner
from zenith_tune.tuners import CommandOutputTuner
import socket

current_hostname = socket.gethostname()

# GPU メモリ使用量が12800MB（12.8GB）を上回った場合にプルーニング
gpu_memory_pruner = AIBoosterGPUMemoryUsedPruner(
    aibooster_server_address="http://localhost:16697",
    threshold=12800.0,  # 12800 MB
    agent_gpu_filter={
        current_hostname: [0],  # ホストの GPU 0 を監視
    },
)

tuner = CommandOutputTuner(
    auto_pruners=[gpu_memory_pruner]
)

AIBoosterDCGMMetricsPruner（カスタムプルーナー）

上記の専用プルーナーに加えて、AIBoosterDCGMMetricsPrunerを直接使用することで、特定のDCGMメトリクスや条件をカスタマイズできます。

from zenith_tune.auto_pruners import AIBoosterDCGMMetricsPruner
from zenith_tune.tuners import CommandOutputTuner

# カスタム条件: GPU電力が200W以上の場合にプルーニング
power_pruner = AIBoosterDCGMMetricsPruner(
    aibooster_server_address="http://localhost:16697",
    metric_name="DCGM_FI_DEV_POWER_USAGE",  # 電力使用量メトリクス
    threshold=200.0,  # 200W
    prune_when="above",  # 閾値を上回った場合
    reduction="max",  # 最大値で判定
    check_interval=10.0,
    warmup_duration=60.0,
)

tuner = CommandOutputTuner(
    auto_pruners=[power_pruner]
)

プルーニング判定の調整

AIBooster連携プルーナーは、実行開始からwarmup_duration（デフォルト60秒）の間は監視をスキップし、その後check_interval（デフォルト10秒）ごとに判定を行います。各判定では、直近のcheck_interval区間のメトリクスから統計値（平均、最小、最大のいずれか）を計算し、設定した閾値条件と比較してプルーニングの要否を判断します。

例えば、デフォルト設定のAIBoosterGPUUtilizationPrunerでは、実行開始60秒後から10秒ごとに「直近10秒間の平均GPU利用率」をチェックし、しきい値を下回るとプルーニングされます。しかし、深層学習のワークロードではデータローディング、評価フェーズ、チェックポイント保存などで10秒以上CPU処理が続くことが一般的です。

ワークロードの特性に合わせて、これらのパラメータを適切に調整してください。データ準備に時間がかかる場合はwarmup_durationを長く、評価が頻繁な場合はcheck_intervalを長く設定することで、意図しないプルーニングを防ぐことができます。

使用例

使用例は以下のファイルを参照してください：

examples/pruners/timeout_pruner_example.py
examples/pruners/gpu_memory_pruner_example.py

プルーナーの種類​

TimeoutPruner​

AIBooster連携プルーナー​

AIBoosterGPUUtilizationPruner​

AIBoosterGPUMemoryUsedPruner​

AIBoosterDCGMMetricsPruner（カスタムプルーナー）​

プルーニング判定の調整​

使用例​