Inferencia por lotes Qwen2.5-32B con Ray Data y vLLM

Abrir la versión 'notebook' de esta página

Utilice Qwen2.5-32B-Instruct para clasificar 16 000 enunciaciones multilingües de asistente de voz en un AI Runtime 8xH100 adjunto. Este cuaderno muestra cómo:

Construye un conjunto de datos de rayos multilingüe balanceado a partir de MASSIVE 1.1.
Ejecuta una réplica persistente del modelo vLLM en cada GPU disponible.
Monitoriza la carga de trabajo con el panel de control Ray y las métricas del sistema MLflow.
Guarda los resultados completos de predicción como Parquet en un volumen de Unity Catalog.

Note

Este ejemplo requiere el entorno de IA Databricks versión 5 o superior.

Conectar al cómputo de GPU sin servidor

Desde el selector de cómputo del portátil, selecciona GPU sin servidor.
En el panel Entorno, selecciona el acelerador 8xH100 y el entorno AI v5.
Haz clic en Aplicar y luego confirma el entorno.

El modelo Qwen es público y no requiere autenticación de Hugging Face. El cuaderno de notas descarga MASSIVE 1.1 del archivo público de Amazon.

Importar bibliotecas

AI v5 incluye los paquetes Ray, vLLM, Hugging Face Datasets, Transformers, PyTorch y MLflow usados en este portátil, por lo que no se requiere instalación de paquetes.

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

Configurar la carga de trabajo

Establece el modelo, los lugares, el tamaño de la muestra y los parámetros de inferencia.

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

Configurar almacenamiento en el Catálogo de Unity

Utiliza los widgets para especificar un catálogo, un esquema y un volumen existentes en Unity Catalog. El cuaderno almacena la ENORME caché y las predicciones de Parquet en este volumen. Necesitas estos privilegios:

USE CATALOG en el catálogo y USE SCHEMA en el esquema.
READ VOLUME y WRITE VOLUME en el volumen.

Cada ejecución de MLflow escribe predicciones en su propio subdirectorio bajo la raíz de salida configurada de Parquet.

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}")

Iniciar Ray

ray_init() inicia Ray en el proceso asociado y muestra la URL del panel de este cuaderno. La conexión Ray permanece activa mientras el portátil permanece conectado. El pool de actores utiliza el recuento de GPU reportado por Ray, por lo que cada GPU disponible ejecuta una réplica de modelo 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.")

Carga y muestrea MASSIVE

Descarga MASSIVE 1.1 en la caché configurada y luego selecciona los mismos 2.000 ejemplos de entrenamiento de cada localidad en cada partida. La primera localización también proporciona los nombres de escenario e intención usados para construir la consigna de clasificación.

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"]))

Crear el conjunto de datos Ray

Combina las muestras de localidad, conserva los campos necesarios para la inferencia y evaluación, y reparticiona los datos para que Ray pueda mantener ocupados a todos los actores predictores.

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.")

Defina el predictor vLLM

MASSIVE agrupa enunciados en 18 escenarios, como alarm, weather, y music. Este cuaderno genera las etiquetas permitidas y la guía de correspondencia entre escenarios e intenciones a partir del conjunto de datos, en lugar de definirlas explícitamente en el código.

El mapeo de escenario a intención ayuda a Qwen a distinguir etiquetas con significados similares. vLLM devuelve una de las etiquetas permitidas, y un paso final de normalización marca cualquier otra respuesta como inválida.

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

Ejecutar y monitorizar la inferencia por lotes

VLLMPredictor carga Qwen una vez cuando cada actor inicia, luego reutiliza ese modelo para cada lote que recibe. Ray Data inicia un actor por cada GPU detectada y programa cada lote en el siguiente actor disponible.

Mientras se ejecuta la inferencia, abre la URL del panel de Ray impresa ray_init() en la celda 10. Utilice el panel para inspeccionar los ocho actores de predicción, las reservas de GPU, el progreso de las tareas, los registros y las tareas rezagadas.

predictions = input_dataset.map_batches(
    VLLMPredictor,
    batch_format="pandas",
    batch_size=BATCH_SIZE,
    compute=ray.data.ActorPoolStrategy(size=ACTOR_COUNT),
    num_gpus=1,
)

Materializa y sigue los resultados

Ray Data construye esta tubería de forma perezosa, así que write_parquet() ejecuta inferencias y guarda los resultados en un solo paso. Spark entonces lee los archivos de Parquet para su evaluación sin ejecutar el modelo de nuevo. La ejecución de MLflow circundante captura parámetros de carga de trabajo, métricas de calidad, tiempo, rendimiento y métricas del sistema, y Databricks añade un enlace clicable (1 MLflow run) debajo de la celda cuando termina.

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.")

Validar los resultados

Las comprobaciones a continuación confirman que la salida contiene una fila por entrada y que cada actor predictor gestionó al menos un lote.

La medición del tiempo comienza antes de que Ray cree los actores y cargue el modelo, por lo que la duración y la capacidad de procesamiento indicadas incluyen el tiempo de arranque en frío.

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']}")

Analizar la calidad de las predicciones

Muestra la precisión por localización, una muestra de predicciones y la distribución de los registros entre los actores de la 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)

Cuaderno de ejemplo

Inferencia por lotes Qwen2.5-32B con Ray Data y vLLM

Obtención del cuaderno

Comentarios

¿Le resultó útil esta página?

Last updated on 2026-09-11

Inferencia por lotes Qwen2.5-32B con Ray Data y vLLM

Conectar al cómputo de GPU sin servidor

Importar bibliotecas

Configurar la carga de trabajo

Configurar almacenamiento en el Catálogo de Unity

Iniciar Ray

Carga y muestrea MASSIVE

Crear el conjunto de datos Ray

Defina el predictor vLLM

Ejecutar y monitorizar la inferencia por lotes

Materializa y sigue los resultados

Validar los resultados

Analizar la calidad de las predicciones

Cuaderno de ejemplo

Inferencia por lotes Qwen2.5-32B con Ray Data y vLLM

Comentarios

Recursos adicionales