MosaicML Composer

ニューラルネットワークをトレーニングするための最先端のアルゴリズム

Composer は、ニューラルネットワークをより高性能に、より高速に、より低コストでトレーニングするためのライブラリです。ニューラルネットワークのトレーニングを高速化し、汎化性能を向上させる最先端の method を多数備えているほか、さまざまな機能拡張を簡単に 組み合わせる ためのオプションの Trainer API も含まれています。

W&B は、ML Experiments をログするための軽量なラッパーを提供しています。ただし、この 2 つを自分で組み合わせる必要はありません。W&B は WandBLogger を通じて Composer ライブラリに直接組み込まれています。

W&B へログし始める

from composer import Trainer
from composer.loggers import WandBLogger

trainer = Trainer(..., logger=WandBLogger())

Composer の WandBLogger を使用する

Composer ライブラリでは、Trainer で WandBLogger クラスを使用して、メトリクスを W&B にログします。ロガーをインスタンス化して Trainer に渡すだけなので、簡単に使用できます。

wandb_logger = WandBLogger(project="gpt-5", log_artifacts=True)
trainer = Trainer(logger=wandb_logger)

ロガーの引数

以下に WandbLogger のパラメーターを示します。完全な一覧と説明については、Composer documentationを参照してください。

Parameter	Description
project	W&B プロジェクト名 (str、optional)
group	W&B グループ名 (str、optional)
name	W&B run 名。指定しない場合は State.run_name が使用されます (str、optional)
entity	W&B entity 名。たとえば、ユーザー名や W&B Team 名です (str、optional)
tags	W&B tags (List[str]、optional)
log_artifacts	checkpoint を wandb にログするかどうか。デフォルト: false (bool、optional)
rank_zero_only	rank 0 のプロセスでのみログするかどうか。Artifacts をログする場合は、すべての rank でログすることを強く推奨します。rank ≥1 の Artifacts は保存されないため、重要な情報が失われる可能性があります。たとえば、Deepspeed ZeRO を使用している場合、すべての rank の Artifacts がなければ checkpoint から復元できません。デフォルト: True (bool、optional)
init_kwargs	config などの、wandb.init() に渡すパラメーターです。wandb.init() が受け取るパラメーターについては、wandb.init() parametersを参照してください。

一般的な使用例は次のとおりです。

init_kwargs = {"notes":"Testing higher learning rate in this experiment", 
               "config":{"arch":"Llama",
                         "use_mixed_precision":True
                         }
               }

wandb_logger = WandBLogger(log_artifacts=True, init_kwargs=init_kwargs)

予測サンプルをログする

WandBLogger を介して W&B へのログのタイミングを制御するには、Composer’s Callbacks システムを使用できます。この例では、検証用画像と予測のサンプルをログします。

import wandb
from composer import Callback, State, Logger

class LogPredictions(Callback):
    def __init__(self, num_samples=100, seed=1234):
        super().__init__()
        self.num_samples = num_samples
        self.data = []
        
    def eval_batch_end(self, state: State, logger: Logger):
        """バッチごとに予測を計算し、self.dataに保存する"""
        
        if state.timer.epoch == state.max_duration: #最後の検証エポック時
            if len(self.data) < self.num_samples:
                n = self.num_samples
                x, y = state.batch_pair
                outputs = state.outputs.argmax(-1)
                data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs[:n]))]
                self.data += data
            
    def eval_end(self, state: State, logger: Logger):
        with wandb.init() as run:
            "wandb.Tableを作成してログする"
            columns = ['image', 'ground truth', 'prediction']
            table = wandb.Table(columns=columns, data=self.data[:self.num_samples])
            run.log({'sample_table':table}, step=int(state.timer.batch))         
...

trainer = Trainer(
    ...
    loggers=[WandBLogger()],
    callbacks=[LogPredictions()]
)