An FP16 parameter ~ 2 bytes
An FP16 gradient ~ 2 bytes
An FP32 optimizer state ~ 8 bytes based on the Adam optimizers
An FP32 copy of parameter ~ 4 bytes (needed for the
optimizer apply(OA) operation)An FP32 copy of gradient ~ 4 bytes (needed for the OA operation)
所以如果 10 Billion 参数的模型, 至少需要 200GB GPU内存, 目前 NVIDIA A100 最大也就 40GB/80GB GPU，而目前很多模型都是100B，所以训练的过程中分片是必须的，否则根本带不起来。
01
—
什么是 Sharding？
在大型语言模型（LLM）中，分片（sharding）是指将训练数据或者模型参数、激活值或梯度等数据分布到多个设备（如 GPU 或 TPU）上，以优化内存使用和计算效率。
Sharding就相当于把上图中的拼图分拆成为很多块，放在不同的位置，但是它们彼此可以联系，对外面的人来说它们仍然是一块完整的拼图。—
常见的Sharding方法：之一数据并行
LLM 中分片优化的关键技术和方法：
1. 数据并行（Data Parallelism）
数据并行是最常见的分片方式，主要通过在不同设备上训练不同的小批量数据（mini-batches）来并行化训练。
1）原理
每个设备拥有完整的模型副本。 输入数据被分成多个小批量（mini-batches），分别分配到不同设备上。 每个设备独立计算前向传播和反向传播，得到梯度。 梯度通过设备之间的 All-Reduce 操作进行同步，更新模型参数。
DP (Data Parallel)：这是最基本的数据并行策略，将模型的副本分布到单机多卡上，每个卡都有一个模型副本，然后每个卡都会处理不同的数据子集。在每个训练步骤结束时，所有卡都会同步模型参数。 DDP (Distributed Data Parallel)：随着数据量的增大，单机多卡的训练效率会变得很低，这时候就需要使用分布式数据并行DDP。DDP是将模型的副本分布到多台机器上，每台机器上有多个卡，每个卡都有一个模型副本。在每个训练步骤结束时，所有卡都会同步模型参数，相当于多机多卡。
All-Reduce操作：数据并行通常涉及到All-Reduce操作，这是一种在所有计算节点之间同步梯度的技术，以确保模型的一致性。
重叠计算和通信（Overlap Communication and Computation）：在设备间同步梯度时，同时进行下一步的前向传播计算，减少通信等待时间。
def setup_distributed(backend: str = 'nccl',# nvidia的GPU通信模块。world_size: int = 1,rank: int = 0,master_addr: str = 'localhost',master_port: str = '12355') -> None:"""初始化分布式环境"""try:os.environ['MASTER_ADDR'] = master_addros.environ['MASTER_PORT'] = master_portdist.init_process_group(backend, rank=rank, world_size=world_size)logger.info(f"Initialized process group: rank {rank}/{world_size}")except Exception as e:logger.error(f"Failed to initialize distributed environment: {e}")raise
# 创建数据集和数据加载器# 每个GPU上面都会保留一份完整的数据集。dataset = RealDataset(length=10000, features=10)# 根据Rank的值来决定自己取到的是数据集的哪一个部分。sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)dataloader = DataLoader(dataset,batch_size=batch_size,sampler=sampler,num_workers=4,#用于加载数据pin_memory=True,prefetch_factor=2,persistent_workers=True)
# 创建模型和优化器model = SimpleModel().to(device)model = DDP(model,device_ids=[rank],output_device=rank,gradient_as_bucket_view=True,find_unused_parameters=False)optimizer = optim.AdamW(model.parameters(), lr=learning_rate)scheduler = StepLR(optimizer, step_size=2, gamma=0.1)
def setup_distributed(local_rank: int,node_rank: int,world_size: int,master_addr: str,master_port: str):# 计算全局 rankglobal_rank = node_rank * torch.cuda.device_count() + local_rankos.environ['MASTER_ADDR'] = str(master_addr)os.environ['MASTER_PORT'] = str(master_port)dist.init_process_group('nccl',rank=global_rank,world_size=world_size)return global_rank
mp.spawn(main_worker,nprocs=gpus_per_node,args=(args.node_rank,world_size,args.master_addr,args.master_port,))
03
—
数据并行优势和劣势
数据并行（Data Parallelism）在LLM（Large Language Models，大型语言模型）训练中具有以下优势和劣势：
优势：
加速训练过程：数据并行通过在多个设备上并行处理不同的数据批次来加快训练速度。
支持更大的批次大小：由于数据被分割到多个GPU上，可以处理比单个GPU内存更大的数据批次
简化实现：数据并行相对容易实现，因为它不需要对模型结构进行修改，只需复制模型到每个计算单元。
通信开销：数据并行需要在每次迭代后同步梯度，这可能导致大量的GPU间通信，特别是在大规模集群中
扩展性问题：在大模型和大规模集群下，数据并行的扩展性较差，因为每个GPU都需要存储整个模型的副本，这限制了模型大小和可扩展性。
内存限制：每个GPU都必须有足够的内存来存储整个模型，对于非常大的模型，这可能是一个限制因素 梯度过时问题：在异步数据并行中，由于计算梯度需要时间，当某个节点完成梯度计算并将其与全局参数向量合并时，全局参数可能已经被更新多次，导致梯度过时问题
进技术交流群请添加AINLP小助手微信（id: ainlp2)
请备注具体方向+所用到的相关技术点 关于AINLP
AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括LLM、预训练模型、自动生成、文本摘要、智能问答、聊天机器人、机器翻译、知识图谱、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP小助手微信(id：ainlp2)，备注工作/研究方向+加群目的。
