W&B のトレーニング run で Weave を使用する

Workspace ダッシュボードで ML メトリクスと並べて関数実行の詳細を表示できるよう、Weave のトレースを W&B のトレーニング run に統合します。

​

モデルのトレーニング run中にトレースをログする

これで、モデルのトレーニング run中に Weave のトレースをログし、個人用Workspaceで確認できるようになりました。

個人用Workspaceは、UI 内でカスタマイズできるダッシュボードで、モデルのトレーニング runを詳細に分析したり、データ可視化を表示したりできます。Workspaceに Weave パネルを追加すると、モデル run中にログされたトレースデータを表示して確認できるため、トレーニング中のモデルの挙動をより深く把握できます。

たとえば、LLM モデルをファインチューニングしていて、W&B Workspace で step 500 で accuracy が低下していることが示された場合、Weave と統合されたトレースを使うと、評価基準を満たさない過度に冗長な応答をモデルが生成し始めた、といったように、何が起きたのかを正確に把握できます。

​

Weave パネルを使用する

ML パイプライン内の関数を @weave.op でデコレートすると、それらの実行情報が自動的に取得され、個人用の Workspace で参照できます。

たとえば、次のスクリプトは、Weave のトレースが W&B のトレーニング run とどのように統合されるかを示しています。これは機械学習のトレーニングループをシミュレートしたもので、各トレーニング step で W&B にメトリクスをログすると同時に、詳細な Weave トレースも生成します。関数に付与された @weave.op デコレータは、入力、出力、実行の詳細を自動的に取得します。トレーニングループが wandb.log() でメトリクスをログすると、Weave のトレース情報もプロジェクトにログされます。

import wandb
import weave
import random

# W&B と Weave を同じプロジェクトで初期化する
project = "my-workspace-project"
weave.init(project)

@weave.op
def evaluate_model(model_state, epoch):
    # 評価メトリクスをシミュレートする
    accuracy = 0.7 + epoch * 0.02 + random.uniform(-0.05, 0.05)
    
    # 特定のモデルの動作をトラッキングする
    test_responses = {
        "conciseness_check": "Model generated 500 words for simple question",
        "accuracy_check": "Model answered 8/10 questions correctly",
        "hallucination_check": "No factual errors detected"
    }
    
    return {
        "accuracy": accuracy,
        "diagnostic_results": test_responses
    }

@weave.op
def training_step(epoch, lr):
    # トレーニングロジックをシミュレートする
    loss = 1.0 / (epoch + 1) + random.uniform(-0.1, 0.1)
    
    # トレースを使った評価
    eval_results = evaluate_model(f"checkpoint_{epoch}", epoch)
    
    return {
        "loss": loss,
        "accuracy": eval_results["accuracy"],
        "diagnostics": eval_results["diagnostic_results"]
    }

# トレーニングループ
with wandb.init(project=project, config={"lr": 0.01}) as run:
    for epoch in range(5):
        # Weave トレースを使ってトレーニングを実行する
        results = training_step(epoch, run.config.lr)
        
        # W&B にログする - インテグレーションポイントを作成する
        run.log({
            "epoch": epoch,
            "loss": results["loss"],
            "accuracy": results["accuracy"]
        })

このサンプルスクリプトを実行すると、トレースがそのほかの run 情報とともにどのようにログされるかを確認できます。

run の実行中または実行後に Workspace でトレース情報を表示するには、run の開始時にターミナルに表示されるリンク (次のような形式です: wandb: 🚀 View run at https://wandb.ai/wandb/my-project/runs/<run-ID>) を開くか、UI から Workspace にアクセスします。

UI から Workspace にアクセスするには、次の手順に従います。

1. W&B UI を開き、Projects タブをクリックします。プロジェクトの一覧が表示されます。
2. プロジェクトの一覧から、run をログしたプロジェクトをクリックします。Workspaces ページが開きます。
3. Workspace を自動 Workspace として設定している場合、Workspace にはデータ可視化と run に関する情報が自動的に表示されます。run の Trace Data は、Workspace の Weave セクションにあります。Workspace が手動 Workspace の場合は、Add panels をクリックし、Add Panels メニューの Weave セクションから新しいパネルを選択して、Weave パネルを追加できます。

Workspace の詳細については、実験結果を表示するを参照してください。

​

特定の W&B run にトレースを関連付ける

デフォルトでは、Weave はアクティブな wandb.run を自動的に検出し、トレースを関連付けます。グローバルな wandb.run ではない特定の run にトレースを関連付ける必要がある場合は、set_wandb_run_context を使用してください。method の run_id 引数と step 引数に値を指定すると、トレースを特定の run および run step に関連付けることができます。

import weave

client = weave.init("<your-team-name>/<your-project-name>")

# 以降のすべてのトレースを特定のrunに関連付けます。
client.set_wandb_run_context(run_id="<your-run-ID>", step=<integer>)

@weave.op
def my_function(x):
    return x * 2

my_function(10)  # このトレースは step 5 の "my-run-id" に関連付けられています。

run_id は run の短縮 ID (例: 38m4t5ja) であり、完全な entity/project/run_id パスではありません。run の短縮 ID は、run の URL (https://wandb.ai/team/project/runs/<run_id>) または コードで wandb.run.id にアクセスする ことで確認できます。

オーバーライドを解除してグローバルな wandb.run に戻すには、clear_wandb_run_context を呼び出します:

client.clear_wandb_run_context()

my_function(20)  # このトレースはグローバルの wandb.run を使用します（存在する場合）。

​

Weave トレースで W&B artifacts を参照する

Weave トレースでは、W&B artifacts (モデル、データセット、チェックポイントなど) を参照できます。これにより、Weave UI にクリック可能なリンクが作成され、artifact の詳細ページへ直接移動できるようになります。特定の操作でどの artifact バージョンが使用されたかをトラッキングするのに役立ちます。

たとえば、モデルの 2 つのバージョン (v1 と v2) をトレーニングしている場合、そのモデルをクエリするときに、v1 artifact への参照をトレースの属性として追加できます。これにより、どの artifact バージョンが特定の出力を生成したのかを正確にトレースできます。

artifact をトレースに関連付けるには、wandb-artifact:/// プレフィックスと artifact の完全名を連結し、wandb-artifact:///<full-name-of-artifact> という構文で artifact の URL を作成します。次に、以下の 2 つの方法のいずれかで Weave に渡します。

• 任意の属性として
• Model の属性として

artifact の完全名は、W&B UI の詳細ページで確認できます。形式は次のとおりです: <entity>/<project>/<artifact_name>:<version>。完全名は、W&B UI の artifact の詳細ページからコピーできます。

​

トレース の属性として追加する

artifact を トレース の属性として追加するには、weave.attributes() コンテキストマネージャにキーと値のペアとして渡します。

Python

import weave

weave.init("your-team-name/your-project-name")

@weave.op
def my_function(name: str):
    return f"Hello, {name}!"

# 任意のキー名を指定し、値に artifact の URL を設定します
with weave.attributes({'artifact_id': 'wandb-artifact:///team-name/project-name/run-38m4t5ja-history:v0'}):
    result = my_function("World")

TypeScript

import {init, op, withAttributes} from 'weave';

async function main() {
  await init('wandb/docs');

  const myFunction = op(async function myFunction(name: string) {
    return `Hello, ${name}!`;
  });

  // 任意のキー名を指定し、値に artifact の URL を設定します
  const result = await withAttributes(
    {artifact_id: "wandb-artifact:///team-name/project-name/run-38m4t5ja-history:v0"},
    async () => myFunction('World')
  );

  console.log('Result:', result);
}

main().catch(console.error);

​

artifactを Model の属性として追加する

artifactを Model の属性として追加するには、Model のインスタンス化時にartifactの URL を渡します。

Model クラスは現在、Weave Python SDK でのみサポートされています。

import weave

weave.init("your-team-name/your-project-name")

class MyModel(weave.Model):
    # artifactの参照をモデルの属性として保存する
    artifact_id: str
    temperature: float = 0.7
    
    @weave.op
    def predict(self, query: str) -> str:
        # モデルの推論ロジックをここに記述する
        return f"Response to: {query}"

# artifactのURLをモデルの属性として渡す
model = MyModel(
    artifact_id="wandb-artifact:///team-name/project-name/run-38m4t5ja-history:v0"
)

result = model.predict("Hello, World!")