Enable continued pretraining (#1109)

Author: awaelchli

config_hub/pretrain/debug.yaml

@@ -11,6 +11,10 @@ model_config:
 # /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
 out_dir: out/pretrain/debug
 
+# Optional path to a checkpoint directory to initialize the model from.
+# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
+initial_checkpoint_dir:
+
 # Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
 # from the latest checkpoint in ``out_dir``. (type: Union[bool, Path], default: False)
 resume: false

config_hub/pretrain/tinyllama.yaml

@@ -11,6 +11,10 @@ model_config:
 # /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
 out_dir: out/pretrain/tiny-llama
 
+# Optional path to a checkpoint directory to initialize the model from.
+# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
+initial_checkpoint_dir:
+
 # Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
 # from the latest checkpoint in ``out_dir``. (type: Union[bool, Path], default: False)
 resume: false

config_hub/pretrain/tinystories.yaml

@@ -27,6 +27,10 @@ model_config:
 # /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
 out_dir: out/pretrain/stories15M
 
+# Optional path to a checkpoint directory to initialize the model from.
+# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
+initial_checkpoint_dir:
+
 # Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
 # from the latest checkpoint in ``out_dir``. (type: Union[bool, Path], default: False)
 resume: false

litgpt/pretrain.py

@@ -39,6 +39,7 @@ def setup(
     model_name: Optional[str] = None,
     model_config: Optional[Config] = None,
     out_dir: Path = Path("out/pretrain"),
+    initial_checkpoint_dir: Optional[Path] = None,
     resume: Union[bool, Path] = False,
     data: Optional[DataModule] = None,
     train: TrainArgs = TrainArgs(
@@ -71,6 +72,8 @@ def setup(
             ``model_config``.
         out_dir: Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
             /teamspace/jobs/<job-name>/share.
+        initial_checkpoint_dir: Optional path to a checkpoint directory to initialize the model from.
+            Useful for continued pretraining. Mutually exclusive with ``resume``.
         resume: Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
             from the latest checkpoint in ``out_dir``.
         data: Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
@@ -107,13 +110,14 @@ def setup(
     if logger_name in ("tensorboard", "wandb"):
         fabric.logger.log_hyperparams(hparams)
 
-    main(fabric, devices, seed, resume, config, data, out_dir, tokenizer_dir, tokenizer, train, eval)
+    main(fabric, devices, seed, initial_checkpoint_dir, resume, config, data, out_dir, tokenizer_dir, tokenizer, train, eval)
 
 
 def main(
     fabric: L.Fabric,
     devices: int,
     seed: int,
+    initial_checkpoint_dir: Optional[Path],
     resume: Union[bool, Path],
     config: Config,
     data: DataModule,
@@ -123,7 +127,7 @@ def main(
     train: TrainArgs,
     eval: EvalArgs,
 ) -> None:
-    validate_args(train, eval)
+    validate_args(train, eval, initial_checkpoint_dir, resume)
 
     if fabric.global_rank == 0:
         out_dir.mkdir(parents=True, exist_ok=True)
@@ -157,6 +161,9 @@ def main(
     train_dataloader, val_dataloader = get_dataloaders(fabric, data, tokenizer, train, model.max_seq_length)
     train_dataloader, val_dataloader = fabric.setup_dataloaders(train_dataloader, val_dataloader)
 
+    if initial_checkpoint_dir:
+        fabric.load_raw(initial_checkpoint_dir / "lit_model.pth", model)
+
     state = {
         "model": model,
         "optimizer": optimizer,
@@ -376,7 +383,7 @@ def init_out_dir(out_dir: Path) -> Path:
     return out_dir
 
 
-def validate_args(train: TrainArgs, eval: EvalArgs) -> None:
+def validate_args(train: TrainArgs, eval: EvalArgs, initial_checkpoint_dir, resume) -> None:
     issues = []
     unsupported = [
         (train, ["max_steps", "epochs"]),
@@ -391,6 +398,8 @@ def validate_args(train: TrainArgs, eval: EvalArgs) -> None:
         for name in names:
             if getattr(args, name) is None:
                 issues.append(f"{__file__} requires the {name!r} argument. This is set in {args}")
+    if initial_checkpoint_dir and resume:
+        issues.append("Can't provide both `--resume` and `--initial_checkpoint_dir`. Choose one.")
     if issues:
         raise ValueError("\n".join(issues))
 

tests/test_pretrain.py

@@ -5,7 +5,7 @@
 from io import StringIO
 from pathlib import Path
 from unittest import mock
-from unittest.mock import Mock
+from unittest.mock import Mock, ANY
 
 import pytest
 import torch
@@ -63,6 +63,32 @@ def test_pretrain(_, tmp_path):
     torch.distributed.barrier()
 
 
+
+@RunIf(min_cuda_gpus=2, standalone=True)
+# Set CUDA_VISIBLE_DEVICES for FSDP hybrid-shard, if fewer GPUs are used than are available
+@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"})
+@mock.patch("litgpt.pretrain.L.Fabric.load_raw")
+def test_initial_checkpoint_dir(load_mock, tmp_path):
+    from litgpt import pretrain
+    from litgpt.config import Config
+
+    model_config = Config(block_size=2, n_layer=2, n_embd=8, n_head=4, padded_vocab_size=8)
+
+    dataset = torch.tensor([[0, 1, 2], [3, 4, 5], [0, 1, 2]])
+    dataloader = DataLoader(dataset)
+    pretrain.get_dataloaders = Mock(return_value=(dataloader, dataloader))
+    pretrain.fit = Mock()
+
+    pretrain.setup(
+        initial_checkpoint_dir=tmp_path,
+        devices=2,
+        model_config=model_config,
+        out_dir=tmp_path,
+    )
+
+    load_mock.assert_called_once_with(tmp_path / "lit_model.pth", ANY)
+
+
 def test_pretrain_model_name_and_config():
     from litgpt import pretrain
     from litgpt.config import Config