作者丨东尼大佬来源丨https://zhuanlan.zhihu.com/p/673671771编辑丨GiantPandaCV1. 能用_all_gather_base的，不用all_gatheroutput = torch.empty(input.numel() * world_size, dtype=input.dtype, device=input.device)torch.distributed._all_gather_base(output, input, group=xxx)vs.output_list = [ torch.empty(input.numel(), dtype=input.dtype, device=input.device) for _ in range(world_size)]torch.distributed.all_gather(output_list, input, group=xxx)output = torch.cat(output_list, dim=0)内存碎片更少，操作更少，性能/内存均有收益！2. 能用专有算子的，不用通用算子如 F.embedding vs. Index-selectMegatron-LM master实现使用的Index-select算子，Index-select会涉及索引展开、内存复用等HostCPU逻辑，效率较低3. 对于生命周期较长的Tensors，可以共用contiguous bufferdata = torch.zeros(global_size, dtype=xx, device=xx)start_idx = 0for i in range(len(item_list)): item_list[i] = data[start_idx:start_idx+item_list[i].nume
………………………………
