分布式框架下的数据处理与模型推理实践

概念

Ray 是一个开源的高性能分布式计算框架，旨在简化大规模分布式应用的开发和运行。它提供了灵活的任务调度、资源管理以及并行计算能力，使开发者能够轻松实现从单机到多节点的计算扩展。Ray 支持多种场景，包括分布式机器学习训练、强化学习、超参数优化、大规模数据处理和实时模型部署等。Ray 的核心概念是通过统一的 API，使开发者专注于逻辑开发，而不必关心底层的分布式实现细节。

通过模块化设计，Ray 集成了多个功能强大的库，如Ray Data（数据处理）、 Ray Train（分布式训练）、Ray Tune（超参数优化）、Ray Serve（模型部署）、RLlib（强化学习）等，为开发者提供了一站式的分布式计算解决方案。无论是构建 AI 应用还是解决复杂并行计算问题，Ray 都是一种高效且易用的选择。Ray 的模块化设计将复杂功能分解为独立子系统，常用模块包括：

• （1）Ray Core：基础的分布式任务调度和资源管理。
• （2）Ray Data：高效的分布式数据处理模块。
• （3）Ray Train：支持主流框架的分布式机器学习训练工具。
• （4）Ray Tune：超参数优化库，支持大规模调优任务。
• （5）Ray Serve：用于实时部署 AI 模型的高性能工具。
• （6）RLlib：分布式强化学习库，适用于复杂环境中的智能决策任务。

数据处理与模型推理

Ray Data 的分布式数据处理

在大模型训练过程中的数据处理阶段，往往有着海量结构化与非结构化数据，这些数据难以快速、有效的处理，造成数据处理效率低下、数据质量不高等问题。智算数据平台基于Ray分布式框架的计算能力，提供了大量常用的分布式数据处理算子，并且实现文本、图片、音频、视频等多种类型数据的流水线处理，为模型训练提供高质量预训练数据。除了基本的数据处理算子外，数据平台还将通过引入多模态大模型，补齐了处理多模态数据能力。

在大模型训练中，数据预处理的效率直接影响训练速度和模型性能。Ray Data 是一个强大的分布式数据处理模块，该模块支持多种文件类型的读写，如json、parquet等格式的文本、图片、视频等。Ray Data存储数据的格式是基于Apache Arrow，一种高效的列式数据格式，相比Python数据处理常用的pandas，在读取和写入大型文件时会更加高效。集成了map，filter等分布式计算算子，提供了高效的数据处理，适用于海量数据的处理任务。此外，Ray Data还支持批处理功能，加速数据处理。在和Actor配合使用的场景下，可以减少请求次数，提高了处理效率。

功能亮点

支持多数据格式：可以处理文本、图片、音频、视频等多模态数据。

分布式操作：通过内置的算子（如 map、filter、groupby 等），实现数据分片和高效并行计算。

与主流工具集成：兼容 Pandas、Spark，方便数据迁移和开发。

示例：分布式数据预处理代码

import ray
from ray.data import read_parquet


ray.init()  # 初始化 Ray


# 加载数据集
dataset = read_parquet("s3://bucket/data.parquet")


# 数据清洗和特征处理
processed_data = (
    dataset.map(lambda x: {"text_length": len(x["text"])})
           .filter(lambda x: x["text_length"] > 50)
)


# 分布式存储
processed_data.write_parquet("s3://bucket/processed_data.parquet")

Ray Data 的设计理念是利用分布式集群的算力，在保留开发者熟悉的数据操作方式的同时，提升数据处理的速度和扩展性。

GPU 分配与分布式推理

在模型推理场景下，特别是针对大模型的批量推理，资源分配和高效利用至关重要。Ray 提供了灵活的 @ray.remote 装饰器，支持任务的分布式调度以及 GPU 资源的高效管理。

示例：使用llava模型进行分布式推理代码

from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
from PIL import Image
import torch
import ray
import os


# 初始化 Ray
ray.init()


@ray.remote(num_gpus=1)  # 请求 1 个 GPU
def run_model_on_gpu(picture):
  processor = LlavaNextProcessor.from_pretrained("llava-hf/llama3-llava-next-8b-hf")
  model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llama3-llava-next-8b-hf", torch_dtype=torch.float16, device_map="auto")
  # prepare image and text prompt, using the appropriate prompt template
  image = Image.open(picture)
  # Define a chat histiry and use `apply_chat_template` to get correctly formatted prompt
  # Each value in "content" has to be a list of dicts with types ("text", "image")
  conversation = [
      {


        "role": "user",
        "content": [
            {"type": "text", "text": "描述这张图片的内容"},
            {"type": "image"},
          ],
      },
  ]
  prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
  inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
  # autoregressively complete prompt
  output = model.generate(**inputs, max_new_tokens=1000)
  return processor.decode(output[0], skip_special_tokens=True)


pic_list = ["/root/llava/img1.png","/root/llava/img2.png"]
results = ray.get([run_model_on_gpu.remote(pic) for pic in pic_list])
print(results)


# 关闭 Ray
ray.shutdown()

在运行时，我们使用nvidia-smi命令，可以查看到有两个ray子进程在并发处理推理任务。

最后输出结果如下，通过以下结果与原图的对比，我们可以直观地看到使用 Ray 进行分布式推理的流程和效果。

输出结果

具体而言，Ray 能够高效地将推理任务拆分到多个 GPU 上进行并行处理，每个 GPU 仅需处理分配的部分数据，这大幅提高了推理效率和整体吞吐量。

例如，在多张图片的分布式推理任务中，每一张图片被分发到不同的 GPU 节点上，分别执行推理操作。最终，结果被汇总生成完整的推理输出。在这一过程中：

推理结果：图片描述的准确性保持一致，基于分布式的推理不会影响结果的质量。

性能提升：推理时间缩短了一半以上，分布式计算的效率得到验证。

可扩展性：只要增加更多的 GPU 资源，Ray 的调度器即可自动分配更多的计算任务，实现：超越单机多卡的算力水平

当然，以上只是Ray分布式的一些基础应用，为了说明Ray在分布式推理中的可行性。我们也可以使用如Ray Data模块的map_batch去做离线的批量分布式推理，提高数据的处理效率；再比如通过Ray Server模块去部署大模型服务，做到支持实时推理，支持多用户并发请求等更多功能。

随着 NPU（如华为 Ascend 系列）等新型加速硬件的崛起，Ray 也扩展了对异构设备的支持。在集群中，Ray 也可以识别并调度到带有 NPU 的节点运行任务，利用国产芯片的高效性能支持大规模 AI 应用。

总结

Ray 作为新一代分布式计算框架，通过模块化设计和简单易用的接口，极大地降低了开发分布式应用的门槛。在数据处理方面，Ray Data 提供了强大的分布式数据处理能力，可以高效地处理海量多模态数据。在模型推理方面，Ray Core的Remote函数和资源调度功能使开发者能够充分利用集群中的 GPU，快速构建分布式推理服务。

随着 AI 的快速发展，Ray 已成为开发者构建大规模分布式系统的核心工具。在未来的应用中，无论是处理复杂的数据管道还是优化模型推理性能，Ray 都将发挥越来越重要的作用。

本文转载自 ​​AI遇见云​​，作者： 景泓斐