chore: import upstream snapshot with attribution

2026-07-13 12:47:19 +08:00
commit d5f64af28f
253 changed files with 46818 additions and 0 deletions
@@ -0,0 +1,27 @@
+## Minimal PyTorch Lightning Trainer Example
+
+
+
+The script in this folder provides minimal examples showing how to train a LitGPT model using LitGPT's `GPT` class with the [PyTorch Lightning](https://github.com/Lightning-AI/pytorch-lightning) Trainer.
+
+You can run the scripts as follows:
+
+&nbsp
+## Small 160M model:
+
+```bash
+# Download the Pythia model
+litgpt download EleutherAI/pythia-160m
+
+python litgpt_ptl_small.py
+```
+
+&nbsp
+## Medium-sized 8B model:
+
+```bash
+# Download the Llama 3.1 model
+litgpt download meta-llama/Meta-Llama-3.1-8B --access_token hf_...
+
+python litgpt_ptl_medium.py
+```
@@ -0,0 +1,58 @@
+import lightning as L
+import torch
+
+import litgpt
+from litgpt.data import Alpaca2k
+from litgpt.lora import GPT, merge_lora_weights
+
+
+class LitLLM(L.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.model = GPT.from_name(
+            name="Llama-3.1-8B",
+            lora_r=32,
+            lora_alpha=16,
+            lora_dropout=0.05,
+            lora_key=False,
+            lora_value=True,
+        )
+        litgpt.lora.mark_only_lora_as_trainable(self.model)
+
+    def on_train_start(self):
+        state_dict = torch.load("checkpoints/meta-llama/Meta-Llama-3.1-8B/lit_model.pth", mmap=True)
+        self.model.load_state_dict(state_dict, strict=False)
+
+    def training_step(self, batch):
+        input_ids, targets = batch["input_ids"], batch["labels"]
+        logits = self.model(input_ids)
+        loss = litgpt.utils.chunked_cross_entropy(logits[..., :-1, :], targets[..., 1:])
+        self.log("train_loss", loss, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        warmup_steps = 10
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=0.0002, weight_decay=0.0, betas=(0.9, 0.95))
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: step / warmup_steps)
+        return [optimizer], [scheduler]
+
+
+if __name__ == "__main__":
+    data = Alpaca2k()
+    tokenizer = litgpt.Tokenizer("checkpoints/meta-llama/Meta-Llama-3.1-8B")
+    data.connect(tokenizer, batch_size=1, max_seq_length=512)
+
+    trainer = L.Trainer(
+        devices=1,
+        max_epochs=2,
+        accumulate_grad_batches=8,
+        precision="bf16-true",
+    )
+    with trainer.init_module(empty_init=True):
+        model = LitLLM()
+
+    trainer.fit(model, data)
+
+    # Save final checkpoint
+    merge_lora_weights(model.model)
+    trainer.save_checkpoint("checkpoints/finetuned.ckpt", weights_only=True)
@@ -0,0 +1,150 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+import lightning as L
+import torch
+
+from litgpt import LLM
+from litgpt.data import Alpaca2k
+
+
+class LitLLM(L.LightningModule):
+    def __init__(self, checkpoint_dir, tokenizer_dir=None, trainer_ckpt_path=None):
+        super().__init__()
+
+        self.llm = LLM.load(checkpoint_dir, tokenizer_dir=tokenizer_dir, distribute=None)
+        self.trainer_ckpt_path = trainer_ckpt_path
+
+    def setup(self, stage):
+        self.llm.trainer_setup(trainer_ckpt=self.trainer_ckpt_path)
+
+    def training_step(self, batch):
+        logits, loss = self.llm(input_ids=batch["input_ids"], target_ids=batch["labels"])
+        self.log("train_loss", loss, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch):
+        logits, loss = self.llm(input_ids=batch["input_ids"], target_ids=batch["labels"])
+        self.log("validation_loss", loss, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        warmup_steps = 10
+        optimizer = torch.optim.AdamW(self.llm.model.parameters(), lr=0.0002, weight_decay=0.0, betas=(0.9, 0.95))
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: step / warmup_steps)
+        return [optimizer], [scheduler]
+
+
+if __name__ == "__main__":
+    batch_size = 8
+    accumulate_grad_batches = 1
+
+    #########################################################
+    # Use case 1: Pretraining from random weights
+    #########################################################
+
+    llm = LLM.load("EleutherAI/pythia-160m", tokenizer_dir="EleutherAI/pythia-160m", init="random")
+    llm.save("pythia-160m-random-weights")
+    del llm
+
+    lit_model = LitLLM(checkpoint_dir="pythia-160m-random-weights", tokenizer_dir="EleutherAI/pythia-160m")
+    data = Alpaca2k()
+
+    data.connect(lit_model.llm.tokenizer, batch_size=batch_size, max_seq_length=512)
+
+    trainer = L.Trainer(
+        devices=1,
+        accelerator="cuda",
+        max_epochs=1,
+        accumulate_grad_batches=accumulate_grad_batches,
+        precision="bf16-true",
+    )
+    trainer.fit(lit_model, data)
+
+    lit_model.llm.model.to(lit_model.llm.preprocessor.device)
+    lit_model.llm.generate("hello world")
+
+    del lit_model
+
+    #############################################################################
+    # Use case 2: Continued pretraining / finetuning from downloaded checkpoint
+    #############################################################################
+
+    lit_model = LitLLM(checkpoint_dir="EleutherAI/pythia-160m")
+    data = Alpaca2k()
+
+    data.connect(lit_model.llm.tokenizer, batch_size=batch_size, max_seq_length=512)
+
+    trainer = L.Trainer(
+        devices=1,
+        accelerator="cuda",
+        max_epochs=1,
+        accumulate_grad_batches=accumulate_grad_batches,
+        precision="bf16-true",
+    )
+    trainer.fit(lit_model, data)
+
+    lit_model.llm.model.to(lit_model.llm.preprocessor.device)
+    lit_model.llm.generate("hello world")
+
+    del lit_model
+
+    #########################################################
+    # Use case 3: Resume training from Trainer checkpoint
+    #########################################################
+
+    import os
+
+    def find_latest_checkpoint(directory):
+        latest_checkpoint = None
+        latest_time = 0
+
+        for root, _, files in os.walk(directory):
+            for file in files:
+                if file.endswith(".ckpt"):
+                    file_path = os.path.join(root, file)
+                    file_time = os.path.getmtime(file_path)
+                    if file_time > latest_time:
+                        latest_time = file_time
+                        latest_checkpoint = file_path
+
+        return latest_checkpoint
+
+    lit_model = LitLLM(
+        checkpoint_dir="EleutherAI/pythia-160m", trainer_ckpt_path=find_latest_checkpoint("lightning_logs")
+    )
+
+    data.connect(lit_model.llm.tokenizer, batch_size=batch_size, max_seq_length=512)
+
+    trainer = L.Trainer(
+        devices=1,
+        accelerator="cuda",
+        max_epochs=1,
+        accumulate_grad_batches=accumulate_grad_batches,
+        precision="bf16-true",
+    )
+    trainer.fit(lit_model, data)
+
+    lit_model.llm.model.to(lit_model.llm.preprocessor.device)
+    lit_model.llm.generate("hello world")
+
+    #################################################################
+    # Use case 4: Resume training after saving a checkpoint manually
+    #################################################################
+
+    lit_model.llm.save("finetuned_checkpoint")
+    del lit_model
+    lit_model = LitLLM(checkpoint_dir="finetuned_checkpoint")
+
+    data.connect(lit_model.llm.tokenizer, batch_size=batch_size, max_seq_length=512)
+
+    trainer = L.Trainer(
+        devices=1,
+        accelerator="cuda",
+        max_epochs=1,
+        accumulate_grad_batches=accumulate_grad_batches,
+        precision="bf16-true",
+    )
+    trainer.fit(lit_model, data)
+
+    lit_model.llm.model.to(lit_model.llm.preprocessor.device)
+    lit_model.llm.generate("hello world")