multi-gpu-deep

多 GPU 与分布式部署深度分析

概述

当单个 GPU 的显存无法容纳完整嵌入表，或需要在多块 GPU / 多节点上部署相同模型时，NVE 提供了多种资源共享和分片方案。本文档从原理到代码实现，详细分析每种方案的设计和用法。

核心理念：嵌入层之间不共享 GPU 缓存；共享仅限于更大的后备存储资源。

---

一、共享场景分类

场景 1：单节点多 GPU 推理

同一节点部署同一模型的多个副本（每个 GPU 一个），各副本读取相同嵌入数据。共享方案：

嵌入层类型	C++	Python	共享方式
GPUEmbeddingLayer	同一 embedding_table 指针	不支持（每个 NVEmbedding 独立权重）	指向同一 GPU/Host 缓冲区
LinearUVMEmbeddingLayer	同一 uvm_table 指针	同一 ManagedMemBlock	UVM 后备表在 Host 内存或跨 GPU 映射
HierarchicalEmbeddingLayer	同一 host table / remote PS	暂不支持 host table 共享	参数服务器天然共享

场景 2：单节点多 GPU 分片

将一个大嵌入表按行分片到多个 GPU，每个 GPU 拥有一部分物理内存，但所有分片映射到统一虚拟地址空间。

场景 3：多节点分片

跨节点共享嵌入表，通过 IMEX（Inter-Node Memory Extension）通道 + NVLink Fabric 实现。

---

二、GPUEmbeddingLayer — 指针共享

GPUEmbeddingLayer 从 config.embedding_table 获取用户传入的指针，不拥有该内存的所有权：

GPUTableConfig config;
config.embedding_table = shared_buffer;  // 同一指针传入多个 layer

多个 GPUEmbeddingLayer 实例使用同一 embedding_table 指针即可共享数据。

更好的方式：使用同一 GPUEmbeddingLayer 的不同执行上下文（ExecutionContext），而不是多个层实例：

// 单个层 + 多个上下文
auto layer = std::make_shared<GPUEmbeddingLayer<int64_t>>(config);
auto ctx1 = layer->create_execution_context(stream1, ...);
auto ctx2 = layer->create_execution_context(stream2, ...);

// 线程 1 使用 ctx1，线程 2 使用 ctx2
pool->submit([&]() { layer->lookup(ctx1, ...); });
pool->submit([&]() { layer->lookup(ctx2, ...); });

共享 GPU 嵌入表的注意事项：

• 所有 GPU kernel 通过 ScopedDevice 切换到目标设备后才能访问指针
• 如果指针指向的是另一个 GPU 的内存，需要 GPU 支持统一地址访问（NVLink 连接）
• 写入时（update() / accumulate()）需要应用层处理读后写冲突

---

三、LinearUVMEmbeddingLayer — UVM 表共享

这是最常用的多 GPU 共享方案。LinearUVMEmbeddingLayer 的 GPU cache 是私有的，但 UVM 后备表可以共享。

Host 内存共享（单节点）

// 1. 分配 UVM 后备表（首选 Host 位置）
void* linear_table = nullptr;
cudaMallocManaged(&linear_table, linear_table_size);

// 2. 优化内存访问：建议所有 GPU 访问，首选位置为 CPU
cudaMemAdvise(linear_table, linear_table_size,
              cudaMemAdviseSetAccessedBy, 0);  // GPU 0
cudaMemAdvise(linear_table, linear_table_size,
              cudaMemAdviseSetAccessedBy, 1);  // GPU 1
cudaMemAdvise(linear_table, linear_table_size,
              cudaMemAdviseSetPreferredLocation, cudaCpuDeviceId);

// 3. 为每块 GPU 创建 GpuTable（独立的 GPU cache）
GPUTableConfig cfg;
cfg.device_id = 0;
cfg.cache_size = 1ULL << 31;      // 2GB cache per GPU
cfg.uvm_table = linear_table;      // ← 共享指针
auto table0 = std::make_shared<GpuTable<int64_t>>(cfg);

cfg.device_id = 1;
auto table1 = std::make_shared<GpuTable<int64_t>>(cfg);

// 4. 创建独立的嵌入层
LinearUVMEmbeddingLayer<int64_t> layer0({}, table0);
LinearUVMEmbeddingLayer<int64_t> layer1({}, table1);

原理：UVM 后备表分配在 Host 内存（cudaMallocManaged + cudaMemAdviseSetPreferredLocation = cudaCpuDeviceId），各 GPU 通过 UVM page fault 按需获取页面。每个 GpuTable 管理自己的 GPU cache，cache miss 时从同一块 Host 后备表读取。

更新冲突：当多个层同时写入 UVM 表时，通过 disable_uvm_update 控制——只允许一个 GpuTable 负责 UVM 写入，其余 GPU 仅读取：

cfg0.disable_uvm_update = false;  // GPU 0 负责写入
cfg1.disable_uvm_update = true;   // GPU 1 只读

Python ManagedMemBlock

import pynve.nve as nve

# 分配 Host 内存块
memblock = nve.ManagedMemBlock(embedding_dim, num_table_rows, data_type)

# 多个 NVEmbedding 共享同一 memblock
emb0 = nve.NVEmbedding(num_table_rows, embedding_dim, data_type,
    cache_type=CacheType.LinearUVM, gpu_cache_size=cache_size,
    memblock=memblock, device=torch.device("cuda:0"))

emb1 = nve.NVEmbedding(num_table_rows, embedding_dim, data_type,
    cache_type=CacheType.LinearUVM, gpu_cache_size=cache_size,
    memblock=memblock, device=torch.device("cuda:1"))

ManagedMemBlock 内部调用 cudaMallocManaged，在 Python 端暴露为线性内存块。

---

四、CUDADistributedBuffer — CUDA 虚拟内存管理

核心概念

CUDADistributedBuffer（include/distributed.hpp）使用 CUDA Virtual Memory Management API 创建跨 GPU / 跨节点的统一虚拟地址映射缓冲区。

与 UVM 的关键区别：

维度	UVM (cudaMallocManaged)	CUDADistributedBuffer
物理分配	GPU 自动管理	每个 GPU 显式分配物理内存
地址映射	统一虚拟地址 + 硬件缺页	虚拟地址 + 显式映射
性能	page fault 有额外开销	直接物理内存访问，无缺页
跨节点	不支持	支持（通过 Fabric 句柄）
适用场景	单节点 Host 内存共享	多节点 NVLink 分片

构造函数流程

CUDADistributedBuffer::CUDADistributedBuffer(
    uint64_t size,
    std::shared_ptr<DistributedEnv> dist_env,
    BufferLocation location) {
    if (single_host_) {
        init_single_host(size, location);
    } else {
        init_multi_host(size);
    }
    env_->barrier();  // 全部初始化完成后同步
}

单节点初始化（init_single_host）

Rank 0 (root)
  │
  ├── collect_devices(): all_gather 收集各 rank 的设备 ID
  │
  ├── 计算 shard_size: 对齐到 num_shards × granularity
  │
  ├── 对每个有 GPU 的 rank:
  │     cuMemCreate(&handle, shard_size, ..., device_id)
  │     cuMemExportToShareableHandle(&fd, handle, POSIX_FILE_DESCRIPTOR)
  │     all_alloc_handles_.push_back(handle)
  │
  ├── broadcast(shareable_fds, root=0)
  │   broadcast(root_pid, root=0)
  │
  └── 非 root 进程:
        pidfd_open(root_pid) → pidfd_getfd() → 获取本地 FD
  │
  └── 所有进程:
        cuMemAddressReserve(&buffer_, total_size_, ...)  // 预留虚拟地址
        cuMemImportFromShareableHandle(&handle, fd, ...) // 导入句柄
        cuMemMap(buf_start, shard_size_, ..., handle, ...) // 映射
        cuMemRelease(handle)
  │
  └── cuMemSetAccess(buffer_, total_size_, access_descs, ...)
        // 设置所有 GPU 的读写权限

多节点初始化（init_multi_host）

每个节点（rank）
  │
  ├── check_imex(): 检查 /dev/nvidia-caps-imex-channels 是否存在
  │                  IMEX = Inter-Node Memory Extension
  │
  ├── cuMemCreate(&handle, shard_size_, ..., device_id)
  │   物理内存分配在本地 GPU
  │
  ├── cuMemExportToShareableHandle(&fabric_handle, handle, FABRIC)
  │   导出为 Fabric 句柄（跨节点共享）
  │
  ├── all_gather: 收集所有 rank 的 Fabric 句柄
  │
  ├── cuMemImportFromShareableHandle() → 导入各 rank 的句柄
  │
  ├── cuMemAddressReserve(&buffer_, total_size_, ...)
  │
  ├── cuMemMap() → 将每个分片映射到虚拟地址的对应位置
  │
  └── cuMemSetAccess() → 设置本地 GPU 访问权限

IMEX 通道 是 NVIDIA 多节点 NVLink 系统的核心组件，为跨节点 GPU 间通信提供硬件路由能力。

析构函数

CUDADistributedBuffer::~CUDADistributedBuffer() {
    env_->barrier();                 // 确保所有进程就绪
    for (auto& h : all_alloc_handles_) {
        cuMemUnmap(buf_start, shard_size_);  // 取消映射
        cuMemRelease(h);                     // 释放物理分配
    }
    env_->barrier();                 // 确保所有进程取消映射完毕
    cuMemAddressFree(buffer_, total_size_);  // 释放虚拟地址空间
}

---

五、MPI 分布式环境

DistributedEnv 接口

class DistributedEnv {
public:
    virtual size_t rank() const = 0;
    virtual size_t world_size() const = 0;
    virtual size_t device_count() const = 0;
    virtual int local_device() const = 0;
    virtual bool single_host() const = 0;
    virtual void barrier() = 0;
    virtual void broadcast(uintptr_t buf, size_t size, int root) = 0;
    virtual void all_gather(uintptr_t send, uintptr_t recv, size_t size) = 0;
};

local_device() 是关键方法：它确定当前进程应该使用哪块 GPU。对于单节点多进程场景，通常 local_device = rank % device_count()。

collect_devices()

uint64_t CUDADistributedBuffer::collect_devices(std::vector<int>& all_devices) {
    int local_device = env_->local_device();
    const auto world_size = env_->world_size();
    all_devices.resize(world_size);
    env_->all_gather(reinterpret_cast<uintptr_t>(&local_device),
                     reinterpret_cast<uintptr_t>(&all_devices[0]),
                     sizeof(local_device));
    uint64_t num_devices = 0;
    for (auto& d : all_devices) if (d >= 0) num_devices++;
    return num_devices;
}

---

六、Python 端的 MemBlock 系列

类层次

MemBlock (抽象基类, C++ binding)
  ├── LinearMemBlock     — cudaMallocHost 分配的 Host 固定内存
  ├── ManagedMemBlock    — cudaMallocManaged 分配的 UVM 内存
  ├── NVLMemBlock        — 单节点多 GPU NVLink 共享内存
  ├── MPIMemBlock        — 多节点 MPI 分布式共享内存
  └── DistMemBlock       — 分布式环境共享内存（别名，同 MPIMemBlock）

MPIMemBlock（Python）

import pynve.nve as nve
from mpi4py import MPI

comm = MPI.COMM_WORLD
mpi_rank = comm.Get_rank()
local_device_id = mpi_rank % torch.cuda.device_count()

# 分配跨节点共享的分布式内存块
memblock = nve.MPIMemBlock(embedding_dim, num_table_rows, data_type)

# 创建嵌入层（每个进程独立 GPU cache，共享后备表）
emb_layer = nve_layers.NVEmbedding(
    num_table_rows, embedding_dim, data_type,
    cache_type=CacheType.LinearUVM,
    gpu_cache_size=cache_size,
    memblock=memblock,
    device=torch.device(f"cuda:{local_device_id}"))

MPIMemBlock 内部绑定到 CUDADistributedBuffer，通过 MPI 通信交换 GPU 物理内存句柄。

NVLMemBlock（Python）

单节点多 GPU 场景下的 NVLink 共享内存：

# 在 GPU 0,1,2 上分配物理内存分片
memblock = nve.NVLMemBlock(embedding_dim, num_table_rows,
                           data_type, gpu_ids=[0, 1, 2])

NVLMemBlock 内部在指定的每块 GPU 上分配物理内存，通过 cuMemCreate / cuMemExportToShareableHandle / cuMemImportFromShareableHandle 将各分片映射到统一虚拟地址空间（与 init_single_host 逻辑相同）。

两种 NoCache 路径

对于 CacheType.NoCache，Python 端使用 LinearMemBlock（Host 固定内存）或 ManagedMemBlock（UVM 内存），嵌入向量直接存储在其中，没有 GPU cache 层。查找时通过 cuEmbed kernel 直接读取。

---

七、数据竞争与读写安全

只读场景

当嵌入表只读时（推理场景），不需要额外同步。每个 GpuTable 独立管理自己的 GPU cache，从共享后备表读取数据。cache 行通过 UVM 或 NVLink 读取。

写入场景

当共享资源被修改时，其他 layer 必须被阻塞直到修改完成：

Layer 0 (GPU 0)          Layer 1 (GPU 1)
    │                        │
    ├ update(keys, values)   │
    │   ↓                    │
    │ UpdateTable kernel     │
    │ (写入 UVM 后备表)       │
    │   ↓                    │
    │ cudaEventRecord        │
    │   ↓                    │
    │ cudaStreamSynchronize  │
    │                        ├ lookup(keys) → (等待完成)
    │                        │   ↓
    │                        │ GpuTable::find()
    │                        │ (读取 UVM 后备表)

同步方式:
  C++: cudaEvent + cudaStreamWaitEvent (GPU 级异步)
  Python: torch.cuda.synchronize() / MPI barrier()

Cache 一致性问题

如果直接修改共享后备表（绕过 layer->update()），GPU cache 中可能还持有旧值：

// 方案 1: 调用 update() 更新缓存和后备表
layer->update(ctx, num_keys, keys, stride, size, new_values);

// 方案 2: 调用 erase() 淘汰缓存条目（下次 lookup 时自动从后备表加载）
layer->erase(ctx, num_keys, keys, table_id);

---

八、性能特征对比

方案总结

方案	C++ 实现	Python 实现	跨节点	性能特征
Host 内存共享	cudaMallocHost + 多 GpuTable	ManagedMemBlock	❌	延迟高（UVM page fault）
UVM 共享	cudaMallocManaged + cudaMemAdvise	ManagedMemBlock	❌	延迟中等（UVM 自动迁移）
NVLink 分片	CUDADistributedBuffer (single_host)	NVLMemBlock	❌	延迟低（硬件直连）
多节点分片	CUDADistributedBuffer (multi_host)	MPIMemBlock	✅	延迟由 NVLink Fabric 决定

多 GPU 基准测试

项目在 benchmarks/multi_gpu_bench.py 中提供了完整的分布式基准测试框架，使用 torch.distributed + mpi4py 启动多进程训练/推理：

# 单节点 4 GPU
mpirun -n 4 python benchmarks/multi_gpu_bench.py \
    --memblock managed

# 跨节点（需要 IMEX 通道）
mpirun -n 8 --host node1:4,node2:4 \
    python benchmarks/multi_gpu_bench.py \
    --memblock mpi

---

九、术语表（GLOSSARY）

术语	含义
UVM	CUDA Unified Virtual Memory。统一虚拟地址空间，可同时映射到 GPU 和 Host 物理内存。
NVLink	NVIDIA 的高带宽 GPU 直连技术，支持多 GPU 间低延迟通信。
NVSwitch	NVLink 交换机，允许多 GPU 全互联。
IMEX	Inter-Node Memory Extension。跨节点 NVLink 通道，使多节点 GPU 能够直接访问对方物理内存。
Fabric Handle	CUDA 虚拟内存管理的跨节点共享句柄类型（CU_MEM_HANDLE_TYPE_FABRIC）。
DDP	Distributed Data Parallel。数据并行训练策略，每个 GPU 持有完整模型副本但处理不同数据分片。
Sharding	将嵌入表按行分割到多个 GPU/节点，每个 GPU 只持有部分物理内存。
page fault	当 GPU 访问尚未映射到本地显存的 UVM 页面时触发的缺页中断。
pidfd	Linux 进程文件描述符（pidfd_open + pidfd_getfd），允许跨进程共享文件句柄。