FSDP2 (fully_shard): Per-Parameter Sharding + DTensor + 2024+ PyTorch Innovation

# === FSDP2 reçetesi ===
import torch
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
import torch.distributed.checkpoint as dcp
 
# Distributed init
dist.init_process_group(backend="nccl")
 
# Model
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.3-70B-Instruct",
    torch_dtype=torch.bfloat16,
)
 
# Mixed precision per-module
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32,
    output_dtype=torch.bfloat16,
)
 
# Wrap her decoder layer'ı ayrı
for layer in model.model.layers:
    fully_shard(layer, mp_policy=mp_policy)
# Root model'i de wrap
fully_shard(model, mp_policy=mp_policy)
 
# === DCP Checkpoint save ===
state_dict = get_state_dict(model, optimizer)
dcp.save(state_dict, checkpoint_id="ckpt/step-1000/")
# /ckpt/step-1000/ klasörü her rank parça parça yazar (parallel I/O)
 
# === Resume — start_rank=4 ile yeniden başlat
state_dict = get_state_dict(model, optimizer)
dcp.load(state_dict, checkpoint_id="ckpt/step-1000/")
set_state_dict(model, optimizer, model_state_dict=state_dict)