(4/n) Data Refactor - Finetuning Scripts

.gitignore

@@ -9,6 +9,7 @@ build
 # data
 data
 datasets
+!lit_gpt/datasets
 checkpoints
 out
 wandb

finetune/full.py

@@ -7,8 +7,10 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
-import lightning as L
 import torch
+from torch.utils.data import DataLoader
+
+import lightning as L
 from lightning.fabric.loggers import CSVLogger
 from lightning.fabric.strategies import FSDPStrategy
 from torchmetrics import RunningMean
@@ -21,24 +23,23 @@
 from lit_gpt.args import EvalArgs, IOArgs, TrainArgs
 from lit_gpt.model import GPT, Block, Config
 from lit_gpt.tokenizer import Tokenizer
+from lit_gpt.datasets import Alpaca, apply_prompt_template
 from lit_gpt.utils import (
     CLI,
     check_valid_checkpoint_dir,
     chunked_cross_entropy,
     get_default_supported_precision,
     load_checkpoint,
     num_parameters,
+    CycleIterator,
 )
-from scripts.prepare_alpaca import generate_prompt
 
 
 def setup(
     precision: Optional[str] = None,
     devices: int = 1,
     resume: Union[bool, Path] = False,
     io: IOArgs = IOArgs(
-        train_data_dir=Path("data/alpaca"),
-        val_data_dir=Path("data/alpaca"),
         checkpoint_dir=Path("checkpoints/stabilityai/stablelm-base-alpha-3b"),
         out_dir=Path("out/full/alpaca"),
     ),
@@ -49,7 +50,6 @@ def setup(
         micro_batch_size=1,
         lr_warmup_steps=100,
         epochs=5,
-        epoch_size=50000,
         learning_rate=3e-3,
         max_seq_length=None,
     ),
@@ -85,7 +85,10 @@ def main(
 ) -> None:
     validate_args(io, train, eval)
 
-    steps_per_epoch = train.epoch_size // devices // train.batch_size(devices)
+    datamodule = Alpaca(io.checkpoint_dir)
+    train_dataloader, val_dataloader = get_dataloaders(fabric, datamodule)
+
+    steps_per_epoch = len(train_dataloader) // train.gradient_accumulation_iters(devices)
     lr_max_steps = train.epochs * steps_per_epoch
 
     check_valid_checkpoint_dir(io.checkpoint_dir)
@@ -95,9 +98,6 @@ def main(
     if fabric.global_rank == 0:
         os.makedirs(io.out_dir, exist_ok=True)
 
-    train_data = torch.load(io.train_data_dir / "train.pt")
-    val_data = torch.load(io.val_data_dir / "test.pt")
-
     checkpoint_path = io.checkpoint_dir / "lit_model.pth"
     fabric.print(f"Loading model {str(checkpoint_path)!r} with {config.__dict__}")
     with fabric.init_module(empty_init=(devices > 1)):
@@ -121,10 +121,8 @@ def main(
     else:
         load_checkpoint(fabric, state["model"], checkpoint_path)
 
-    fabric.seed_everything(1337 + fabric.global_rank)
-
     train_time = time.perf_counter()
-    fit(fabric, state, train_data, val_data, devices, resume, io, train, eval)
+    fit(fabric, state, train_dataloader, val_dataloader, devices, resume, io, train, eval)
     fabric.print(f"Training time: {(time.perf_counter()-train_time):.2f}s")
     if fabric.device.type == "cuda":
         fabric.print(f"Memory used: {torch.cuda.max_memory_allocated() / 1e9:.02f} GB")
@@ -136,8 +134,8 @@ def main(
 def fit(
     fabric: L.Fabric,
     state: Dict,
-    train_data: List[Dict],
-    val_data: List[Dict],
+    train_dataloader: DataLoader,
+    val_dataloader: DataLoader,
     devices: int,
     resume: Union[bool, Path],
     io: IOArgs,
@@ -155,14 +153,15 @@ def fit(
         f" {model.max_seq_length} and context length is {model.config.block_size}"
     )
 
-    validate(fabric, model, val_data, tokenizer, dataclasses.replace(eval, max_iters=2), train)  # sanity check
+    validate(fabric, model, val_dataloader, tokenizer, dataclasses.replace(eval, max_iters=2), train)  # sanity check
     initial_iter = state["iter_num"]
+    train_iterator = CycleIterator(train_dataloader)
 
     # resume data loader state by fast-forwarding through all seen batches
     if resume:
         resume_t0 = time.perf_counter()
         for resume_iter in range(initial_iter):
-            get_batch(fabric, train_data, None)
+            next(train_iterator)
             if resume_iter % 1000 == 0:
                 fabric.print(f"Resuming dataset: {resume_iter} / {initial_iter}")
         fabric.barrier()
@@ -176,16 +175,11 @@ def fit(
     )
     fabric.barrier()
 
-    for state["iter_num"] in range(state["iter_num"] + 1, train.max_iters(devices) + 1):
+    while state["iter_num"] <= train.max_iters(devices) and train_iterator.epoch < train.epochs:
+        state["iter_num"] += 1
         iter_t0 = time.perf_counter()
-
-        input_ids, targets = get_batch(
-            fabric,
-            train_data,
-            train.micro_batch_size,
-            train.max_seq_length,
-            longest_seq_ix if state["iter_num"] == 1 else None,
-        )
+        batch = next(train_iterator)
+        input_ids, targets = batch["input_ids"], batch["labels"]
 
         is_accumulating = state["iter_num"] % train.gradient_accumulation_iters(devices) != 0
         with fabric.no_backward_sync(model, enabled=is_accumulating):
@@ -224,7 +218,7 @@ def fit(
 
         if not is_accumulating and state["step_count"] % eval.interval == 0:
             t0 = time.perf_counter()
-            val_loss = validate(fabric, model, val_data, tokenizer, eval, train)
+            val_loss = validate(fabric, model, val_dataloader, tokenizer, eval, train)
             t1 = time.perf_counter() - t0
             fabric.print(f"iter {state['iter_num']}: val loss {val_loss.item():.4f}, val time: {t1 * 1000:.2f} ms")
             metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
@@ -239,22 +233,23 @@ def fit(
 # FSDP has issues with `inference_mode`
 @torch.no_grad()
 def validate(
-    fabric: L.Fabric, model: GPT, val_data: List[Dict], tokenizer: Tokenizer, eval: EvalArgs, train: TrainArgs
+    fabric: L.Fabric, model: GPT, val_dataloader: DataLoader, tokenizer: Tokenizer, eval: EvalArgs, train: TrainArgs
 ) -> torch.Tensor:
     fabric.print("Validating ...")
     model.eval()
     losses = torch.zeros(eval.max_iters)
+    val_iterator = iter(val_dataloader)
     for k in range(eval.max_iters):
-        input_ids, targets = get_batch(fabric, val_data, train.micro_batch_size, train.max_seq_length)
+        batch = next(val_iterator)
+        input_ids, targets = batch["input_ids"], batch["labels"]
         logits = model(input_ids)
         losses[k] = chunked_cross_entropy(logits[..., :-1, :], targets[..., 1:], chunk_size=0)
     val_loss = losses.mean()
 
     # produce an example:
     instruction = "Recommend a movie for me to watch during the weekend and explain the reason."
     fabric.print(instruction)
-    sample = {"instruction": instruction, "input": ""}
-    prompt = generate_prompt(sample)
+    prompt = apply_prompt_template(val_dataloader.dataset.prompt_template, {"instruction": instruction})
     encoded = tokenizer.encode(prompt, device=fabric.device)
     with fabric.init_tensor():
         # do not set `max_seq_length=max_returned_token` because memory is not a concern here
@@ -270,51 +265,24 @@ def validate(
     return val_loss
 
 
-def get_batch(
-    fabric: L.Fabric,
-    data: List[Dict],
-    micro_batch_size: int,
-    max_seq_length: Optional[int],
-    longest_seq_ix: Optional[int] = None,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    ix = torch.randint(len(data), (micro_batch_size,))
-    if longest_seq_ix is not None:
-        # force the longest sample at the beginning so potential OOMs happen right away
-        ix[0] = longest_seq_ix
-
-    input_ids = [data[i]["input_ids"].type(torch.int64) for i in ix]
-    labels = [data[i]["labels"].type(torch.int64) for i in ix]
-
-    # this could be `longest_seq_length` to have a fixed size for all batches
-    max_len = max(len(s) for s in input_ids)
-
-    def pad_right(x, pad_id):
-        # pad right based on the longest sequence
-        n = max_len - len(x)
-        return torch.cat((x, torch.full((n,), pad_id, dtype=x.dtype)))
-
-    x = torch.stack([pad_right(x, pad_id=0) for x in input_ids])
-    y = torch.stack([pad_right(x, pad_id=-1) for x in labels])
-
-    # Truncate if needed
-    if max_seq_length:
-        x = x[:, :max_seq_length]
-        y = y[:, :max_seq_length]
-
-    if fabric.device.type == "cuda" and x.device.type == "cpu":
-        x, y = fabric.to_device((x.pin_memory(), y.pin_memory()))
-    else:
-        x, y = fabric.to_device((x, y))
-    return x, y
-
-
 def get_lr_scheduler(optimizer, warmup_steps: int, max_steps: int):
     # linear warmup followed by cosine annealing
     scheduler1 = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: step / warmup_steps)
     scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=(max_steps - warmup_steps))
     return torch.optim.lr_scheduler.SequentialLR(optimizer, [scheduler1, scheduler2], milestones=[warmup_steps])
 
 
+def get_dataloaders(fabric: L.Fabric, datamodule: L.LightningDataModule) -> Tuple[DataLoader, DataLoader]:
+    if fabric.global_rank == 0:
+        datamodule.prepare_data()
+    fabric.barrier()
+    datamodule.setup()
+    train_dataloader = datamodule.train_dataloader()
+    val_dataloader = datamodule.val_dataloader()
+    train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)
+    return train_dataloader, val_dataloader
+
+
 def get_longest_seq_length(data: List[Dict]) -> Tuple[int, int]:
     # find out the minimum max_seq_length required during fine-tuning (saves memory!)
     lengths = [len(d["input_ids"]) for d in data]

lit_gpt/datasets/__init__.py

@@ -0,0 +1,5 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+from lit_gpt.datasets.base import SFTDataset, apply_prompt_template, sft_collate_fn
+from lit_gpt.datasets.alpaca import Alpaca
+from lit_gpt.datasets.lima import LIMA