01-PyTorch分布式源码解析

# torch/distributed/distributed_c10d.py

def init_process_group(
    backend: str,
    init_method: Optional[str] = None,
    timeout: timedelta = default_pg_timeout,
    world_size: int = -1,
    rank: int = -1,
    store: Optional[Store] = None,
    group_name: str = "",
    pg_options: Optional[ProcessGroupOptions] = None,
):
    """
    初始化分布式进程组

    Args:
        backend: 通信后端 ("nccl", "gloo", "mpi")
        init_method: 初始化方法 ("env://", "tcp://", "file://")
        world_size: 进程总数
        rank: 当前进程rank
        store: 分布式KV存储
    """
    global _pg_group_ranks
    global _backend
    global _default_pg_init_method

    # 1. 参数验证
    if backend not in ["nccl", "gloo", "mpi", "ucc"]:
        raise ValueError(f"Invalid backend: {backend}")

    # 2. 创建或获取Store (用于进程间协调)
    if store is None:
        if init_method is None:
            init_method = "env://"  # 默认使用环境变量

        # 根据init_method创建Store
        if init_method == "env://":
            # 从环境变量获取: MASTER_ADDR, MASTER_PORT
            store = _create_store_from_env()
        elif init_method.startswith("tcp://"):
            # TCP方式: tcp://master_ip:port
            store = TCPStore(...)
        elif init_method.startswith("file://"):
            # 文件方式: file:///path/to/file
            store = FileStore(...)

    # 3. 如果world_size和rank未指定,从环境变量获取
    if world_size == -1:
        world_size = int(os.environ.get("WORLD_SIZE", -1))
    if rank == -1:
        rank = int(os.environ.get("RANK", -1))

    # 4. 创建ProcessGroup
    if backend == "nccl":
        pg = ProcessGroupNCCL(store, rank, world_size, pg_options)
    elif backend == "gloo":
        pg = ProcessGroupGloo(store, rank, world_size, pg_options)
    elif backend == "mpi":
        pg = ProcessGroupMPI(rank, world_size)

    # 5. 设置为默认ProcessGroup
    _set_default_pg(pg)
    _backend = backend
    _default_pg_init_method = init_method

    # 6. 初始化同步屏障
    barrier()

    return pg


def _create_store_from_env():
    """从环境变量创建Store"""
    master_addr = os.environ.get("MASTER_ADDR", "localhost")
    master_port = int(os.environ.get("MASTER_PORT", "29500"))
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))

    # 创建TCPStore
    # rank 0 是server, 其他rank是client
    is_master = (rank == 0)

    store = TCPStore(
        host_name=master_addr,
        port=master_port,
        world_size=world_size,
        is_master=is_master,
        timeout=timedelta(seconds=300),
    )

    return store

# torch/distributed/distributed_c10d.py

def all_reduce(tensor, op=ReduceOp.SUM, group=None, async_op=False):
    """
    对tensor执行AllReduce操作

    Args:
        tensor: 输入/输出tensor (in-place操作)
        op: 归约操作 (SUM, PRODUCT, MIN, MAX, etc.)
        group: ProcessGroup (默认使用全局默认组)
        async_op: 是否异步执行

    Returns:
        如果async_op=True, 返回Work对象用于同步
        否则返回None
    """
    # 获取ProcessGroup
    if group is None:
        group = _get_default_group()

    # 参数验证
    if not tensor.is_cuda and group._get_backend_name() == "nccl":
        raise ValueError("NCCL backend requires CUDA tensors")

    # 执行AllReduce
    opts = AllreduceOptions()
    opts.reduceOp = op

    work = group.allreduce([tensor], opts)

    if async_op:
        return work
    else:
        work.wait()
        return None


# 异步AllReduce使用示例
async_work = dist.all_reduce(tensor, async_op=True)
# ... 做其他计算 ...
async_work.wait()  # 等待完成

# torch/nn/parallel/distributed.py

class DistributedDataParallel(Module):
    """
    分布式数据并行

    将模型复制到多个GPU,每个GPU处理不同的数据batch
    梯度通过AllReduce同步
    """

    def __init__(
        self,
        module,
        device_ids=None,
        output_device=None,
        dim=0,
        broadcast_buffers=True,
        process_group=None,
        bucket_cap_mb=25,
        find_unused_parameters=False,
        check_reduction=False,
        gradient_as_bucket_view=False,
        static_graph=False,
    ):
        super(DistributedDataParallel, self).__init__()

        # 1. 保存模块引用
        self.module = module

        # 2. 获取ProcessGroup
        if process_group is None:
            self.process_group = _get_default_group()
        else:
            self.process_group = process_group

        # 3. 设置设备
        if device_ids is None:
            device_ids = [torch.cuda.current_device()]
        self.device_ids = device_ids
        self.output_device = output_device or device_ids[0]

        # 4. 同步模型参数 (从rank 0广播)
        self._sync_params_and_buffers(authoritative_rank=0)

        # 5. 构建Reducer
        self._build_reducer(
            bucket_cap_mb=bucket_cap_mb,
            find_unused_parameters=find_unused_parameters,
            gradient_as_bucket_view=gradient_as_bucket_view,
        )

        # 6. 注册梯度钩子
        self._register_grad_hooks()

    def _sync_params_and_buffers(self, authoritative_rank=0):
        """同步所有rank的模型参数"""
        # 收集所有参数和buffer
        module_states = []
        for param in self.module.parameters():
            module_states.append(param.detach())
        for buffer in self.module.buffers():
            module_states.append(buffer.detach())

        # 从authoritative_rank广播到所有rank
        for state in module_states:
            dist.broadcast(state, src=authoritative_rank,
                          group=self.process_group)

    def _build_reducer(self, bucket_cap_mb, find_unused_parameters,
                       gradient_as_bucket_view):
        """构建Reducer用于梯度同步"""

        # 收集需要同步的参数
        parameters = list(self.module.parameters())
        params_for_reduction = [p for p in parameters if p.requires_grad]

        # 构建参数桶
        # 将参数按大小分组,减少通信次数
        self.reducer = dist.Reducer(
            params_for_reduction,
            list(range(len(params_for_reduction))),  # 参数索引
            self.process_group,
            [[]],  # 预期稀疏梯度
            bucket_cap_mb * 1024 * 1024,  # 桶大小 (bytes)
            find_unused_parameters,
            gradient_as_bucket_view,
            self.module.parameters,  # 参数名称
        )

    def _register_grad_hooks(self):
        """为每个参数注册梯度钩子"""
        for param in self.module.parameters():
            if param.requires_grad:
                # 注册backward钩子
                param.register_hook(self._make_param_hook(param))

    def _make_param_hook(self, param):
        """创建参数钩子"""
        def hook(grad):
            # 通知Reducer该参数的梯度已ready
            self.reducer.mark_variable_ready(param)
            return grad
        return hook

# torch/nn/parallel/distributed.py (续)

class DistributedDataParallel(Module):

    def forward(self, *inputs, **kwargs):
        """
        DDP forward流程:
        1. 同步模型参数 (如果需要)
        2. 准备Reducer
        3. 执行模型forward
        """
        # 1. 如果配置了broadcast_buffers,同步buffer
        if self.broadcast_buffers:
            self._sync_buffers()

        # 2. 准备Reducer (重置状态)
        if self.require_backward_grad_sync:
            self.reducer.prepare_for_backward([])

        # 3. 执行forward
        if self.device_ids:
            # 移动输入到GPU
            inputs = self._move_inputs_to_device(inputs, self.device_ids[0])
            kwargs = self._move_kwargs_to_device(kwargs, self.device_ids[0])

        # 4. 调用原始模块的forward
        output = self.module(*inputs, **kwargs)

        # 5. 如果是第一次迭代,记录计算图
        if self._is_first_iteration():
            self.reducer.prepare_for_backward(
                list(_find_tensors(output))
            )

        return output

    def _sync_buffers(self):
        """同步模型buffer (如BatchNorm的running_mean)"""
        buffers = [b.data for b in self.module.buffers()]
        if buffers:
            # 从rank 0广播
            dist.broadcast_coalesced(buffers, group=self.process_group)

# torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py

def allreduce_hook(
    process_group: dist.ProcessGroup,
    bucket: dist.GradBucket
) -> torch.futures.Future[torch.Tensor]:
    """
    默认AllReduce钩子

    Args:
        process_group: 通信组
        bucket: 梯度桶 (包含展平后的梯度tensor)

    Returns:
        Future对象, 完成后包含归约后的梯度
    """
    return (
        dist.all_reduce(bucket.buffer(), group=process_group, async_op=True)
        .get_future()
        .then(lambda fut: fut.value()[0])
    )


def fp16_compress_hook(
    process_group: dist.ProcessGroup,
    bucket: dist.GradBucket
) -> torch.futures.Future[torch.Tensor]:
    """
    FP16压缩钩子

    将梯度转换为FP16传输, 减少通信量50%
    """
    # 转换为FP16
    compressed = bucket.buffer().half()

    # AllReduce
    future = dist.all_reduce(compressed, group=process_group, async_op=True).get_future()

    # 转回FP32
    def decompress(fut):
        return fut.value()[0].float()

    return future.then(decompress)


def powerSGD_hook(
    process_group: dist.ProcessGroup,
    bucket: dist.GradBucket,
    state: PowerSGDState
) -> torch.futures.Future[torch.Tensor]:
    """
    PowerSGD压缩钩子

    使用低秩分解压缩梯度
    """
    input_tensor = bucket.buffer()
    device = input_tensor.device

    # 1. 重塑为矩阵
    matrix = input_tensor.view(state.matrix_approximation_rank, -1)

    # 2. 正交化
    if state.use_error_feedback:
        matrix.add_(state.error_dict[bucket.index()])

    # 3. 低秩分解
    p, q = torch.qr(matrix)
    # 只传输p和q, 大大减少通信量

    # 4. AllReduce p和q
    dist.all_reduce(p, group=process_group)
    dist.all_reduce(q, group=process_group)

    # 5. 重建梯度
    reconstructed = torch.mm(p, q.t())

    # 6. 误差反馈
    if state.use_error_feedback:
        state.error_dict[bucket.index()] = matrix - reconstructed

    return torch.futures.Future().set_result(reconstructed.view_as(input_tensor))


# 使用通信钩子
model = DDP(model)
model.register_comm_hook(process_group, fp16_compress_hook)

# torch/distributed/fsdp/fully_sharded_data_parallel.py

class FullyShardedDataParallel(nn.Module):
    """
    全分片数据并行

    将模型参数、梯度、优化器状态分片到多个GPU
    """

    def __init__(
        self,
        module: nn.Module,
        process_group: ProcessGroup = None,
        sharding_strategy: ShardingStrategy = ShardingStrategy.FULL_SHARD,
        cpu_offload: CPUOffload = None,
        auto_wrap_policy: Optional[Callable] = None,
        backward_prefetch: BackwardPrefetch = BackwardPrefetch.BACKWARD_PRE,
        mixed_precision: Optional[MixedPrecision] = None,
        ignored_modules: Optional[Iterable[nn.Module]] = None,
        param_init_fn: Optional[Callable] = None,
        sync_module_states: bool = False,
    ):
        super().__init__()

        self.process_group = process_group or dist.distributed_c10d._get_default_group()
        self.rank = dist.get_rank(self.process_group)
        self.world_size = dist.get_world_size(self.process_group)

        # 分片策略
        self.sharding_strategy = sharding_strategy

        # CPU卸载
        self.cpu_offload = cpu_offload

        # 混合精度
        self.mixed_precision = mixed_precision

        # 1. 自动wrap子模块 (可选)
        if auto_wrap_policy is not None:
            module = self._auto_wrap(module, auto_wrap_policy)

        self._fsdp_wrapped_module = module

        # 2. 收集并分片参数
        self._shard_parameters()

        # 3. 注册forward/backward钩子
        self._register_hooks()

    def _shard_parameters(self):
        """分片模型参数"""
        self._handles = []

        for param_name, param in self._fsdp_wrapped_module.named_parameters():
            # 1. 创建FlatParameter (展平参数)
            flat_param = FlatParameter(param)

            # 2. 计算每个rank的分片大小
            param_numel = flat_param.numel()
            chunk_size = (param_numel + self.world_size - 1) // self.world_size

            # 3. 只保留本rank的分片
            start_idx = self.rank * chunk_size
            end_idx = min(start_idx + chunk_size, param_numel)

            # 4. 创建分片视图
            sharded_param = flat_param.data[start_idx:end_idx].clone()

            # 5. 替换原参数
            setattr(self._fsdp_wrapped_module, param_name,
                   nn.Parameter(sharded_param))

            # 6. 创建Handle管理分片
            handle = FlatParamHandle(
                flat_param=flat_param,
                sharded_param=sharded_param,
                param_name=param_name,
            )
            self._handles.append(handle)

    def _unshard_params(self, handles=None):
        """取消分片, 收集完整参数"""
        if handles is None:
            handles = self._handles

        for handle in handles:
            # AllGather收集所有分片
            full_param = torch.empty(
                handle.flat_param.shape,
                dtype=handle.sharded_param.dtype,
                device=handle.sharded_param.device,
            )

            # 收集所有rank的分片
            dist.all_gather_into_tensor(
                full_param,
                handle.sharded_param,
                group=self.process_group,
            )

            # 更新模型参数为完整参数
            handle.flat_param.data = full_param

    def _reshard_params(self, handles=None):
        """重新分片, 释放非本地参数"""
        if handles is None:
            handles = self._handles

        for handle in handles:
            # 只保留本地分片
            handle.flat_param.data = handle.sharded_param

    def forward(self, *args, **kwargs):
        """
        FSDP forward流程:
        1. AllGather收集完整参数
        2. 执行forward
        3. 释放参数 (重新分片)
        """
        # 1. Unshard参数
        self._unshard_params()

        # 2. Forward
        output = self._fsdp_wrapped_module(*args, **kwargs)

        # 3. Reshard参数 (如果不需要保留用于backward)
        if self.sharding_strategy == ShardingStrategy.FULL_SHARD:
            self._reshard_params()

        return output

# torch/distributed/fsdp/fully_sharded_data_parallel.py (续)

class FullyShardedDataParallel(nn.Module):

    def _register_hooks(self):
        """注册forward/backward钩子"""

        def _pre_forward_hook(module, args):
            """Forward前: Unshard参数"""
            if module._is_root:
                self._unshard_params()

        def _post_forward_hook(module, args, output):
            """Forward后: Reshard参数"""
            if self.sharding_strategy == ShardingStrategy.FULL_SHARD:
                self._reshard_params()

        def _pre_backward_hook(module, grad_output):
            """Backward前: Unshard参数 (如果已reshard)"""
            self._unshard_params()

            # 预取下一层的参数 (优化)
            if self.backward_prefetch == BackwardPrefetch.BACKWARD_PRE:
                self._prefetch_next_layer()

        def _post_backward_hook(param, grad):
            """
            Backward后: ReduceScatter梯度

            不是AllReduce! 每个rank只需要自己分片的梯度
            """
            # 1. ReduceScatter: 每个rank得到梯度的不同分片
            grad_sharded = torch.empty(
                self._chunk_size,
                dtype=grad.dtype,
                device=grad.device,
            )

            dist.reduce_scatter_tensor(
                grad_sharded,
                grad,
                group=self.process_group,
            )

            # 2. 更新本地梯度
            param.grad = grad_sharded

            # 3. Reshard参数
            self._reshard_params()

        # 注册钩子
        self.register_forward_pre_hook(_pre_forward_hook)
        self.register_forward_hook(_post_forward_hook)
        # backward钩子通过autograd hooks注册


class FlatParamHandle:
    """
    管理分片参数的Handle
    """

    def __init__(self, flat_param, sharded_param, param_name):
        self.flat_param = flat_param      # 原始完整参数
        self.sharded_param = sharded_param  # 本rank的分片
        self.param_name = param_name

        # 梯度累积
        self.grad_accumulator = None

        # 通信handle (用于异步)
        self._all_gather_work = None
        self._reduce_scatter_work = None

    def all_gather_params(self, process_group, async_op=True):
        """异步AllGather参数"""
        full_param = torch.empty(
            self.flat_param.shape,
            dtype=self.sharded_param.dtype,
            device=self.sharded_param.device,
        )

        work = dist.all_gather_into_tensor(
            full_param,
            self.sharded_param,
            group=process_group,
            async_op=async_op,
        )

        if async_op:
            self._all_gather_work = work
            self._full_param_buffer = full_param
        else:
            self.flat_param.data = full_param

    def wait_all_gather(self):
        """等待AllGather完成"""
        if self._all_gather_work is not None:
            self._all_gather_work.wait()
            self.flat_param.data = self._full_param_buffer
            self._all_gather_work = None

    def reduce_scatter_grad(self, process_group, async_op=True):
        """ReduceScatter梯度"""
        grad_sharded = torch.empty(
            self.sharded_param.shape,
            dtype=self.flat_param.grad.dtype,
            device=self.flat_param.grad.device,
        )

        work = dist.reduce_scatter_tensor(
            grad_sharded,
            self.flat_param.grad,
            group=process_group,
            async_op=async_op,
        )

        if async_op:
            self._reduce_scatter_work = work
            self._grad_sharded_buffer = grad_sharded
        else:
            self.sharded_param.grad = grad_sharded

    def wait_reduce_scatter(self):
        """等待ReduceScatter完成"""
        if self._reduce_scatter_work is not None:
            self._reduce_scatter_work.wait()
            self.sharded_param.grad = self._grad_sharded_buffer
            self._reduce_scatter_work = None

import os
import logging

# 启用PyTorch分布式日志
os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"

# 或通过代码设置
import torch.distributed as dist
dist.set_debug_level(dist.DebugLevel.DETAIL)

# NCCL调试
os.environ["NCCL_DEBUG"] = "INFO"
os.environ["NCCL_DEBUG_SUBSYS"] = "ALL"

# 查看DDP内部状态
model = DDP(model)
print(model._module_copies)
print(model.reducer.buckets)

# 检查通信状态
def check_dist_state():
    print(f"Rank: {dist.get_rank()}")
    print(f"World Size: {dist.get_world_size()}")
    print(f"Backend: {dist.get_backend()}")
    print(f"Is Initialized: {dist.is_initialized()}")

# 检查DDP状态
def check_ddp_state(model):
    print(f"Device IDs: {model.device_ids}")
    print(f"Output Device: {model.output_device}")
    print(f"Bucket Cap: {model.bucket_bytes_cap}")
    print(f"Find Unused Params: {model.find_unused_parameters}")

# 检查FSDP状态
def check_fsdp_state(model):
    print(f"Sharding Strategy: {model.sharding_strategy}")
    print(f"CPU Offload: {model.cpu_offload}")
    print(f"Mixed Precision: {model.mixed_precision}")

# 同步检查点
def sync_check(msg):
    rank = dist.get_rank()
    print(f"[Rank {rank}] {msg}")
    dist.barrier()