ChatGPT内部使用PyTorch还是TensorFlow？

许多人好奇，像ChatGPT这样的大型语言模型（LLM）底层究竟使用了PyTorch还是TensorFlow？答案是：虽然两者都极其优秀，但在大型生成式AI（尤其是OpenAI/Meta/Hugging Face生态）领域，PyTorch占据了绝对的主导地位。这得益于其动态计算图、强大的社区支持，以及更易于集成的分布式训练工具（如DeepSpeed和原生FSDP）。

对于训练一个数十亿甚至千亿参数的模型，如GPT系列，最大的挑战不是计算力本身，而是单卡GPU内存的限制。为了解决这一问题，PyTorch生态系统提供了核心工具：FSDP (Fully Sharded Data Parallel)。

核心技术聚焦：PyTorch FSDP

传统的DDP（Distributed Data Parallel）只复制参数，但在每个GPU上仍需要存储完整的参数、梯度和优化器状态。对于参数量巨大的LLM，这很快就会耗尽A100或H100的显存。

FSDP通过以下方式彻底解决了这个问题：

1. 完全分片 (Fully Sharded)：FSDP将模型的参数（P）、梯度（G）和优化器状态（O）在所有可用的GPU上进行分片存储。
2. 动态传输：在特定层需要进行前向或后向传播时，相关的参数分片会被动态地聚集（All-Gather）到当前GPU上进行计算，计算完成后立即释放内存（又称零冗余优化器——Zero Redundancy Optimizer, Stage 3）。

这是构建类似ChatGPT基础设施的关键一步，因为它允许我们在可用的GPU资源上训练远超单卡容量的模型。

实战演练：使用PyTorch FSDP进行分布式训练

以下代码展示了如何初始化FSDP环境，并对一个简单的Transformer模型进行FSDP封装和训练。

环境准备：

确保安装了PyTorch和torchrun工具：

pip install torch

代码示例：fsdp_llm_train.py****

import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import torch.nn as nn

# 1. 初始化分布式环境
def setup():
    # 使用nccl后端进行GPU通信
    dist.init_process_group(backend="nccl")
    # 设置当前进程对应的GPU ID
    torch.cuda.set_device(dist.get_rank())

# 2. 简单的Transformer块 (模拟LLM组件)
class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim)
        )

    def forward(self, x):
        res = x
        x = self.norm1(x)
        attn_output, _ = self.attn(x, x, x)
        x = res + attn_output

        res = x
        x = self.norm2(x)
        x = self.mlp(x)
        return res + x

# 3. 构建模型并使用FSDP封装
def setup_model(local_rank):
    # 假设我们有一个由12个Transformer块组成的模型
    model = nn.Sequential(*[
        TransformerBlock(embed_dim=1024, num_heads=16) 
        for _ in range(12)
    ]).to(local_rank)

    # 定义FSDP自动封装策略：对TransformerBlock实例进行分片
    transformer_wrap_policy = transformer_auto_wrap_policy(
        TransformerBlock
    )

    # 封装模型
    fsdp_model = FSDP(model, 
                      auto_wrap_policy=transformer_wrap_policy,
                      sharding_strategy=dist.fsdp.ShardingStrategy.FULL_SHARD)

    return fsdp_model

# 4. 训练循环
def train(rank, world_size):
    setup()

    model = setup_model(rank)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
    loss_fn = nn.MSELoss()

    print(f"Rank {rank}: Model initialized with FSDP.")

    # 模拟数据输入: Batch size 4, 序列长度 512, Embedding维度 1024
    data = torch.randn(4, 512, 1024, device=rank)
    target = torch.randn(4, 512, 1024, device=rank)

    # 训练步骤
    model.train()
    for step in range(3):
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)

        # 后向传播：FSDP会自动处理梯度的All-Reduce和参数分片
        loss.backward()

        optimizer.step()
        print(f"Rank {rank}, Step {step}: Loss = {loss.item():.4f}")

    # 清理分布式环境
    dist.destroy_process_group()

if __name__ == "__main__":
    # 假设使用两个GPU进行训练
    world_size = torch.cuda.device_count()
    if world_size < 1:
        print("需要至少1个CUDA设备")
    else:
        # FSDP通常通过torchrun启动
        # 实际运行时，torchrun会自动设置rank和world_size
        # 这里仅为单进程调试演示，实际运行请使用下面的torchrun命令
        print("请使用torchrun启动脚本以启用分布式训练。")

如何运行 FSDP 示例 (使用2个GPU)：

为了实际执行分布式训练，我们需要使用torchrun（或旧版torch.distributed.launch）来启动多个进程。

# 假设你有两块可用的GPU (nproc_per_node=2)
torchrun --nproc_per_node=2 fsdp_llm_train.py

通过FSDP，我们解决了大型模型训练中的内存瓶颈问题，这正是构建大规模AI系统（如ChatGPT的基础训练平台）时必须掌握的核心基础设施技术。从PyTorch生态出发，开发者能够更快速、更高效地迭代和部署下一代LLM模型。