FSDP

Что такое FSDP

FSDP (Fully Sharded Data Parallel) — встроенный в PyTorch (с версии 1.12) механизм шардирования состояния обучения между GPU. По эффекту аналогичен ZeRO-3: каждый GPU хранит только 1/N весов, градиентов и состояний оптимизатора. Перед forward/backward выполняется AllGather для сбора нужных параметров.

Ключевое преимущество перед DeepSpeed — нативная интеграция в PyTorch без внешних зависимостей, чистый Python API.

Режимы шардирования

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

# FULL_SHARD = ZeRO-3: веса + градиенты + оптимизатор
model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)

# SHARD_GRAD_OP = ZeRO-2: только градиенты + оптимизатор  
model = FSDP(model, sharding_strategy=ShardingStrategy.SHARD_GRAD_OP)

# NO_SHARD = стандартный DDP (без шардирования)
model = FSDP(model, sharding_strategy=ShardingStrategy.NO_SHARD)

FSDP с Hugging Face Trainer

from transformers import TrainingArguments

args = TrainingArguments(
    fsdp="full_shard",
    fsdp_config={
        "fsdp_transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
        "fsdp_backward_prefetch": "backward_pre",
    },
    bf16=True,
    ...
)

FSDP vs DeepSpeed ZeRO

Для большинства задач fine-tuning выбор между FSDP и DeepSpeed — вопрос предпочтения. При нужде в CPU offload — DeepSpeed.

Связанные термины

• ZeRO — аналогичный алгоритм в DeepSpeed
• DeepSpeed — альтернативный фреймворк
• gradient checkpointing — совместно используется с FSDP
• tensor parallelism — другой тип параллелизма