# FSDP2

[Fully Sharded Data Parallel (FSDP2)](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) 将模型、梯度和优化器状态分片到多个 GPU 上。在计算之前，每个 GPU 会从所有分片中收集完整的参数集合，计算完成后释放这些参数。分片机制允许你训练超出单个 GPU 显存容量的模型，但代价是比 [DDP](./ddp) 产生更多的通信开销。当你的模型或优化器状态无法放入单个 GPU 时，请使用 FSDP。

```text
                      ┌─────────────────┐
                      │  training data  │
                      └────────┬────────┘
            ┌──────────────────┼──────────────────┐
            │ shard 0          │ shard 1          │ shard 2
            ▼                  ▼                  ▼
     ┌─────────────┐    ┌─────────────┐    ┌─────────────┐
     │  param      │    │  param      │    │  param      │
     │  shard 0    │    │  shard 1    │    │  shard 2    │
     │  GPU 0      │    │  GPU 1      │    │  GPU 2      │
     └──────┬──────┘    └──────┬──────┘    └──────┬──────┘
            │                  │                  │
            └──────── all-gather (params) ────────┘
                               │
                    full params on each GPU
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
         forward             forward             forward
            │                  │                  │
            └───── reduce-scatter (grads) ────────┘
                               │
            ┌──────────────────┼──────────────────┐
            ▼                  ▼                  ▼
     grad shard 0       grad shard 1       grad shard 2
     optim shard 0      optim shard 1      optim shard 2
        step               step               step
```

## 分片策略

FSDP2 通过 `~TrainingArguments.fsdp_config` 控制分片。设置 `fsdp=True` 以启用 FSDP，并在 FSDP 配置中设置 `reshard_after_forward` 来选择内存与吞吐量之间的权衡。

| `reshard_after_forward` | 行为 |
|---|---|
| `true` | 在前向传播后重新分片参数以节省更多内存 |
| `false` | 在正向传播和反向传播之间保持参数聚合，以避免重新执行 all-gather，代价是更高的峰值内存 |

`auto_wrap_policy` 控制模块如何被封装为 FSDP 单元。其默认值为 `"TRANSFORMER_BASED_WRAP"`，该策略会封装模型的 transformer 层。如果不进行封装（`"NO_WRAP"`），整个模型将作为一个 FSDP 单元，从而失去分片带来的内存优势。

## 配置 FSDP

这些字段控制 FSDP2 如何包装、分片和加载模型。`reshard_after_forward` 和 `auto_wrap_policy` 在 [分片策略](#sharding-strategies) 中进行了介绍。

- `cpu_offload` 会在参数和梯度未被使用时将其卸载到 CPU，以节省 GPU 显存。

- `transformer_layer_cls_to_wrap` 定义了当 `auto_wrap_policy` 为 `"TRANSFORMER_BASED_WRAP"` 时，需要封装为 FSDP 单元的 transformer 层。每个单元独立管理自身的 gather 和 scatter 操作。在前向传播过程中，仅收集当前单元的参数，而释放之前单元的参数以节省内存。

  仅包装顶层模型不会带来任何 GPU 显存节省。包装每一个单独的 `Linear` 层会使单元间通信变得非常昂贵。将此字段留空，FSDP 将从模型定义中读取该值。

- `min_num_params` 设置了基于大小进行包装时每个模块的最小参数量。它仅在 `auto_wrap_policy` 为 `"SIZE_BASED_WRAP"` 时使用。

- `state_dict_type` 控制检查点的格式。默认值为 `"FULL_STATE_DICT"`，对应单个兼容 Transformers 的检查点。对于大型模型，使用 `"SHARDED_STATE_DICT"` 可以为每个 rank 生成一个检查点文件，速度更快。分片检查点只能加载回 FSDP，因此对于最终需要共享或在 FSDP 之外加载的检查点，请保存为 `"FULL_STATE_DICT"`。

- `cpu_ram_efficient_loading` 仅在 rank 0 上从磁盘加载检查点。其他 GPU 初始化一个空模型，并通过广播接收权重，从而避免多个进程将大型模型加载到 CPU 内存中。

- `activation_checkpointing` 会在反向传播过程中重新计算激活值，而不是将其存储。请在 [TrainingArguments](/docs/transformers/main/zh/main_classes/trainer#transformers.TrainingArguments) 中使用此参数来替代 [gradient checkpointing](./grad_checkpointing)。同时设置两者会引发错误。

通过 [Accelerate 配置文件](./accelerate#accelerate-config-file) 或传递给 `fsdp_config` 的 FSDP 配置文件来配置 FSDP 训练。

运行 [accelerate config](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-config) 命令，并根据你的硬件和训练配置回答相关问题。这会在你的缓存中创建一个 `default_config.yaml` 文件。

使用 [accelerate launch](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-launch) 运行基于 [Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 的脚本。`fsdp_config` 是多余的，因为 Accelerate 配置文件已涵盖相同的设置。

```cli
accelerate launch train.py
```

```json
{
  "version": 2,
  "reshard_after_forward": true,
  "cpu_offload": false,
  "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
  "transformer_layer_cls_to_wrap": ["LlamaDecoderLayer"],
  "state_dict_type": "FULL_STATE_DICT",
  "cpu_ram_efficient_loading": true,
  "activation_checkpointing": true
}
```

设置 `fsdp=True`，并将 FSDP 配置文件传递给 `fsdp_config`。

```py
from transformers import TrainingArguments

TrainingArguments(
    ...,
    fsdp=True,
    fsdp_config="path/to/fsdp.json",
)
```

> [!TIP]
> 对于混合专家模型，`fsdp_size` 可以与 `tp_size` 和 `ep_size` 结合使用，将专家分片到一个网格维度上，其余部分分片到另一个维度上。参见 [专家并行](./expert_parallelism#combining-with-fsdp2)。

## 使用上下文并行进行评估

通过 Accelerate 的 `ParallelismConfig` 中的 `cp_size` 来设置 FSDP2 的上下文并行（CP）。启动时，需同时传入 `--use_parallelism_config` 和 `--parallelism_config_cp_size`。如果 Accelerate YAML 中的 `parallelism_config` 非空，也会自动启用 `use_parallelism_config`。如果仅传入 `--parallelism_config_cp_size`，该参数将被忽略，`cp_size=1` 保持不变。

在 FSDP2 下，CP 将序列分片到各个 rank，使每个 GPU 持有上下文的一个切片（`allgather` 或 `alltoall`）。这与 [Ulysses 序列并行](./deepspeed_alst) 不同，后者通过 DeepSpeed ALST 运行，并在同一配置中使用 `sp_size`。请在 CP 和 Ulysses SP 之间二选一。不要同时将 `cp_size` 和 `sp_size` 设置为大于 1。

[evaluate()](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer.evaluate) 通过 [prediction_step()](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer.prediction_step) 使用与 [train()](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer.train) 相同的上下文并行上下文。在该上下文激活期间，每个 rank 运行其序列分片。`eval_loss` 是正确的损失值（未乘以 `cp_size`）。早停和 `load_best_model_at_end` 使用该指标。

Logits 沿序列维度保持分片状态，且没有任何机制将其重组为完整长度的预测。在 CP 下，不支持基于 logits 的预测以及 `compute_metrics`。

当您需要 `eval_loss` 时，请复用现有的 FSDP2 + `cp_size` 启动配置。

如果您子类化了 [prediction_step()](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer.prediction_step) 并自行统计 token 数量，请在进入上下文并行上下文之前统计 `num_items_in_batch`。在该上下文激活期间，序列缓冲区会在各 rank 之间就地分片。之后再进行统计会导致批次数量被低估。默认的 [Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 已经会先进行统计，只有当您重写 `prediction_step` 时才需要这样做。

## 后续步骤

- 如果模型可以放入单个 GPU，请参阅 [DDP](./ddp) 以进行数据并行训练。
- 有关 ZeRO 优化和 NVMe 卸载，请参阅 [DeepSpeed](./deepspeed)。
- 对于使用 PyTorch/XLA 在 TPU 上运行 FSDP，请在 `~TrainingArguments.fsdp_config` 中设置 `xla`、`xla_fsdp_settings` 和 `xla_fsdp_grad_ckpt`。
- Read the [FSDP chapter](https://nanotron-ultrascale-playbook.static.hf.space/index.html#zero-3:_adding_parameter_partitioning_(fsdp)) from The Ultra-Scale Playbook for more information about how FSDP works.