PyTorch FSDP（完全分片数据并行）的内存优化：参数、梯度与优化器状态的分片策略

PyTorch FSDP 中的内存优化：参数、梯度与优化器状态的分片策略

大家好！今天我们来深入探讨 PyTorch 中 FSDP（Fully Sharded Data Parallel，完全分片数据并行）的内存优化策略。在大规模深度学习模型的训练中，内存瓶颈是一个常见的问题。FSDP 旨在通过将模型参数、梯度和优化器状态分片到不同的 GPU 设备上，从而显著降低每个设备的内存占用，实现更大模型的训练。 本次讲座将围绕以下几个方面展开：

1. FSDP 的基本原理与优势： 简单回顾 FSDP 的核心思想，强调其在内存优化方面的作用。
2. 参数分片策略： 详细讲解不同的参数分片策略，包括 FULL_SHARD 和 SHARD_GRAD_OP，以及它们对内存和通信的影响。
3. 梯度分片策略： 深入分析梯度累积和梯度通信的机制，以及如何通过梯度分片进一步优化内存使用。
4. 优化器状态分片策略： 讨论如何将优化器状态进行分片，以减少每个设备的内存负担。
5. 混合精度训练与 FSDP： 结合混合精度训练（AMP）技术，进一步降低内存占用，提高训练效率。
6. 代码示例与实践： 通过具体的代码示例，演示如何在 PyTorch 中使用 FSDP 进行内存优化。
7. 性能评估与调优： 提供一些性能评估和调优的建议，帮助大家更好地使用 FSDP。

1. FSDP 的基本原理与优势

FSDP 是一种数据并行训练策略，它将模型参数、梯度和优化器状态分片到多个 GPU 设备上，从而减少每个设备的内存占用。相比于传统的 DataParallel，FSDP 的优势在于：

• 更大的模型容量： 由于参数和优化器状态被分片，每个设备只需存储部分模型，可以训练更大的模型。
• 更好的内存效率： 通过分片，显著降低了每个设备的内存占用，避免了 OOM（Out of Memory）错误。
• 可扩展性： FSDP 可以很好地扩展到更多的 GPU 设备，实现更快的训练速度。

简单来说，FSDP 的核心思想可以概括为：分而治之，将大模型分解成小块，分散到不同的设备上进行计算。

2. 参数分片策略

FSDP 提供了多种参数分片策略，最常用的两种是 FULL_SHARD 和 SHARD_GRAD_OP。

• FULL_SHARD： 将模型参数完全分片到各个 GPU 上。每个 GPU 只存储一部分参数。在需要计算时，通过 all-gather 操作收集完整的参数。计算完成后，梯度会被分片，并用于更新本地参数。

  • 优点： 内存占用最小，可以训练最大的模型。
  • 缺点： 通信开销较大，需要频繁的 all-gather 操作。

• SHARD_GRAD_OP： 该策略仅对梯度进行分片，保持模型参数完整存储在每个设备上，并在反向传播过程中对梯度进行分片和归约。

  • 优点： 通信开销相对较小，因为不需要 all-gather 参数。
  • 缺点： 内存占用相对较大，因为每个设备需要存储完整的模型参数。

以下是一个使用 FULL_SHARD 策略的示例：

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.sharding_strategy import ShardingStrategy

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型
model = SimpleModel(input_size=10, hidden_size=20, output_size=5)

# 使用 FSDP 进行封装，并指定 FULL_SHARD 策略
fsdp_model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

# 示例输入
input_tensor = torch.randn(32, 10).cuda() # 假设使用 CUDA
output = fsdp_model(input_tensor)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(fsdp_model.parameters(), lr=0.001)

# 示例训练循环
for i in range(10):
    optimizer.zero_grad()
    output = fsdp_model(input_tensor)
    target = torch.randint(0, 5, (32,)).cuda() # 假设使用 CUDA
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(f"Iteration {i}: Loss = {loss.item()}")

在这个示例中，我们使用 ShardingStrategy.FULL_SHARD 指定了 FULL_SHARD 策略。这意味着模型的参数将被完全分片到各个 GPU 上。

以下是一个使用 SHARD_GRAD_OP 策略的示例：

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.sharding_strategy import ShardingStrategy

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型
model = SimpleModel(input_size=10, hidden_size=20, output_size=5)

# 使用 FSDP 进行封装，并指定 SHARD_GRAD_OP 策略
fsdp_model = FSDP(model, sharding_strategy=ShardingStrategy.SHARD_GRAD_OP)

# 示例输入
input_tensor = torch.randn(32, 10).cuda() # 假设使用 CUDA
output = fsdp_model(input_tensor)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(fsdp_model.parameters(), lr=0.001)

# 示例训练循环
for i in range(10):
    optimizer.zero_grad()
    output = fsdp_model(input_tensor)
    target = torch.randint(0, 5, (32,)).cuda() # 假设使用 CUDA
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(f"Iteration {i}: Loss = {loss.item()}")

在这个示例中，我们使用 ShardingStrategy.SHARD_GRAD_OP 指定了 SHARD_GRAD_OP 策略。这意味着模型的参数将完整存储在每个 GPU 上，而梯度将被分片。

选择哪种策略取决于具体的模型大小、硬件环境和通信带宽。通常情况下，对于非常大的模型，FULL_SHARD 策略是更好的选择。

3. 梯度分片策略

梯度分片是 FSDP 中另一个重要的内存优化技术。在反向传播过程中，每个 GPU 计算出的梯度会被分片，并用于更新本地的参数。

梯度累积（Gradient Accumulation）是一种常用的技巧，可以在有限的 GPU 资源下，模拟更大的 batch size。在 FSDP 中，梯度累积可以与梯度分片结合使用，进一步降低内存占用。

以下是一个使用梯度累积的示例：

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.sharding_strategy import ShardingStrategy

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型
model = SimpleModel(input_size=10, hidden_size=20, output_size=5)

# 使用 FSDP 进行封装，并指定 FULL_SHARD 策略
fsdp_model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(fsdp_model.parameters(), lr=0.001)

# 梯度累积的步数
accumulation_steps = 4

# 示例训练循环
for i in range(10):
    optimizer.zero_grad()
    for j in range(accumulation_steps):
        input_tensor = torch.randn(32, 10).cuda() # 假设使用 CUDA
        output = fsdp_model(input_tensor)
        target = torch.randint(0, 5, (32,)).cuda() # 假设使用 CUDA
        loss = criterion(output, target)
        loss = loss / accumulation_steps  # 归一化损失
        loss.backward()

    optimizer.step()
    print(f"Iteration {i}: Loss = {loss.item()}")

在这个示例中，我们将 batch size 设置为 32，并将 accumulation_steps 设置为 4。这意味着我们将在 4 个小 batch 上累积梯度，然后再进行一次优化器更新。这相当于使用了一个 batch size 为 128 的训练。

梯度通信是 FSDP 中另一个重要的方面。在反向传播过程中，每个 GPU 计算出的梯度需要进行 all-reduce 操作，以确保所有 GPU 上的梯度一致。FSDP 提供了多种梯度通信策略，可以根据具体的硬件环境和模型结构进行选择。

4. 优化器状态分片策略

优化器状态（Optimizer State）是优化器在更新模型参数时需要维护的一些信息，例如动量（Momentum）和方差（Variance）。对于大型模型，优化器状态可能会占用大量的内存。

FSDP 允许将优化器状态进行分片，以减少每个设备的内存负担。这意味着每个 GPU 只需存储一部分优化器状态。

以下是一个使用优化器状态分片的示例：

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.sharding_strategy import ShardingStrategy
from torch.distributed.optim import FSDPAdam

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型
model = SimpleModel(input_size=10, hidden_size=20, output_size=5)

# 使用 FSDP 进行封装，并指定 FULL_SHARD 策略
fsdp_model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

# 使用 FSDPAdam 优化器，它会自动处理优化器状态的分片
optimizer = FSDPAdam(fsdp_model.parameters(), lr=0.001)

# 损失函数
criterion = nn.CrossEntropyLoss()

# 示例训练循环
for i in range(10):
    optimizer.zero_grad()
    input_tensor = torch.randn(32, 10).cuda() # 假设使用 CUDA
    output = fsdp_model(input_tensor)
    target = torch.randint(0, 5, (32,)).cuda() # 假设使用 CUDA
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(f"Iteration {i}: Loss = {loss.item()}")

在这个示例中，我们使用了 FSDPAdam 优化器，它是 PyTorch 提供的专门用于 FSDP 的优化器。FSDPAdam 会自动处理优化器状态的分片，无需手动进行配置。

5. 混合精度训练与 FSDP

混合精度训练（Automatic Mixed Precision，AMP）是一种常用的加速训练和降低内存占用的技术。它通过使用半精度浮点数（FP16）来存储模型参数和激活值，从而减少内存占用，并提高计算速度。

FSDP 可以与 AMP 结合使用，进一步降低内存占用，提高训练效率。

以下是一个使用 AMP 和 FSDP 的示例：

import torch
import torch.nn as nn
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.sharding_strategy import ShardingStrategy
from torch.cuda.amp import GradScaler, autocast

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 初始化模型
model = SimpleModel(input_size=10, hidden_size=20, output_size=5)

# 使用 FSDP 进行封装，并指定 FULL_SHARD 策略
fsdp_model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(fsdp_model.parameters(), lr=0.001)

# 初始化 GradScaler
scaler = GradScaler()

# 示例训练循环
for i in range(10):
    optimizer.zero_grad()
    with autocast():
        input_tensor = torch.randn(32, 10).cuda() # 假设使用 CUDA
        output = fsdp_model(input_tensor)
        target = torch.randint(0, 5, (32,)).cuda() # 假设使用 CUDA
        loss = criterion(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    print(f"Iteration {i}: Loss = {loss.item()}")

在这个示例中，我们使用了 torch.cuda.amp.GradScaler 和 torch.cuda.amp.autocast 来启用 AMP。GradScaler 用于缩放损失，以避免梯度下溢。autocast 用于自动将模型和输入的数据类型转换为半精度浮点数。

6. 代码示例与实践

在实际应用中，使用 FSDP 需要考虑以下几个方面：

• 初始化分布式环境： 在使用 FSDP 之前，需要初始化 PyTorch 的分布式环境。
• 封装模型： 使用 FullyShardedDataParallel 类将模型封装起来。
• 选择合适的分片策略： 根据模型大小、硬件环境和通信带宽选择合适的分片策略。
• 使用 FSDP 兼容的优化器： 建议使用 FSDPAdam 等 FSDP 兼容的优化器。
• 调整 batch size 和梯度累积步数： 根据 GPU 内存大小调整 batch size 和梯度累积步数。
• 启用 AMP： 如果条件允许，可以启用 AMP，进一步降低内存占用，提高训练效率。

以下是一个更完整的示例，演示了如何在多个 GPU 上使用 FSDP 进行训练：

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.sharding_strategy import ShardingStrategy
from torch.distributed.optim import FSDPAdam
from torch.cuda.amp import GradScaler, autocast
import os

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return 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):
    setup(rank, world_size)
    torch.cuda.set_device(rank)

    # 初始化模型
    model = SimpleModel(input_size=10, hidden_size=20, output_size=5).to(rank) # 将模型移动到当前设备

    # 使用 FSDP 进行封装，并指定 FULL_SHARD 策略
    fsdp_model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

    # 损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = FSDPAdam(fsdp_model.parameters(), lr=0.001)

    # 初始化 GradScaler
    scaler = GradScaler()

    # 示例训练循环
    for i in range(10):
        optimizer.zero_grad()
        with autocast():
            input_tensor = torch.randn(32, 10).to(rank) # 将输入移动到当前设备
            output = fsdp_model(input_tensor)
            target = torch.randint(0, 5, (32,)).to(rank) # 将目标移动到当前设备
            loss = criterion(output, target)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        if rank == 0:
            print(f"Iteration {i}: Loss = {loss.item()}")

    cleanup()

if __name__ == "__main__":
    world_size = torch.cuda.device_count() # 获取可用 GPU 数量
    torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size, join=True)

这个示例演示了如何在多个 GPU 上使用 FSDP 进行训练。需要注意的是，在使用 FSDP 之前，需要初始化 PyTorch 的分布式环境。

7. 性能评估与调优

在使用 FSDP 时，需要进行性能评估和调优，以确保获得最佳的训练效率。以下是一些建议：

• 选择合适的分片策略： 不同的分片策略对内存占用和通信开销有不同的影响。需要根据具体的模型大小、硬件环境和通信带宽选择合适的分片策略。
• 调整 batch size 和梯度累积步数： 调整 batch size 和梯度累积步数可以平衡内存占用和训练速度。
• 启用 AMP： 如果条件允许，可以启用 AMP，进一步降低内存占用，提高训练效率。
• 使用性能分析工具： 可以使用 PyTorch Profiler 等性能分析工具来分析训练过程中的瓶颈，并进行相应的优化。
• 监控 GPU 内存使用情况： 可以使用 torch.cuda.memory_allocated() 和 torch.cuda.max_memory_allocated() 等函数来监控 GPU 内存使用情况，以便及时发现 OOM 错误。

通过合理的配置和调优，可以充分发挥 FSDP 的优势，实现更大规模模型的训练。

FSDP 的核心优势

FSDP 通过参数、梯度和优化器状态的分片，显著降低了每个设备的内存占用，实现了更大模型的训练。 结合梯度累积、混合精度训练等技术，可以进一步优化内存使用和训练效率。 通过仔细的性能评估和调优，可以充分发挥 FSDP 的优势。

希望今天的讲座对大家有所帮助！

更多IT精英技术系列讲座，到智猿学院