如何基于Ray通过张量并行实现自定义大模型离线批量推理？

针对无法单GPU容纳的自定义大模型，结合Ray生态实现无OOM的离线批量推理，核心思路是用Ray Actor托管模型张量分片，替代会复制全模型的map_batches API，以下是落地流程：

1. 预先完成模型张量分片 #

先根据模型结构设计分片策略，确保每个分片能被单GPU容纳：

• 若模型是Transformer架构，可按注意力头、FFN层维度拆分，或直接将不同Transformer层分配到不同GPU；
• 自定义模型需手动实现分片逻辑：拆分模型参数为多个独立分片，每个分片对应一个可独立运行的子模型单元，保证分片间的计算依赖可通过张量传递衔接。

2. 用Ray Actor托管单个模型分片 #

每个Ray Actor绑定1个GPU，仅加载对应模型分片，避免全模型复制：

import ray
import torch

# 初始化Ray集群（若未初始化）
ray.init()

# 自定义分片模型类：实现单分片的推理逻辑
class CustomModelShard(torch.nn.Module):
    def __init__(self, shard_id):
        super().__init__()
        # 根据shard_id定义对应分片的模型层/参数
        self.layers = self._build_shard_layers(shard_id)

    def _build_shard_layers(self, shard_id):
        # 替换为自定义模型的分片层构建逻辑
        return torch.nn.Sequential(...)

    def forward(self, x):
        # 替换为自定义分片的前向推理逻辑
        return self.layers(x)

# Ray Actor：托管单个模型分片，绑定1个GPU
@ray.remote(num_gpus=1)
class ModelShardActor:
    def __init__(self, shard_id, shard_state_dict):
        self.device = torch.device("cuda")
        self.shard = CustomModelShard(shard_id)
        self.shard.load_state_dict(shard_state_dict)
        self.shard.to(self.device).eval()

    def infer(self, input_tensor):
        input_tensor = input_tensor.to(self.device)
        with torch.no_grad():
            output = self.shard(input_tensor)
        return output.cpu()

# 初始化所有分片Actor：假设拆分为4个分片
shard_count = 4
# 预先拆分好的模型参数字典列表，每个元素对应一个分片的state_dict
shard_state_dicts = self._split_model_into_shards(original_model.state_dict())

shard_actors = [
    ModelShardActor.remote(i, shard_state_dicts[i])
    for i in range(shard_count)
]

3. 构建张量并行推理流水线 #

针对离线批量数据，按张量并行逻辑完成推理：

• 数据预处理：将批量输入（如tokenized文本）转换为模型所需张量，根据分片策略拆分输入张量（例如，若按注意力头分片，需将输入的query/key/value张量按头拆分）；
• 分片推理调度：按模型计算依赖顺序，将拆分后的输入张量传递给对应分片Actor，获取中间结果后传递给下一个分片；
• 结果聚合：收集所有分片的最终输出，按并行规则合并为完整推理结果。

示例批量推理逻辑：

def batch_infer(batch_inputs):
    # 1. 预处理批量输入，拆分为对应分片的输入张量
    shard_inputs = _split_batch_inputs(batch_inputs, shard_count)
    
    # 2. 调用所有分片Actor进行并行推理
    shard_outputs = ray.get([
        shard_actors[i].infer.remote(shard_inputs[i])
        for i in range(shard_count)
    ])
    
    # 3. 聚合分片输出为最终结果
    final_outputs = _aggregate_shard_outputs(shard_outputs)
    return final_outputs

4. 替代方案：流水线并行 #

若自定义模型更适合按层拆分（如多层串行结构），可改用流水线并行：

• 将模型的每一层（或连续几层）作为一个分片，用Ray Actor托管；
• 批量数据按流水线方式依次经过各层Actor，最终输出结果；
• 这种方式适合模型层数多但单层级参数较小的场景，同样能避免单GPU加载全模型。

关键优化点 #

• 禁用梯度计算：推理时始终用torch.no_grad()减少内存占用；
• 优化数据传输：尽量在GPU间直接传递张量（Ray支持GPU-aware内存共享），避免频繁CPU-GPU数据拷贝；
• 批量大小调优：根据单GPU内存容量调整推理批量，避免分片处理时OOM；
• 资源隔离：确保每个Ray Actor独占1个GPU，避免多进程抢占GPU内存。

内容的提问来源于stack exchange，提问作者ganto