Ray DataとvLLMを使用したQwen2.5-32Bのバッチ推論
アタッチされた 8xH100 AI Runtime 上で Qwen2.5-32B-Instruct を使用し、16,000 件の多言語音声アシスタントの発話を分類します。このノートブックでは、以下の方法を示します：
- MASSIVE 1.1 からバランスの取れた多言語 Ray データセットを構築する。
- 利用可能な各 GPU で 1 つの永続的な vLLM モデルレプリカを実行します。
- Ray ダッシュボードと MLflow システムメトリクスを使用してワークロードを監視します。
- 完全な予測結果を Unity Catalog ボリュームに Parquet として保存します。
注記
この例では、Databricks AI 環境バージョン 5 以降が必要です。
Serverless GPU コンピュートへの接続
- ノートブックのコンピュートセレクターから、 [Serverless GPU] を選択します。
- [環境] パネルで、 8xH100 アクセラレータと AI v5 環境を選択します。
- [適用] をクリックし、環境を確認します。
Qwenモデルは公開されており、Hugging Face認証は不要です。このノートブックは、AmazonのパブリックアーカイブからMASSIVE 1.1をダウンロードします。
ライブラリのインポート
AI v5には、このノートブックで使用されるRay、vLLM、Hugging Face Datasets、Transformers、PyTorch、およびMLflowパッケージが含まれているため、パッケージのインストールは不要です。
Python
import json
import re
import time
from pathlib import Path
import mlflow
import pandas as pd
from datasets import DownloadConfig, DownloadManager, concatenate_datasets, load_dataset
from datasets.utils.logging import disable_progress_bar
from pyspark.sql import functions as F
from vllm import LLM, SamplingParams
from vllm.sampling_params import StructuredOutputsParams
ワークロードを構成する
モデル、ロケール、サンプルサイズ、および推論パラメーターを設定します。
Python
MODEL_NAME = "Qwen/Qwen2.5-32B-Instruct"
DATASET_NAME = "AmazonScience/massive"
MASSIVE_ARCHIVE_URL = "https://amazon-massive-nlu-dataset.s3.amazonaws.com/amazon-massive-dataset-1.1.tar.gz"
LOCALES = ["en-US", "es-ES", "de-DE", "ar-SA", "hi-IN", "ja-JP", "sw-KE", "zh-CN"]
ROWS_PER_LOCALE = 2_000
BATCH_SIZE = 64
MAX_MODEL_LEN = 512
MAX_OUTPUT_TOKENS = 8
SEED = 42
Unity Catalogストレージの構成
ウィジェットを使用して、既存の Unity Catalog のカタログ、スキーマ、およびボリュームを指定します。ノートブックは、MASSIVE キャッシュと Parquet 予測をこのボリュームに保存します。以下の権限が必要です：
USE CATALOGカタログ上、およびスキーマ上のUSE SCHEMA。READ VOLUMEおよびボリューム上のWRITE VOLUME。
各 MLflow ランは、構成された Parquet 出力ルートの下にある独自のサブディレクトリに予測を書き込みます。
Python
widget_defaults = {
"uc_catalog": "main",
"uc_schema": "default",
"uc_volume": "ray_data",
}
for widget_name, default_value in widget_defaults.items():
dbutils.widgets.text(widget_name, default_value)
CATALOG = dbutils.widgets.get("uc_catalog")
SCHEMA = dbutils.widgets.get("uc_schema")
VOLUME = dbutils.widgets.get("uc_volume")
volume_path = f"/Volumes/{CATALOG}/{SCHEMA}/{VOLUME}"
parquet_output_root = f"{volume_path}/sgc-raydata-vllm-batch-inference"
massive_cache_path = f"{volume_path}/hf-cache/amazon-massive-1.1"
print(f"Parquet output root: {parquet_output_root}")
print(f"Dataset cache: {massive_cache_path}")
Ray を起動する
ray_init() アタッチされたコンピュート上で Ray を起動し、このノートブックのダッシュボード URL を出力します。ノートブックが接続されている間、Ray 接続はアクティブなままになります。アクタープールは Ray によって報告された GPU 数を使用するため、利用可能な各 GPU で 1 つの vLLM モデルレプリカが実行されます。
Python
import ray
from serverless_gpu import ray_init
ray_context = ray_init()
ACTOR_COUNT = int(ray.cluster_resources().get("GPU", 0))
if ACTOR_COUNT < 1:
raise RuntimeError("Ray did not detect a GPU. Attach GPU compute and run the notebook again.")
print(f"Ray detected {ACTOR_COUNT} GPUs; using {ACTOR_COUNT} predictor actors.")
MASSIVE の読み込みとサンプリング
Download MASSIVE 1.1を構成済みのキャッシュに、その後、ランのたびに各ロケールから同じ2,000件のトレーニング例を選択してください。最初のロケールは、分類プロンプトの構築に使用されるシナリオ名とインテント名も提供します。
Python
disable_progress_bar()
download_config = DownloadConfig(cache_dir=f"{massive_cache_path}/downloads")
download_manager = DownloadManager(download_config=download_config)
massive_archive_dir = Path(download_manager.download_and_extract(MASSIVE_ARCHIVE_URL))
massive_data_dir = massive_archive_dir / "1.1" / "data"
locale_datasets = []
scenario_names = None
scenario_intents = None
for locale in LOCALES:
locale_dataset = load_dataset(
"json",
data_files=str(massive_data_dir / f"{locale}.jsonl"),
split="train",
cache_dir=f"{massive_cache_path}/datasets",
)
locale_dataset = locale_dataset.filter(lambda row: row["partition"] == "train")
locale_scenarios = sorted(locale_dataset.unique("scenario"))
if scenario_names is not None and locale_scenarios != scenario_names:
raise ValueError(f"Scenario labels differ for locale {locale}.")
if scenario_names is None:
scenario_names = locale_scenarios
label_frame = locale_dataset.select_columns(["scenario", "intent"]).to_pandas()
scenario_intents = {
scenario: sorted(group["intent"].unique())
for scenario, group in label_frame.groupby("scenario")
}
sample = locale_dataset.shuffle(seed=SEED).select(range(ROWS_PER_LOCALE))
locale_datasets.append(sample.select_columns(["id", "locale", "utt", "scenario"]))
Ray データセットを作成する
ロケールのサンプルを結合し、推論と評価に必要なフィールドを保持して、Ray がすべての予測子アクタを稼働させ続けられるようにデータを再パーティション化します。
Python
massive_sample = concatenate_datasets(locale_datasets)
records = [
{
"input_id": f"{row['locale']}:{row['id']}",
"locale": row["locale"],
"utterance": row["utt"],
"expected_scenario": row["scenario"],
}
for row in massive_sample
]
input_dataset = ray.data.from_items(records).repartition(ACTOR_COUNT * 8)
print(f"Prepared {len(records):,} records across {len(LOCALES)} locales and {len(scenario_names)} scenarios.")
vLLM 予測子を定義する
MASSIVE は、alarm、weather、music などの 18 のシナリオに発話（utterances）をグループ化します。このノートブックは、許可されたラベルとシナリオからインテントへのガイダンスをハードコーディングするのではなく、データセットから構築します。
シナリオとインテントのマッピングは、Qwenが類似した意味を持つラベルを区別するのに役立ちます。vLLMは許可されたラベルのいずれかを返し、最終的な正規化ステップでその他の回答を無効としてマークします。
Python
scenario_set = set(scenario_names)
scenario_guidance = "\n".join(
f"- {scenario}: {', '.join(scenario_intents[scenario])}"
for scenario in scenario_names
)
system_prompt = (
"Classify the user utterance into exactly one MASSIVE scenario. "
"Use these scenario-to-intent mappings to distinguish similar labels:\n"
f"{scenario_guidance}\n"
"Return only the scenario label."
)
def format_prompt(tokenizer, utterance: str) -> str:
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": utterance},
]
return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
def normalize_label(response: str) -> str | None:
normalized = re.sub(r"[^a-z]+", " ", response.lower()).strip()
return normalized if normalized in scenario_set else None
Python
class VLLMPredictor:
def __init__(self):
gpu_ids = ray.get_runtime_context().get_accelerator_ids().get("GPU", [])
if len(gpu_ids) != 1:
raise RuntimeError(f"Expected one GPU per actor, but received {gpu_ids}.")
self.gpu_assignment = str(gpu_ids[0])
self.llm = LLM(
model=MODEL_NAME,
tensor_parallel_size=1,
dtype="bfloat16",
max_model_len=MAX_MODEL_LEN,
max_num_seqs=BATCH_SIZE,
gpu_memory_utilization=0.90,
enable_prefix_caching=True,
)
self.tokenizer = self.llm.get_tokenizer()
self.sampling_params = SamplingParams(
temperature=0.0,
max_tokens=MAX_OUTPUT_TOKENS,
structured_outputs=StructuredOutputsParams(choice=scenario_names),
)
def __call__(self, batch: pd.DataFrame) -> pd.DataFrame:
prompts = [format_prompt(self.tokenizer, utterance) for utterance in batch["utterance"]]
outputs = self.llm.generate(prompts, self.sampling_params, use_tqdm=False)
raw_responses = [output.outputs[0].text.strip() for output in outputs]
predicted_scenarios = [normalize_label(response) for response in raw_responses]
result = batch.copy()
result["raw_response"] = raw_responses
# Preserve invalid responses as nulls with a stable string type across batches.
result["predicted_scenario"] = pd.array(predicted_scenarios, dtype="string")
result["valid_prediction"] = result["predicted_scenario"].notna()
result["correct"] = (result["predicted_scenario"] == result["expected_scenario"]).fillna(False)
result["model_name"] = MODEL_NAME
result["ray_gpu_assignment"] = self.gpu_assignment
return result
バッチ推論の実行と監視
VLLMPredictor 各アクターの起動時に Qwen を 1 回読み込み、その後受信するすべてのバッチでそのモデルを再利用します。Ray Data は検出された GPU ごとに 1 つのアクターを起動し、次に利用可能なアクターで各バッチをスケジュールします。
推論の実行中に、セル 10 で ray_init() によって出力された Ray ダッシュボード URL を開きます。ダッシュボードを使用して、8 つの予測アクター、GPU 予約、タスクの進行状況、Logs、およびストラッグラーを検査します。
Python
predictions = input_dataset.map_batches(
VLLMPredictor,
batch_format="pandas",
batch_size=BATCH_SIZE,
compute=ray.data.ActorPoolStrategy(size=ACTOR_COUNT),
num_gpus=1,
)
結果をマテリアライズして追跡する
Ray Data はこのパイプラインを遅延評価で構築するため、write_parquet() は推論を実行し、結果を 1 つのステップで保存します。その後、Spark はモデルを再実行することなく、評価のために Parquet ファイルを読み取ります。周辺の MLflow ランは、ワークロードのパラメーター、品質メトリクス、タイミング、throughput、およびシステムメトリクスをキャプチャし、終了時に Databricks がセルの下にクリック可能な (1 MLflow run) Link を追加します。
Python
mlflow.set_system_metrics_sampling_interval(2)
with mlflow.start_run(run_name="raydata-massive-qwen25-32b", log_system_metrics=True) as active_run:
parquet_output_path = f"{parquet_output_root}/{active_run.info.run_id}"
print(f"Parquet output: {parquet_output_path}")
mlflow.log_params(
{
"model": MODEL_NAME,
"dataset": DATASET_NAME,
"dataset_version": "1.1",
"locales": json.dumps(LOCALES),
"record_count": len(records),
"actor_count": ACTOR_COUNT,
"batch_size": BATCH_SIZE,
"max_model_len": MAX_MODEL_LEN,
"max_output_tokens": MAX_OUTPUT_TOKENS,
"temperature": 0.0,
"output_constraint": "scenario_choices",
"system_metrics_interval_seconds": 2,
"gpu_memory_utilization": 0.90,
}
)
mlflow.set_tags(
{
"dataset_source": MASSIVE_ARCHIVE_URL,
"parquet_output_path": parquet_output_path,
}
)
start_time = time.perf_counter()
predictions.write_parquet(parquet_output_path)
cold_start_inclusive_duration_seconds = time.perf_counter() - start_time
results_df = spark.read.parquet(parquet_output_path)
aggregate = results_df.agg(
F.count("*").alias("record_count"),
F.avg(F.col("correct").cast("double")).alias("overall_accuracy"),
F.avg(F.col("valid_prediction").cast("double")).alias("valid_prediction_rate"),
F.countDistinct("ray_gpu_assignment").alias("unique_gpu_assignments"),
).first()
scenario_accuracy_df = results_df.groupBy("expected_scenario").agg(
F.count("*").alias("record_count"),
F.avg(F.col("correct").cast("double")).alias("accuracy"),
F.avg(F.col("valid_prediction").cast("double")).alias("valid_prediction_rate"),
).orderBy("expected_scenario")
macro_scenario_accuracy = scenario_accuracy_df.agg(F.avg("accuracy")).first()[0]
cold_start_inclusive_records_per_second = (
aggregate["record_count"] / cold_start_inclusive_duration_seconds
)
mlflow.log_metrics(
{
"overall_accuracy": aggregate["overall_accuracy"],
"macro_scenario_accuracy": macro_scenario_accuracy,
"valid_prediction_rate": aggregate["valid_prediction_rate"],
"cold_start_inclusive_duration_seconds": cold_start_inclusive_duration_seconds,
"cold_start_inclusive_records_per_second": cold_start_inclusive_records_per_second,
}
)
mlflow_run_id = active_run.info.run_id
print(f"MLflow run ID: {mlflow_run_id}")
print("Open the '(1 MLflow run)' link attached to this cell for parameters and metrics.")
結果を検証する
以下のチェックにより、出力に入力ごとの行が 1 つ含まれていること、およびすべての予測アクターが少なくとも 1 つのバッチを処理したことが確認されます。
タイミングはRayがアクタを作成してモデルをロードする前に開始されるため、報告される期間とthroughputにはコールド起動時間が含まれます。
Python
if aggregate["record_count"] != len(records):
raise RuntimeError("The persisted result count does not match the input count.")
if aggregate["unique_gpu_assignments"] != ACTOR_COUNT:
raise RuntimeError(f"Expected results from {ACTOR_COUNT} Ray GPU assignments.")
print(f"Records: {aggregate['record_count']:,}")
print(f"Overall accuracy: {aggregate['overall_accuracy']:.2%}")
print(f"Macro scenario accuracy: {macro_scenario_accuracy:.2%}")
print(f"Valid prediction rate: {aggregate['valid_prediction_rate']:.2%}")
print(f"Inference duration including actor and model cold start: {cold_start_inclusive_duration_seconds:.1f} seconds")
print(f"Throughput including actor and model cold start: {cold_start_inclusive_records_per_second:.1f} records/second")
print(f"Unique GPU assignments: {aggregate['unique_gpu_assignments']}")
予測品質の分析
ロケール別の精度、予測のサンプル、および GPU アクタ全体でのレコードの分布を表示します。
Python
locale_accuracy_df = (
results_df.groupBy("locale")
.agg(
F.count("*").alias("record_count"),
F.avg(F.col("correct").cast("double")).alias("accuracy"),
F.avg(F.col("valid_prediction").cast("double")).alias("valid_prediction_rate"),
)
.orderBy("locale")
)
print("Accuracy by locale:")
locale_accuracy_df.show(truncate=False)
Python
prediction_columns = [
"locale", "utterance", "expected_scenario", "predicted_scenario",
"correct", "ray_gpu_assignment",
]
sample_predictions_df = (
results_df.select(prediction_columns)
.orderBy(F.rand(SEED))
.limit(16)
)
actor_distribution_df = (
results_df.groupBy("ray_gpu_assignment")
.agg(F.count("*").alias("record_count"))
.orderBy("ray_gpu_assignment")
)
print("Sample predictions:")
sample_predictions_df.show(truncate=80)
print("Records by Ray GPU assignment:")
actor_distribution_df.show(truncate=False)