lightning example

Author: pmcurtin

docs/source/examples/lightning.md

@@ -0,0 +1,5 @@
+# Pytorch Lightning
+
+```{eval-rst}
+.. literalinclude:: ./lightning.py
+```

dsp.py renamed to docs/source/examples/lightning_example.py

@@ -1,10 +1,10 @@
-from dataclasses import dataclass
+import os
 from pathlib import Path
 
-import deepspeed
+import lightning as L
 import torch
-
 from datasets import load_dataset
+
 from torch import nn
 from torch.utils.data import Dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -42,47 +42,54 @@ def __getitem__(self, idx):
         }
 
 
-@dataclass
-class DSPArgs:
-    deepspeed_config: str
-    # train_batch_size: int
-    # batch_size: int
+class GPT2LightningWrapper(L.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+    def training_step(self, batch, batch_idx):
+        device_batch = {k: v.to(self.model.device) for k, v in batch.items()}
+        loss = self.model(**device_batch).loss
+        self.log("train_loss", loss)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-5)
+        return optimizer
 
 
 def train():
-    model = AutoModelForCausalLM.from_pretrained("gpt2")
-    # optimizer = torch.optim.Adam(model.parameters())
+    lightning_model = GPT2LightningWrapper()
+
     wikitext_train = load_dataset("Salesforce/wikitext", "wikitext-2-v1", split="train")
     train_dataset = GPT2CausalLMDataset(text_dataset=wikitext_train)
-
-    loader = torch.utils.data.DataLoader(train_dataset, batch_size=8)
-
-    model_engine, optimizer, _, _ = deepspeed.initialize(
-        args=DSPArgs(deepspeed_config="dsp_config.json"),
-        model=model,
-        model_parameters=model.parameters(),
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8)
+
+    trainer = L.Trainer(
+        accelerator="gpu",
+        limit_train_batches=10,
+        max_epochs=1,
+        devices=2,
+        num_nodes=1,
+        strategy="ddp",
     )
 
-    model.train()
-    for batch_idx, batch in enumerate(loader):
-        if batch_idx == 10:
-            break
-        print(f"Step {batch_idx}")
-
-        device_batch = {k: v.to(model.device) for k, v in batch.items()}
-
-        model.zero_grad()
+    trainer.fit(model=lightning_model, train_dataloaders=train_loader)
 
-        loss = model_engine(**device_batch).loss
-        model_engine.backward(loss)
-
-        model_engine.step()
+    if int(os.environ["RANK"]) == 0:
+        return trainer.model.model
+    return None
 
 
 if __name__ == "__main__":
+    # hack to prevent lightning from recognizing SLURM environment...
+    os.environ["SLURM_JOB_NAME"] = "bash"
     Path("output").mkdir(exist_ok=True)
     results = torchrunx.launch(
         func=train,
         hostnames=["localhost"],
-        workers_per_host=1,
+        workers_per_host=2,
     )
+
+    trained_model: nn.Module = results.rank(0)
+    torch.save(trained_model.state_dict(), "output/model.pth")