PyTorch FSDP Anatomy: FULL_SHARD vs SHARD_GRAD_OP vs HYBRID_SHARD + Mixed Precision Policy

# === FSDP + Llama 3.3 70B Full FT (8×H100 SXM) ===
import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
    BackwardPrefetch,
    CPUOffload,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
 
# 1. Init distributed
dist.init_process_group(backend="nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 
# 2. Load model on each GPU (her GPU önce full model'i yükler, sonra shard'lanır)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.3-70B-Instruct",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
 
# 3. Mixed precision policy
bf16_policy = MixedPrecision(
    param_dtype=torch.bfloat16,         # param compute'ta bf16
    reduce_dtype=torch.float32,         # gradient reduce fp32 (numerik kararlılık)
    buffer_dtype=torch.bfloat16,        # buffer'lar bf16
)
 
# 4. Auto-wrap policy — transformer layer-wise
auto_wrap = transformer_auto_wrap_policy(
    transformer_layer_cls={LlamaDecoderLayer},
)
 
# 5. FSDP wrap
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=bf16_policy,
    auto_wrap_policy=auto_wrap,
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,    # next param prefetch
    use_orig_params=True,                                # PyTorch 2.0+ recommended
    device_id=torch.cuda.current_device(),
    sync_module_states=True,                             # init sync across ranks
    cpu_offload=CPUOffload(offload_params=False),        # disable for speed
)
 
# 6. Optimizer — FSDP-aware
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=5e-6,                            # full FT için DÜŞÜK lr (LoRA 2e-4'ten)
    weight_decay=0.01,
)
 
# 7. Training loop
for batch in loader:
    out = model(**batch)
    out.loss.backward()
    optimizer.step()
    optimizer.zero_grad()
 
# 8. Save FSDP state — full state_dict consolidate
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):
    if dist.get_rank() == 0:
        torch.save(model.state_dict(), "llama-3.3-70b-tr-finetuned.pt")