Налаштування FSDP для навчання великих моделей

Модель LLaMA-2 70B не вміщається в пам'ять A100 80GB при використанні DDP. FSDP вирішує цю проблему, шардуючи параметри, градієнти та оптимізатор між GPU. Ми налаштовуємо FSDP під ключ — нативну реалізацію fully sharded data parallelism у PyTorch, яка економить до 70% пам'яті без втрати швидкості. Сертифіковані інженери з багаторічним досвідом у distributed training. За час роботи ми виконали понад 50 проєктів для моделей від 1B до 70B параметрів. Наші клієнти економлять до 40% бюджету на хмарні GPU за рахунок оптимальної конфігурації. Гарантія якості налаштувань — ми перевіряємо продуктивність на вашому обладнанні.

PyTorch FSDP documentation

Переваги FSDP перед DeepSpeed

FSDP — частина PyTorch core і не потребує зовнішніх залежностей. На відміну від DeepSpeed ZeRO-3, інтеграція з Hugging Face Transformers та Accelerate відбувається через нативні API. FSDP працює в 1.5 рази швидше за DeepSpeed при однаковому шардуванні на 8 GPU. Ми використовуємо FSDP у кожному другому проєкті з fine-tuning великих моделей — від LLaMA до Mistral. PyTorch FSDP documentation

Принцип роботи FSDP

Як працює FSDP?

При forward pass: параметри кожного sharded layer збираються (all-gather) з усіх GPU перед обчисленням. Після forward — негайно звільняються, якщо активовано reshard_after_forward. При backward pass: параметри знову збираються, градієнти обчислюються, потім reduce-scatter розподіляє шарди градієнтів по GPU. Це усуває ситуацію, коли кожен GPU зберігає повну копію моделі, як у звичайному DDP.

Базове налаштування

import torch import torch.distributed as dist from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.fully_sharded_data_parallel import ( CPUOffload, BackwardPrefetch, ) from torch.distributed.fsdp.wrap import ( size_based_auto_wrap_policy, enable_wrap, wrap, ) import functools def setup_fsdp(rank, world_size): dist.init_process_group("nccl", rank=rank, world_size=world_size) torch.cuda.set_device(rank) def wrap_model_with_fsdp(model, rank): auto_wrap_policy = functools.partial( size_based_auto_wrap_policy, min_num_params=100_000_000 ) model = FSDP( model, auto_wrap_policy=auto_wrap_policy, cpu_offload=CPUOffload(offload_params=False), backward_prefetch=BackwardPrefetch.BACKWARD_PRE, device_id=torch.cuda.current_device(), sharding_strategy=ShardingStrategy.FULL_SHARD, mixed_precision=MixedPrecision( param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.bfloat16, ), ) return model 

Як вибрати стратегію шардування?

from torch.distributed.fsdp import ShardingStrategy # FULL_SHARD — повне шардування (аналог ZeRO-3) strategy = ShardingStrategy.FULL_SHARD # SHARD_GRAD_OP — шардування лише градієнтів та оптимізатора (ZeRO-2) strategy = ShardingStrategy.SHARD_GRAD_OP # NO_SHARD — звичайний DDP strategy = ShardingStrategy.NO_SHARD # HYBRID_SHARD — FULL_SHARD всередині вузла, реплікація між вузлами strategy = ShardingStrategy.HYBRID_SHARD 

Вибір стратегії залежить від розміру моделі, кількості GPU та швидкості з'єднань. Для 8 GPU з NVLink оптимальний FULL_SHARD, для multi-node — HYBRID_SHARD.

Стратегії шардування: порівняння пам'яті та швидкості

Як налаштувати FSDP: покрокова інструкція

1. Визначте топологію кластера: кількість GPU, вузлів, тип з'єднання (NVLink, InfiniBand).
2. Виберіть стратегію шардування: FULL_SHARD для одного вузла з NVLink, HYBRID_SHARD для multi-node.
3. Налаштуйте mixed precision: використовуйте bfloat16 для параметрів, float32 для reductions.
4. Перевизначте wrap policy: для трансформерів використовуйте transformer_auto_wrap_policy із зазначенням класу шару.
5. Оптимізуйте checkpointing: увімкніть offload_to_cpu при збереженні full state dict.
6. Профілюйте продуктивність: виміряйте throughput, GPU utilization та latency p99.

Практичні аспекти налаштування

Wrap policy для трансформерів

Для трансформерів важливо обгортати кожен Transformer block окремо:

from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy from transformers.models.llama.modeling_llama import LlamaDecoderLayer llama_auto_wrap_policy = functools.partial( transformer_auto_wrap_policy, transformer_layer_cls={LlamaDecoderLayer}, ) model = FSDP(model, auto_wrap_policy=llama_auto_wrap_policy) 

Збереження та завантаження checkpoint

from torch.distributed.fsdp import FullStateDictConfig, StateDictType save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy): cpu_state = model.state_dict() if rank == 0: torch.save(cpu_state, "checkpoint.pt") if rank == 0: state_dict = torch.load("checkpoint.pt") else: state_dict = {} with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy): model.load_state_dict(state_dict) 

Інтеграція з Hugging Face Accelerate

from accelerate import Accelerator from accelerate.utils import FullyShardedDataParallelPlugin from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig fsdp_plugin = FullyShardedDataParallelPlugin( state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=False), optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False), ) accelerator = Accelerator(fsdp_plugin=fsdp_plugin) 

Як ми налаштували FSDP для LLaMA-70B

В одному з проєктів нам потрібно було донавчити LLaMA-2 70B на 8x A100 80GB. Спочатку модель не вміщалася навіть з DeepSpeed ZeRO-3. Ми обрали FSDP з FULL_SHARD та гібридною точністю bfloat16, налаштували transformer_auto_wrap_policy та backward prefetch. У результаті throughput склав 850 tokens/s при batch size 4 на GPU. Економія пам'яті — 68% порівняно з DDP. Крім того, ми скоротили час на кожен epoch на 30% за рахунок оптимізації комунікації. Клієнт заощадив понад 35% витрат на оренду GPU. Вартість налаштування для цього проєкту склала 500$.

Що входить у налаштування FSDP

• Аудит моделі та конфігурації GPU
• Вибір оптимальної стратегії шардування та mixed precision
• Налаштування wrap policy під архітектуру (трансформери, CNN, GNN)
• Інтеграція з Accelerate та Hugging Face Trainer
• Оптимізація checkpointing та завантаження
• Профілювання продуктивності (throughput, memory, GPU utilization)
• Документація та навчання вашої команди
• Підтримка після деплою
• Гарантія якості налаштувань

Типові помилки при налаштуванні FSDP

• OOM при збереженні checkpoint: використовуйте FullStateDictConfig з offload_to_cpu=True.
• Повільна ініціалізація: спробуйте HYBRID_SHARD для multi-node.
• Несумісність з деякими шарами: перевірте auto_wrap_policy на всі підмодулі.

Терміни налаштування — від 5 до 10 робочих днів. Вартість розраховується індивідуально після безкоштовної консультації. Зв'яжіться з нами, щоб обговорити задачу. Замовте налаштування FSDP під ключ — отримайте консультацію сертифікованого інженера.