Dağıtık Eğitim
DeepSpeed & FSDP.

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import os

def setup(rank: int, world_size: int):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train(rank: int, world_size: int, model, dataset, epochs: int = 3):
    setup(rank, world_size)

    # Modeli GPU'ya taşı ve DDP'ye sar
    model = model.to(rank)
    model = DDP(model, device_ids=[rank])

    # DistributedSampler: her GPU farklı alt-küme görür
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=4, sampler=sampler, num_workers=4, pin_memory=True
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # Her epoch farklı shuffle
        for batch in dataloader:
            input_ids = batch["input_ids"].to(rank)
            labels    = batch["labels"].to(rank)

            loss = model(input_ids, labels=labels).loss
            loss.backward()    # All-reduce otomatik
            optimizer.step()
            optimizer.zero_grad()

    cleanup()

# Başlatma (torchrun kullan)
# torchrun --nproc_per_node=4 ddp_training.py
if __name__ == "__main__":
    import torch.multiprocessing as mp
    world_size = torch.cuda.device_count()
    mp.spawn(train, args=(world_size, model, dataset), nprocs=world_size)

# Basitleştirilmiş tensor parallelism örneği
# W: [d_model, 4*d_model] matrisini iki GPU'ya böl

# GPU 0: W_0 = W[:, :2*d_model]   → ilk yarı sütunlar
# GPU 1: W_1 = W[:, 2*d_model:]   → son yarı sütunlar

import torch
import torch.distributed as dist

def column_parallel_linear(x, weight_shard, bias_shard=None):
    # Her GPU kendi sütun dilimini çarpar
    local_out = torch.nn.functional.linear(x, weight_shard, bias_shard)
    return local_out   # All-gather sonradan yapılır

def row_parallel_linear(x_shard, weight_shard, bias=None):
    # Her GPU kendi satır dilimini çarpar, sonra all-reduce
    local_out = torch.nn.functional.linear(x_shard, weight_shard)
    dist.all_reduce(local_out, op=dist.ReduceOp.SUM)
    if bias is not None:
        local_out += bias
    return local_out

# N = GPU sayısı, Ψ = parametre sayısı
# Adam: fp32 weights(4) + fp32 momentum(4) + fp32 variance(4) + fp16 weights(2) + fp16 grad(2) = 16 byte/param

# Baseline (DDP):
baseline = "16Ψ bytes per GPU"

# ZeRO Stage 1 (optimizer state split):
stage1 = "4Ψ + 12Ψ/N bytes per GPU"

# ZeRO Stage 2 (gradient + optimizer split):
stage2 = "2Ψ + 14Ψ/N bytes per GPU"

# ZeRO Stage 3 (full split):
stage3 = "16Ψ/N bytes per GPU"  # Lineer ölçekleme!

def zero_memory(params_b: float, n_gpus: int) -> dict:
    """params_b: parametre sayısı milyar olarak"""
    P = params_b * 1e9 * 16 / 1e9  # GB (baseline)
    return {
        "baseline_gb":  P,
        "stage1_gb": 4/16 * P + (12/16 * P) / n_gpus,
        "stage2_gb": 2/16 * P + (14/16 * P) / n_gpus,
        "stage3_gb": P / n_gpus,
    }

# LLaMA-3 8B, 8 GPU:
for k, v in zero_memory(8, 8).items():
    print(f"{k:12s}: {v:.1f} GB/GPU")
# baseline_gb :  80.0 GB/GPU
# stage1_gb   :  28.8 GB/GPU
# stage2_gb   :  18.0 GB/GPU
# stage3_gb   :  10.0 GB/GPU

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "bf16": {
    "enabled": true
  },
  "gradient_accumulation_steps": 4,
  "gradient_clipping": 1.0,
  "steps_per_print": 100,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer
)
from peft import LoraConfig, get_peft_model
from datasets import load_dataset

model_id = "meta-llama/Meta-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype="auto"
)

# LoRA config
lora_config = LoraConfig(
    r=16, lora_alpha=32,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    lora_dropout=0.05, bias="none"
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 83,886,080 || total: 8,114,896,896 || 1.034%

# Training argümanları — deepspeed config'e işaret et
training_args = TrainingArguments(
    output_dir="./llama3-lora",
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-4,
    bf16=True,
    logging_steps=10,
    save_steps=200,
    deepspeed="ds_config_zero3.json",  # Tek satır!
    report_to="wandb"
)

dataset = load_dataset("tatsu-lab/alpaca", split="train")
trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()

# Başlatma:
# deepspeed --num_gpus=4 train_deepspeed.py
# veya:
# torchrun --nproc_per_node=4 train_deepspeed.py

import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    ShardingStrategy,
    BackwardPrefetch,
    MixedPrecision,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
import functools

def create_fsdp_model(model, rank: int):
    # Mixed precision policy
    mp_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    )

    # LlamaDecoderLayer'ı wrap policy ile FSDP birimlerine böl
    wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={LlamaDecoderLayer},
    )

    fsdp_model = FSDP(
        model,
        auto_wrap_policy=wrap_policy,
        mixed_precision=mp_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
        device_id=rank,
        use_orig_params=True,   # LoRA için gerekli
    )
    return fsdp_model

# Checkpoint kaydetme (FSDP ile özel state_dict gerekir)
from torch.distributed.fsdp import FullStateDictConfig, StateDictType

def save_fsdp_model(model, output_dir: str, rank: int):
    cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
        state_dict = model.state_dict()
    if rank == 0:
        torch.save(state_dict, f"{output_dir}/model.pt")

import torch
from torch.cuda.amp import autocast, GradScaler

model = MyModel().to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()   # Yalnızca fp16 için gerekli

for batch in dataloader:
    # bf16: torch_dtype=torch.bfloat16 — scaler gerekmez
    with autocast(device_type="cuda", dtype=torch.bfloat16):
        output = model(batch["input_ids"].to("cuda"))
        loss = output.loss

    # fp16 için scaler kullan
    # scaler.scale(loss).backward()
    # scaler.step(optimizer)
    # scaler.update()

    # bf16 için direkt
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad()

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")

# Tek satırda aktif et
model.gradient_checkpointing_enable()

# Veya manuel PyTorch checkpoint
from torch.utils.checkpoint import checkpoint

class CheckpointedLayer(torch.nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def forward(self, x):
        # Aktivasyonları saklamaz, backward'da yeniden hesaplar
        return checkpoint(self.layer, x, use_reentrant=False)

# Bellek karşılaştırması (LLaMA-3 8B, seq=2048, bs=1):
# Gradient checkpointing kapalı:  ~60 GB
# Gradient checkpointing açık:    ~28 GB
# Fiyat: ~%30 daha yavaş eğitim

# Kurulum
# pip install flash-attn --no-build-isolation

from transformers import AutoModelForCausalLM
import torch

# HuggingFace ile tek parametre
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",  # ← burada
    device_map="auto"
)

# Manuel kullanım
from flash_attn import flash_attn_func

def flash_attention_forward(q, k, v, dropout_p=0.0, causal=True):
    """
    q, k, v: (batch, seq, heads, head_dim) — float16 veya bfloat16
    """
    return flash_attn_func(
        q, k, v,
        dropout_p=dropout_p,
        softmax_scale=q.shape[-1] ** (-0.5),
        causal=causal    # Decoder için True
    )

# Bellek karşılaştırma (seq_len=4096, 8B model):
# Standart attn:  ~45 GB aktivasyon
# Flash attn 2:   ~12 GB aktivasyon

import torch
from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, DataCollatorForLanguageModeling
)
from trl import SFTTrainer
from peft import LoraConfig, get_peft_model, TaskType
from datasets import load_dataset
import os

MODEL_ID  = "meta-llama/Meta-Llama-3-8B"
OUTPUT    = "./llama3-8b-lora-finetuned"
DS_CONFIG = "./ds_config_zero3.json"  # yukarıdaki JSON

# 1. Model ve tokenizer yükle
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
model.enable_input_require_grads()

# 2. LoRA config
lora_cfg = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.05,
    bias="none",
    modules_to_save=["embed_tokens", "lm_head"],
)
model = get_peft_model(model, lora_cfg)
model.print_trainable_parameters()

# 3. Veri seti
dataset = load_dataset("tatsu-lab/alpaca", split="train")

def format_prompt(example):
    if example["input"]:
        text = f"### Talimat:\n{example['instruction']}\n\n### Giriş:\n{example['input']}\n\n### Yanıt:\n{example['output']}"
    else:
        text = f"### Talimat:\n{example['instruction']}\n\n### Yanıt:\n{example['output']}"
    return {"text": text}

dataset = dataset.map(format_prompt)

# 4. Training arguments
args = TrainingArguments(
    output_dir=OUTPUT,
    num_train_epochs=3,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    gradient_checkpointing=True,
    learning_rate=2e-4,
    weight_decay=0.01,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,
    logging_steps=10,
    save_strategy="epoch",
    deepspeed=DS_CONFIG,      # ← ZeRO3 aktif
    dataloader_num_workers=4,
    report_to="wandb",
    run_name="llama3-8b-lora-zero3",
)

# 5. SFT Trainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=2048,
    args=args,
)

trainer.train()
trainer.save_model(OUTPUT)

# Başlatma:
# deepspeed --num_gpus=4 finetune_llama3_lora.py
# W&B ile izle: wandb.ai/your_project

import torch
from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=tensorboard_trace_handler("./profiler_logs"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    for step, batch in enumerate(dataloader):
        if step >= 5:
            break
        train_step(batch)
        prof.step()

# TensorBoard'da izle:
# tensorboard --logdir=./profiler_logs