1F1B 交错式调度模式

1F1B 交错式调度模式

1F1B 交错式调度模式是一种更高效的流水线并行策略，它通过交错执行前向和后向计算，提高了资源利用率和计算效率。这种模式要求 micro-batch 的数量是流水线阶段的整数倍。每个设备不再仅负责连续多个层的计算，而是可以处理多个层的子集，这些子集被称为模型块。具体来说：

1. 模型块分配（Model Block Allocation）

在传统的1F1B非交错式调度模式中，设备1可能负责层1-4，设备2负责层5-8，以此类推。然而，在1F1B交错式调度模式下，设备1可以处理层1、2、9、10，设备2处理层3、4、11、12，以此类推。这种分配方式使得每个设备在流水线中被分配到多个阶段。

假设我们有4个设备（设备0、设备1、设备2、设备3），模型被分成16个层，每个设备负责4个层的子集。具体分配如下：

• 设备0：负责层1、2、9、10
• 设备1：负责层3、4、11、12
• 设备2：负责层5、6、13、14
• 设备3：负责层7、8、15、16

2. 多阶段任务分配（Multi-Stage Task Allocation）

在1F1B交错式调度模式下，每个设备在流水线中参与多个阶段的任务。例如，设备1可能参与热身阶段、前向计算阶段和后向计算阶段的某些子集任务。具体步骤如下：

1. 热身阶段（Warm-up Phase）：
   • 每个设备依次执行前向计算，直到所有设备都完成一次前向计算。这个阶段的目的是为后续的前向-后向阶段做好准备，确保数据在设备之间正确传递。
   • 例如，设备0先执行层1、2的前向计算，将结果传递给设备1；设备1执行层3、4的前向计算，将结果传递给设备2；设备2执行层5、6的前向计算，将结果传递给设备3；设备3执行层7、8的前向计算，得到部分模型输出。
2. 前向-后向阶段（Forward-Backward Phase）：
   • 每个设备按顺序执行一次前向计算，然后进行一次后向计算。每个设备在完成前向计算后，会等待前一个设备的梯度信息，然后进行后向计算。
   • 例如，设备0执行层1、2的前向计算，将结果传递给设备1；设备1执行层3、4的前向计算，将结果传递给设备2；设备2执行层5、6的前向计算，将结果传递给设备3；设备3执行层7、8的前向计算，得到部分模型输出，并计算损失函数的梯度。
   • 然后，设备3执行层15、16的后向计算，将梯度传递给设备2；设备2执行层13、14的后向计算，将梯度传递给设备1；设备1执行层11、12的后向计算，将梯度传递给设备0；设备0执行层9、10的后向计算。
3. 后向阶段（Backward Phase）：
   • 每个设备完成最后一次后向计算。最后一个设备完成后向计算后，会将梯度信息传递给前一个设备，依次类推，直到所有设备都完成后向计算。
   • 例如，设备3完成层15、16的后向计算，将梯度传递给设备2；设备2完成层13、14的后向计算，将梯度传递给设备1；设备1完成层11、12的后向计算，将梯度传递给设备0；设备0完成层9、10的后向计算。

3. 并行执行（Parallel Execution）

在1F1B交错式调度模式下，每个设备可以并行执行不同阶段的计算任务，从而更好地利用流水线并行的优势。具体来说：

• 设备0：在执行层1、2的前向计算的同时，可以并行执行层9、10的后向计算。
• 设备1：在执行层3、4的前向计算的同时，可以并行执行层11、12的后向计算。
• 设备2：在执行层5、6的前向计算的同时，可以并行执行层13、14的后向计算。
• 设备3：在执行层7、8的前向计算的同时，可以并行执行层15、16的后向计算。

代码示例

以下是一个简化的代码示例，展示了如何实现1F1B交错式调度模式：

Python复制

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size, num_micro_batches, num_pipeline_stages):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    # 创建模型并封装为DDP模型
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # 定义损失函数和优化器
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # 创建数据集和数据加载器
    dataset = datasets.FakeData(size=1000, image_size=(1, 10), num_classes=5)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    # 1F1B 交错式调度模式
    for epoch in range(5):
        sampler.set_epoch(epoch)
        for micro_batch_idx in range(num_micro_batches):
            # 热身阶段
            if micro_batch_idx < num_pipeline_stages:
                # 前向计算
                data, target = next(iter(dataloader))
                output = ddp_model(data.to(rank))
                loss = loss_fn(output, target.to(rank))
                loss.backward()
            else:
                # 前向-后向阶段
                data, target = next(iter(dataloader))
                output = ddp_model(data.to(rank))
                loss = loss_fn(output, target.to(rank))
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

        # 后向阶段
        for _ in range(num_pipeline_stages):
            optimizer.step()
            optimizer.zero_grad()

        print(f"Rank {rank}, Epoch {epoch} finished")

    cleanup()

def main():
    world_size = torch.cuda.device_count()
    num_micro_batches = 10  # 示例中假设每个epoch有10个micro-batch
    num_pipeline_stages = 4  # 示例中假设有4个流水线阶段
    mp.spawn(train, args=(world_size, num_micro_batches, num_pipeline_stages), nprocs=world_size, join=True)

if __name__ == "__main__":
    main()

代码说明

1. 模型定义：ToyModel是一个简单的线性模型，包含两个全连接层和一个ReLU激活函数。
2. 训练函数：train函数是每个进程的主函数，负责模型的训练。它首先设置分布式环境和设备，然后创建模型并封装为DDP模型。接着定义损失函数和优化器，并创建数据集和数据加载器。
3. 1F1B 交错式调度模式：
   • 在每个epoch中，首先进行热身阶段，每个设备依次执行前向计算。
   • 接下来是前向-后向阶段，每个设备按顺序执行一次前向计算，然后进行一次后向计算。
   • 最后是后向阶段，每个设备完成最后一次后向计算。
4. 主函数：main函数使用mp.spawn启动多个进程，每个进程调用train函数进行训练。

总结

1F1B 交错式调度模式通过交错执行前向和后向计算，