Inférence par batch Qwen2.5-32B avec Ray Data et vLLM
Utilisez Qwen2.5-32B-Instruct pour classer 16 000 énoncés d'assistant vocal multilingue sur un Runtime IA 8xH100 attaché. Ce Notebook montre comment :
- Créez un dataset Ray multilingue équilibré à partir de MASSIVE 1.1.
- Exécutez une réplique de modèle vLLM persistante sur chaque GPU disponible.
- Surveillez la charge de travail avec le tableau de bord Ray et les métriques système MLflow.
- Enregistrez les résultats de prévision complets au format Parquet dans un volume Unity Catalog.
Cet exemple nécessite la version 5 ou une version ultérieure de l’environnement Databricks AI.
Se connecter au compute GPU serverless
- Dans le sélecteur de compute du Notebook, sélectionnez Serverless GPU .
- Dans le panneau Environment , sélectionnez l’accélérateur 8xH100 et l’environnement AI v5 .
- Cliquez sur Appliquer , puis confirmez l’environnement.
Le modèle Qwen est public et ne nécessite pas d’authentification Hugging Face. Le notebook download MASSIVE 1.1 depuis l’archive publique d’Amazon.
Importer des bibliothèques
AI v5 inclut les packages Ray, vLLM, Hugging Face Datasets, Transformers, PyTorch et MLflow utilisés dans ce Notebook ; aucune installation de package n'est donc requise.
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
Configurer le workload
Définissez le modèle, les locales, la taille de l’échantillon et les paramètres d’inférence.
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
Configurer le stockage Unity Catalog
Utilisez les widgets pour spécifier un catalogue, un schéma et un volume Unity Catalog existants. Le notebook stocke le cache MASSIVE et les prédictions Parquet dans ce volume. Vous avez besoin de ces privilèges :
USE CATALOGsur le catalogue etUSE SCHEMAsur le schéma.READ VOLUMEetWRITE VOLUMEsur le volume.
Chaque exécution MLflow écrit des prédictions dans son propre sous-répertoire sous la racine de sortie Parquet configurée.
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}")
start Ray
ray_init() démarre Ray sur le compute associé et affiche l’URL du tableau de bord pour ce notebook. La connexion Ray reste active tant que le notebook est connecté. Le pool d’acteurs utilise le nombre de GPU rapporté par Ray, de sorte que chaque GPU disponible exécute une réplique de modèle vLLM.
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.")
Charger et échantillonner MASSIVE
download MASSIVE 1.1 dans le cache configuré, puis sélectionnez les 2 000 mêmes exemples d’entraînement de chaque locale à chaque exécution. La première locale fournit également les noms de scénario et d’intention utilisés pour créer le prompt de classification.
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"]))
Créer le dataset Ray
Combinez les échantillons locaux, conservez les champs nécessaires à l'inférence et à l'évaluation, et repartitionnez les données afin que Ray puisse maintenir tous les acteurs prédicteurs occupés.
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.")
Définir le prédicteur vLLM
MASSIVE regroupe les énoncés en 18 scénarios, tels que alarm, weather et music. Ce notebook construit les étiquettes autorisées et les conseils de scénario vers intention à partir du dataset au lieu de les coder en dur.
Le mappage scénario-vers-intention aide Qwen à distinguer les étiquettes ayant des significations similaires. vLLM renvoie l'une des étiquettes autorisées et une étape de normalisation finale marque toute autre réponse comme invalide.
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
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
Exécuter et surveiller l’inférence par batch
VLLMPredictor charge Qwen une fois au start de chaque acteur, puis réutilise ce modèle pour chaque batch qu’il reçoit. Ray Data start un acteur par GPU détecté et planifie chaque batch sur l’acteur disponible suivant.
Pendant l’exécution de l’inférence, ouvrez l’URL du tableau de bord Ray affichée par ray_init() dans la cellule 10. Utilisez le tableau de bord pour inspecter les huit acteurs prédicteurs, les réservations de GPU, la progression des tâches, les logs et les stragglers.
predictions = input_dataset.map_batches(
VLLMPredictor,
batch_format="pandas",
batch_size=BATCH_SIZE,
compute=ray.data.ActorPoolStrategy(size=ACTOR_COUNT),
num_gpus=1,
)
Matérialiser et suivre les résultats
Ray Data construit ce pipeline de manière paresseuse, ainsi write_parquet() exécute l’inférence et enregistre les résultats en une seule étape. Spark lit ensuite les fichiers Parquet pour l’évaluation sans réexécuter le modèle. Le run MLflow environnant capture les paramètres de workload, les métriques de qualité, le timing, le throughput et les métriques système, et Databricks ajoute un (1 MLflow run) Link cliquable sous la cellule une fois terminé.
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.")
Valider les résultats
Les vérifications ci-dessous confirment que la sortie contient une ligne par entrée et que chaque acteur prédicteur a traité au moins un batch.
Le chronométrage commence avant que Ray ne crée les acteurs et ne charge le modèle ; la durée et le throughput rapportés incluent donc le temps de cold-start.
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']}")
Analyser la qualité des prédictions
Afficher la précision par locale, un échantillon de prédictions et la distribution des enregistrements entre les acteurs GPU.
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)
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)