W&B の学習 run で Weave を使用する

Weave のトレースを W&B の学習 run と統合し、学習中のモデル動作を詳細に可視化します。関数実行の詳細や診断情報を、従来の ML メトリクスと並べて、カスタマイズ可能な Workspace のダッシュボードに記録・表示できます。

​

モデル学習 run 中にトレースをログする

モデル学習 run 中に Weave のトレースをログし、Personal Workspace で表示できるようになりました。

Personal Workspace は UI 上のカスタマイズ可能なダッシュボードで、モデル学習 run の詳細な分析やデータの可視化を行うことができます。Workspace に Weave パネルを追加すると、run 中にログされたトレースデータを確認およびアクセスでき、モデルが学習中にどのように動作しているかをより深く把握できます。

たとえば、LLM モデルをファインチューニングしていて、W&B Workspace 上でステップ 500 で精度が低下していることがわかった場合、連携された Weave のトレースから、評価基準を満たさない冗長な応答をモデルが生成し始めた、といった事象を正確に突き止めることができます。

​

Weave Panels を使用する

ML パイプライン内の関数に @weave.op デコレータを付けることで、その実行情報が自動的にキャプチャされ、自分用のワークスペースからアクセスできるようになります。

たとえば、次のスクリプトは、Weave のトレースが W&B の学習 run とどのように連携するかを示しています。これは、各学習ステップで 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 の実行中または完了後にワークスペースでトレース情報を確認するには、run 開始時にターミナルに表示されるリンク（次のような形式: wandb: 🚀 View run at https://wandb.ai/wandb/my-project/runs/<run-ID>）を開くか、UI からワークスペースに移動します。

UI からワークスペースに移動するには:

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

ワークスペースの詳細については、View experiments results を参照してください。

​

Weave トレースで W&B アーティファクトを参照する

Weave トレース内で、W&B artifacts（モデル、データセット、チェックポイントなど）を参照できます。これにより Weave UI にクリック可能なリンクが作成され、アーティファクトの詳細ページへ直接遷移できるようになります。どのアーティファクトのバージョンが、どの特定のオペレーションで使用されたかを追跡するのに役立ちます。

たとえば、2 つのバージョンのモデル（v1 と v2）を学習した場合、そのモデルをクエリするときに、v1 アーティファクトの参照をトレースの属性として追加できます。これにより、どのアーティファクトのバージョンが特定の出力を生成したかを正確にトレースできます。

トレースにアーティファクトを関連付けるには、wandb-artifact:/// プレフィックスとアーティファクトのフルネームを連結して、次の構文でアーティファクトの URL を構築します: wandb-artifact:///<full-name-of-artifact>。その後、次のいずれかの方法で Weave に渡します。

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

アーティファクトのフルネームは、W&B UI のアーティファクト詳細ページで確認できます。形式は次のとおりです: <entity>/<project>/<artifact_name>:<version>。W&B UI のアーティファクト詳細ページからフルネームをコピーできます。

​

trace の属性として追加する

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

Python

import weave

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

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

# 任意のキー名と、その値としてアーティファクトの 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}!`;
  });

  // 任意のキー名と、その値としてアーティファクトの 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);

​

Model の属性としてアーティファクトを追加する

Model の属性としてアーティファクトを追加するには、モデルのインスタンスを生成するときにアーティファクトの URL を渡します。

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

import weave

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

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

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

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