import pytorch_lightning as pl
import torch
from torch import nn
import torch.nn.functional as F

class LitMNIST(pl.LightningModule):
    def __init__(self, hidden_dim=128, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.l1 = nn.Linear(28 * 28, self.hparams.hidden_dim)
        self.l2 = nn.Linear(self.hparams.hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.l1(x))
        return self.l2(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        self.log("train_loss", loss, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(1) == y).float().mean()
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.lr)

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([transforms.ToTensor()])
train_ds = datasets.MNIST(".", train=True, download=True, transform=transform)
val_ds = datasets.MNIST(".", train=False, transform=transform)

train_loader = DataLoader(train_ds, batch_size=128, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=128, num_workers=4)

model = LitMNIST()
trainer = pl.Trainer(
    max_epochs=10,
    accelerator="gpu" if torch.cuda.is_available() else "cpu",
    devices=1,
    log_every_n_steps=10
)
trainer.fit(model, train_loader, val_loader)

from pytorch_lightning.callbacks import (
    EarlyStopping, ModelCheckpoint, LearningRateMonitor
)

callbacks = [
    ModelCheckpoint(
        save_top_k=1,
        monitor="val_acc",
        mode="max",
        filename="best-{epoch:02d}-{val_acc:.3f}"
    ),
    EarlyStopping(monitor="val_loss", patience=3, mode="min"),
    LearningRateMonitor(logging_interval="step")
]

import wandb
from pytorch_lightning.loggers import WandbLogger

wandb.login()
logger = WandbLogger(project="mnist-lightning", name="exp-1")
trainer = pl.Trainer(logger=logger, callbacks=callbacks)

trainer = pl.Trainer(
    strategy="ddp",           # Distributed Data Parallel
    accelerator="gpu",
    devices=8,
    precision=16,             # Mixed precision (FP16)
    max_epochs=50
)

trainer = pl.Trainer(
    strategy="fsdp",          # Shards model across GPUs
    accelerator="gpu",
    devices=8,
    precision="bf16-mixed"    # BFloat16
)

model = torch.compile(model)  # Before trainer.fit()

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(r=8, lora_alpha=32, target_modules=["q", "v"])
model = get_peft_model(model, lora_config)

import torch.quantization

model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model = torch.quantization.prepare(model)
# Calibrate, then:
model = torch.quantization.convert(model)

from lightning.fabric import Fabric

fabric = Fabric(accelerator="gpu", devices=4, strategy="ddp")
fabric.launch()

model, optimizer = fabric.setup(model, optimizer)
dataloader = fabric.setup_dataloaders(dataloader)

for batch in dataloader:
    output = model(batch)
    loss = loss_fn(output, batch)
    fabric.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

def test_model_output_shape():
    model = LitMNIST()
    x = torch.randn(1, 1, 28, 28)
    out = model(x)
    assert out.shape == (1, 10)

name: CI
on: [push]
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3
      - name: Install
        run: pip install -r requirements.txt
      - name: Test
        run: pytest tests/

from transformers import AutoTokenizer
from lightning.pytorch import LightningDataModule

class TweetDataModule(LightningDataModule):
    def __init__(self, batch_size=32):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
        self.batch_size = batch_size

    def setup(self, stage):
        # Load 10M tweets
        pass

    def train_dataloader(self):
        return DataLoader(..., batch_size=self.batch_size, num_workers=8)

trainer = pl.Trainer(
    strategy="fsdp",
    accelerator="gpu",
    devices=8,
    precision="bf16-mixed",
    max_epochs=3,
    accumulate_grad_batches=4,
    logger=WandbLogger(project="tweet-sentiment"),
    callbacks=[ModelCheckpoint(monitor="val_acc", mode="max")]
)
trainer.fit(model, datamodule)

# Export to ONNX
torch.onnx.export(model, dummy_input, "bert_sentiment.onnx")

pip install lightning wandb peft

import pytorch_lightning as pl
print(pl.__version__)  # 2.2+