集合通信总结和 mpi4py 实践

# mpi_hello.py
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

print(f"Hello from process {rank}/{size}")

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def main():
    # 初始化分布式环境，并使用 NCCL 作为通信后端。
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    device = rank
    
    # 1. 验证数据是否相同
    data = torch.randn(20, 10).to(device)  # 每个进程独立生成随机数据
    print(f"Rank {rank}: Data mean={data.mean().item():.6f}, std={data.std().item():.6f}")
    
    # 2. 验证模型初始参数是否相同
    model = torch.nn.Linear(10, 10).to(device) # 各进程独立初始化
    print(f"Rank {rank}: Weight mean={model.weight.mean().item():.6f}")
    
    # 3. DDP包装后参数同步验证
    ddp_model = DDP(model, device_ids=[device]) # DDP同步参数
    print(f"Rank {rank}: DDP Weight mean={ddp_model.module.weight.mean().item():.6f}")
    
    # 4. 计算输出
    result = ddp_model(data)
    print(f"Rank {rank}: Output mean={result.mean().item():.6f}")

if __name__ == "__main__":
    main()

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

def demo_broadcast():
    """
    Broadcast：root 进程准备一份数据，把它广播给通信域内所有进程。
    """
    if rank == 0:
        data = {"msg": "Hello", "vec": np.arange(4)}
        print(f"[BCAST] Rank 0 初始化数据 -> {data}")
    else:
        data = None                 # 非 root 必须占位
    data = comm.bcast(data, root=0) # 广播
    print(f"[BCAST] Rank {rank} 收到数据 -> {data}")

if __name__ == "__main__":
    if rank == 0:
        print(f"=== 进程总数: {size} ===\n")

    comm.Barrier()   # 让所有进程同步后再开始演示
    if rank == 0: print("\n*** Broadcast 演示 ***")
    demo_broadcast()

# comm_ops.py
from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

def demo_broadcast():
    """
    Broadcast：root 进程准备一份数据，把它广播给通信域内所有进程。
    """
    if rank == 0:
        data = {"msg": "Hello", "vec": np.arange(4)}
        print(f"[BCAST] Rank 0 初始化数据 -> {data}")
    else:
        data = None                 # 非 root 必须占位
    data = comm.bcast(data, root=0) # 广播
    print(f"[BCAST] Rank {rank} 收到数据 -> {data}")

def demo_scatter():
    """
    Scatter：root 进程按进程数把一个可迭代对象切块，分别送到各进程。
    最终每个进程仅拿到自己的“那一块”。
    """
    if rank == 0:
        big_array = np.arange(size * 3, dtype='i')   # 举例：总长度 = 进程数 × 3
        chunks = np.split(big_array, size)           # 均匀切成 size 份
        print(f"[SCATTER] Rank 0 切割后 chunks = {chunks}")
    else:
        chunks = None
    local_arr = comm.scatter(chunks, root=0)  # 本进程获得一个形状 (3,) 的子数组
    print(f"[SCATTER] Rank {rank} 拿到 {local_arr}, 平均={local_arr.mean():.1f}")

def demo_gather(local_data):
    """
    Gather：把各进程局部结果收集到 root 进程，root 获得列表，其他进程得 None
    Allgather：所有进程都得到完整列表
    """
    gathered = comm.gather(local_data, root=0)
    if rank == 0:
        print(f"[GATHER] Root 收到来自所有进程的数据 -> {gathered}\n")

    gathered_all = comm.allgather(local_data)
    print(f"[ALLGATHER] Rank {rank} 收到完整数据列表 -> {gathered_all}")

if __name__ == "__main__":
    if rank == 0:
        print(f"\n=== 进程总数: {size} ===")

    comm.Barrier()
    if rank == 0: print("\n*** Broadcast 演示 ***")
    demo_broadcast()

    comm.Barrier()
    if rank == 0: print("\n*** Scatter + Gather 演示 ***")
    local = demo_scatter()

    comm.Barrier()
    demo_gather(local)

# collective_ops.py
from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank, size = comm.Get_rank(), comm.Get_size()

def demo_allgather():
    """
    每个进程发出局部向量，所有进程都得到完整列表
    """
    local_vec = np.arange(rank*3, rank*3 + 3, dtype='i')  # eg Rank2->[6 7 8]
    gathered = comm.allgather(local_vec)
    print(f"[ALLGATHER] Rank {rank}: local={local_vec} → gathered={gathered}")

def demo_reduce():
    """
    将标量归约到 root；这里只做求和
    """
    local_val = (rank + 1) ** 2          # 1,4,9,16
    total = comm.reduce(local_val, op=MPI.SUM, root=0)
    if rank == 0:
        print(f"\n[REDUCE] 汇总结果 (sum) = {total}")  # 1+4+9+16=30

def demo_allreduce():
    """
    所有进程都同时得到归约值；这里做 max
    """
    local_val = (rank + 1) * 2           # 2,4,6,8
    global_max = comm.allreduce(local_val, op=MPI.MAX)
    print(f"[ALLREDUCE] Rank {rank}: local={local_val}, global_max={global_max}")

if __name__ == "__main__":
    if rank == 0:
        print(f"\n=== 进程总数: {size} ===\n")

    comm.Barrier();  demo_allgather()
    comm.Barrier();  demo_reduce()
    comm.Barrier();  demo_allreduce()

# alltoall_equal.py
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
rank, size = comm.Get_rank(), comm.Get_size()

# 准备发送缓冲区：一个二维数组 shape=(size, k)
k = 2
sendbuf = np.arange(rank*size*k, (rank+1)*size*k, dtype='i').reshape(size, k)
# 发送给自己的那一行可用作“自环”数据
print(f"Rank {rank} sendbuf=\n{sendbuf}")

# 预先分配接收同形数组
recvbuf = np.empty_like(sendbuf)
comm.Alltoall(sendbuf, recvbuf)

print(f"Rank {rank} recvbuf=\n{recvbuf}")