メインコンテンツまでスキップ
バージョン: v2506

高度なチューニング

ZenithTuneを用いた高度なチューニング方法について説明します。

チューニングスクリプトの並列実行​

ZenithTuneは、分散学習など複数プロセスを使用したアプリケーションのチューニングを積極的にサポートしています。 計算機システムやジョブスケジューラに依存して分散実行方法に制約があるため、複数プロセスアプリケーションに対するZenithTuneを用いたチューニングの実行では以下の方法を順に検討してください。

方法1. チューニングスクリプト：1プロセス、チューニング対象：マルチプロセス​

最もシンプルな例です。 通常通り1プロセスで動作するチューニングスクリプトを記述し、コマンド生成関数でマルチプロセスを立ち上げるコマンドを定義します。

単一のエントリーポイントから複数ノードにプロセスを立ち上げられる場合はこの方法が適しています。

def command_generator(trial, **kwargs):
    num_workers = trial.suggest_int("num_workers", low=1, high=10)
    command = f"mpirun --np 8 train.py --num-workers {num_workers}"
    # command = f"torchrun --nproc-per-node 8 train.py --num-workers {num_workers}"
    return command

tuner = CommandOutputTuner()
tuner.optimize(command_generator, value_extractor, n_trials=10)

チューニング実行方法

python optimize.py

方法2. チューニングスクリプト：マルチプロセス、チューニング対象：1プロセス​

ジョブスケジューラやクラスタシステムによってはシステム自体がマルチプロセス実行のエントリーポイントを担う場合があり、ユーザ自身が単一エントリーポイントから複数ノードにプロセスを立ち上げることが煩雑な場合があります。 このような場合には、チューニングスクリプト自体をマルチプロセス実行し、シングルプロセスのコマンドを生成することも可能です。 Tunerにはマルチプロセス実行のための排他制御や通信設定が実装されているため、ユーザは通常の分散実行エントリーポイント（mpirunやtorchrun）を使用してチューニングすることができます。

def command_generator(trial, **kwargs):
    num_workers = trial.suggest_int("num_workers", low=1, high=10)
    command = f"python train.py --num-workers {num_workers}"
    return command

tuner = CommandOutputTuner()
tuner.optimize(command_generator, value_extractor, n_trials=10)

チューニング実行方法

$ mpirun --np 8 python optimize.py
# $ torchrun --nproc-per-node 8 python optimize.py

方法3. チューニングスクリプト：各ノードに1プロセス、チューニング対象：マルチプロセス​

Kubernetesなどのコンテナ型の分散環境などで、各計算ノードごとにエントリーポイントを記述する場合があります。 この場合は、各計算ノードごとに1プロセスのチューニングスクリプトを実行、コマンド生成は各計算ノードの並列度（例えばGPU数）に従ったプロセス数を立ち上げるよう記述することも可能です。

def command_generator(trial, **kwargs):
    num_workers = trial.suggest_int("num_workers", low=1, high=10)
    command = f"torchrun --nproc-per-node 8 train.py --num-workers {num_workers}"
    return command

tuner = CommandOutputTuner()
tuner.optimize(command_generator, value_extractor, n_trials=10)

チューニング実行方法

torchrun --nproc-per-node 1 python optimize.py

CPU Affinityチューニング​

分散学習において学習プロセスとデータローダプロセスの処理が競合し、全体のスループットが低下してしまう事があります。 この問題に対してデータローダプロセスの使用可能なCPUコアを適切に制限することで、学習プロセスの処理性能が安定し全体のスループットを向上できる場合があります。 一方で、データローダプロセスの使用可能なCPUコアを制限しすぎると、学習本体へのデータ供給処理が遅れてむしろ低速化してしまいます。 したがって、データローダプロセスの使用可能なCPUコア数をチューニングすることで最大のスループットを得られる可能性があります。

CPU Affinityチューニング

そのために、ZenithTuneはデータローダプロセスのCPU Affinityをチューニングするためのコンポーネントを提供しています。 チューニング手順は以下の通りです。

  1. zenith_tune.tuning_component.dataloader_affinity.worker_affinity_init_fnをimportする
  2. Torch準拠のデータローダのworker_init_fnにworker_affinity_init_fnを設定する
  3. worker_affinity_init_fnの引数available_cpusをチューニングする

worker_affinity_init_fnは自身のデータローダプロセスに対してavailable_cpusの値に応じたCPU Affinityを設定します。 available_cpusの最適値を求めることで、最高性能を得ることができます。

また、本機能はMMEngineレジストリをサポートしています。 MMEngine学習スクリプト上でimport zenith_tuneを記述することで、zenith_tune.worker_affinity_init_fnとしてモジュールを参照できます。 この統合により、MMengineの--cfg-optionsを用いてデータローダのworker_init_fnを強制的に上書きすることができます。

python train.py --cfg-options train_dataloader.worker_init_fn.type=zenith_tune.worker_affinity_init_fn \
    train_dataloader.worker_init_fn.available_cpus={available_cpus}

学習データサンプル数の制限​

チューニングを実行するにあたり1回の試行はなるべく高精度かつコンパクトに実行すべきです。 通常、各種フレームワークの設定方法に従ってエポック数やイテレーション数を適切な範囲に設定しますが、そのような調節機能がサポートされていない場合に、zenith_tune.tuning_component.LimitedSamplerを使用することで学習に使用するデータサンプル数を強制的に制限します。 LimitedSamplerのlimited_samplesに制限するデータサンプル数を与え、Samplerに指定してください。