Repository: whittle-org/whittle (language: python)

Issue:
adapt litgpt workflow for pre-training
provide a new pretrain.py script to support super-network training.


Discussion / hints:

Fix patch:
diff --git a/whittle/pretrain_super_network.py b/whittle/pretrain_super_network.py
new file mode 100644
index 0000000..b461319
--- /dev/null
+++ b/whittle/pretrain_super_network.py
@@ -0,0 +1,514 @@
+import math
+import pprint
+import time
+from datetime import timedelta
+from pathlib import Path
+from typing import Optional, Union
+
+import lightning as L
+import torch
+from lightning.fabric.strategies import FSDPStrategy
+from lightning.fabric.utilities.throughput import ThroughputMonitor, measure_flops
+from torch.utils.data import DataLoader
+from torchmetrics.aggregation import RunningMean
+from typing_extensions import Literal
+from litgpt import Tokenizer
+from litgpt.args import EvalArgs, TrainArgs
+from litgpt.config import name_to_config
+from litgpt.data import DataModule, TinyLlama
+from litgpt.model import Config
+from litgpt.utils import (
+    CycleIterator,
+    capture_hparams,
+    check_nvlink_connectivity,
+    choose_logger,
+    chunked_cross_entropy,
+    extend_checkpoint_dir,
+    find_resume_path,
+    get_default_supported_precision,
+    init_out_dir,
+    instantiate_torch_optimizer,
+    num_parameters,
+    parse_devices,
+)
+from litgpt.pretrain import (
+    get_dataloaders,
+    validate_args,
+    initialize_weights,
+    save_checkpoint,
+    validate,
+    get_lr,
+)
+
+from whittle.training_strategies import (
+    SandwichStrategy,
+    RandomStrategy,
+    StandardStrategy,
+)
+from whittle.training_strategies.base_strategy import BaseTrainingStrategy
+from whittle.sampling.random_sampler import RandomSampler
+from whittle.models.gpt import GPT
+from whittle.models.gpt.blocks import Block
+from syne_tune.config_space import randint, lograndint
+
+
+training_strategies_cls = {
+    "sandwich": SandwichStrategy,
+    "random": RandomStrategy,
+    "standard": StandardStrategy,
+}
+
+
+def get_search_space(config):
+    return {
+        "embed_dim": lograndint(1, config.n_embd),
+        "num_heads": randint(1, config.n_head),
+        "mlp_ratio": randint(1, 4),
+        "depth": randint(1, config.n_layer),
+    }
+
+
+def setup(
+    model_name: str,
+    model_config: Optional[Config] = None,
+    out_dir: Path = Path("../examples/gpt/out/pretrain"),
+    precision: Literal["bf16-true", "bf16-mixed", "32-true", None] = None,
+    initial_checkpoint_dir: Optional[Path] = None,
+    resume: Union[bool, Literal["auto"], Path] = False,
+    data: Optional[DataModule] = None,
+    train: TrainArgs = TrainArgs(
+        save_interval=1000,
+        log_interval=1,
+        global_batch_size=512,
+        micro_batch_size=4,
+        max_tokens=int(3e12),  # 3 trillion
+        max_norm=1.0,
+        min_lr=4e-5,
+        lr_warmup_steps=2000,
+        tie_embeddings=False,
+    ),
+    eval: EvalArgs = EvalArgs(interval=1000, max_iters=100),
+    optimizer: Union[str, dict] = "AdamW",
+    devices: Union[int, str] = "auto",
+    num_nodes: int = 1,
+    training_strategy: str = "sandwich",
+    tokenizer_dir: Optional[Path] = None,
+    logger_name: Literal["wandb", "tensorboard", "csv"] = "tensorboard",
+    seed: int = 42,
+):
+    """Pretrain a model.
+
+    Arguments:
+        model_name: The name of the model to pretrain. Choose from names in ``litgpt.config``. Use "list" to list the supported models.
+        model_config: A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
+            ``model_config``. Overrides the `model_name` if specified.
+        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.
+        precision: The precision to use for finetuning. Determines a compatible precision setting by default.
+        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``. An error will be raised if no checkpoint is found. Passing
+            ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
+        data: Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
+        train: Training-related arguments. See ``litgpt.args.TrainArgs`` for details.
+        eval: Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details.
+        optimizer: An optimizer name (such as "AdamW") or config.
+        training_strategy:
+        devices: How many devices/GPUs to use. Uses all GPUs by default.
+        num_nodes: How many nodes the code is being run on.
+        tokenizer_dir: Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
+            module require this.
+        logger_name: The name of the logger to send metrics to.
+        seed: The random seed to use for reproducibility.
+    """
+    if model_name == "list":
+        available_models = "\n".join(sorted(name_to_config))
+        print(f"Available values:\n{available_models}")
+        quit()
+
+    if initial_checkpoint_dir is not None:
+        initial_checkpoint_dir = extend_checkpoint_dir(initial_checkpoint_dir)
+
+    if tokenizer_dir is not None:
+        tokenizer_dir = extend_checkpoint_dir(tokenizer_dir)
+
+    if model_config is None:
+        # Support both model_name options: meta-llama/Meta-Llama-3-8B & Meta-Llama-3-8B
+        try:
+            model_config = Config.from_name(model_name)
+        except ValueError:
+            print(f"Model name {model_name} is not supported.\n")
+            available_models = "\n".join(sorted(name_to_config))
+            print(f"Available values:\n{available_models}")
+            quit()
+
+    assert training_strategy in training_strategies_cls, print(
+        f"Training strategy is {training_strategy}. Should be in {list(training_strategies_cls)}"
+    )
+
+    hparams = capture_hparams()
+    data = TinyLlama() if data is None else data
+
+    config = Config.from_name(model_name) if model_config is None else model_config
+    config.fix_head_size = True
+
+    precision = precision or get_default_supported_precision(training=True)
+    num_devices = parse_devices(devices)
+    out_dir = init_out_dir(out_dir)
+    # in case the dataset requires the Tokenizer
+    tokenizer = Tokenizer(tokenizer_dir) if tokenizer_dir is not None else None
+
+    logger = choose_logger(
+        logger_name,
+        out_dir,
+        name=f"pretrain-{config.name}",
+        resume=bool(resume),
+        log_interval=train.log_interval,
+    )
+
+    if num_devices * num_nodes > 1:
+        strategy = FSDPStrategy(
+            auto_wrap_policy={Block},
+            state_dict_type="full",
+            sharding_strategy="HYBRID_SHARD",
+        )
+    else:
+        strategy = "auto"
+
+    fabric = L.Fabric(
+        devices=num_devices,
+        num_nodes=num_nodes,
+        strategy=strategy,
+        precision=precision,
+        loggers=[logger],
+    )
+
+    if torch.cuda.is_available() and num_devices > 1:
+        check_nvlink_connectivity(fabric)
+
+    fabric.launch()
+
+    fabric.print(pprint.pformat(hparams))
+    if logger_name in ("tensorboard", "wandb"):
+        fabric.logger.log_hyperparams(hparams)
+
+    main(
+        fabric,
+        num_devices,
+        seed,
+        initial_checkpoint_dir,
+        resume,
+        config,
+        data,
+        out_dir,
+        tokenizer_dir,
+        tokenizer,
+        train,
+        eval,
+        optimizer,
+        training_strategy,
+    )
+
+
+def fit(
+    fabric: L.Fabric,
+    devices: int,
+    state: dict,
+    train_dataloader: DataLoader,
+    val_dataloader: DataLoader,
+    out_dir: Path,
+    tokenizer_dir: Optional[Path],
+    train: TrainArgs,
+    eval: EvalArgs,
+    training_strategy: BaseTrainingStrategy,
+) -> None:
+    model = state["model"]
+    optimizer = state["optimizer"]
+
+    if eval.initial_validation:
+        model.reset_super_network()
+        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
+        val_loss = f"{val_loss:.3f}"
+    else:
+        model.reset_super_network()
+        fabric.print("Verifying settings ...")
+        validate(
+            fabric, model, val_dataloader, max_iters=2, verbose=False
+        )  # sanity check
+        val_loss = "n/a"
+
+    throughput = ThroughputMonitor(fabric, window_size=5)
+
+    with torch.device("meta"):
+        meta_model = GPT(model.config)
+        x = torch.randint(0, 1, (train.micro_batch_size, meta_model.max_seq_length))
+
+        def model_fwd():
+            return meta_model(x)  # noqa: F821
+
+        def model_loss(y):
+            return chunked_cross_entropy(y, x, chunk_size=0)  # noqa: F821
+
+        measured_flops = measure_flops(meta_model, model_fwd, model_loss)
+        fabric.print(
+            f"Measured TFLOPs: {measured_flops * fabric.world_size / 1e12:.2f}"
+        )
+        del meta_model, x
+
+    max_tokens_per_device = train.max_tokens // fabric.world_size
+    tokens_per_iter = train.micro_batch_size * model.max_seq_length
+    max_iters = max_tokens_per_device // tokens_per_iter
+    log_iter_interval = train.log_interval * train.gradient_accumulation_iters(devices)
+    initial_iter = state["iter_num"]
+    train_iterator = CycleIterator(train_dataloader)
+
+    running_loss = RunningMean(
+        window=train.gradient_accumulation_iters(devices), sync_on_compute=False
+    ).to(fabric.device)
+    fabric.barrier()
+    total_t0 = time.perf_counter()
+
+    warmup_iters = train.warmup_iters(devices, max_iters, train_dataloader)
+
+    for train_data in train_iterator:
+        if state["iter_num"] >= max_iters:
+            break
+
+        # determine and set the learning rate for this iteration
+        lr = get_lr(
+            optimizer.defaults["lr"],
+            state["iter_num"],
+            warmup_iters,
+            max_iters,
+            train.min_lr,
+        )
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
+
+        state["iter_num"] += 1
+        iter_t0 = time.perf_counter()
+
+        input_ids = train_data[:, 0 : model.max_seq_length].contiguous().long()
+        targets = train_data[:, 1 : (model.max_seq_length + 1)].contiguous().long()
+
+        is_accumulating = (
+            state["iter_num"] % train.gradient_accumulation_iters(devices) != 0
+        )
+
+        if hasattr(training_strategy, "random_samples"):
+            # if we update multiple sub-networks in each step, we need to further scale the gradient
+            scale_loss = 1 / (
+                train.gradient_accumulation_iters(devices)
+                * training_strategy.random_samples
+            )
+        else:
+            scale_loss = 1 / train.gradient_accumulation_iters(devices)
+
+        with fabric.no_backward_sync(model, enabled=is_accumulating):
+            loss = training_strategy(model, input_ids, targets, scale_loss=scale_loss)
+        #            fabric.backward(loss / train.gradient_accumulation_iters(devices))
+
+        running_loss.update(loss)
+
+        if not is_accumulating:
+            fabric.clip_gradients(model, optimizer, max_norm=train.max_norm)
+            optimizer.step()
+            optimizer.zero_grad()
+            state["step_count"] += 1
+
+        if state["iter_num"] % log_iter_interval == 0:
+            loss = (
+                running_loss.compute().item()
+            )  # expensive device-to-host synchronization
+            t1 = time.perf_counter()
+            throughput.update(
+                time=(t1 - total_t0),
+                flops=(measured_flops * log_iter_interval),
+                batches=state["iter_num"],
+                samples=(state["iter_num"] * train.micro_batch_size),
+                lengths=(
+                    state["iter_num"] * train.micro_batch_size * model.max_seq_length
+                ),
+            )
+            metrics = {
+                "loss": loss,
+                "iter": state["iter_num"],
+                "step": state["step_count"],
+                "epoch": train_iterator.epoch,
+                "iter_time": t1 - iter_t0,
+                "remaining_time": (
+                    (t1 - total_t0)
+                    / (state["iter_num"] - initial_iter)
+                    * (max_iters - state["iter_num"])
+                ),
+                "tokens": state["iter_num"]
+                * train.micro_batch_size
+                * model.max_seq_length,
+                "total_tokens": (
+                    state["iter_num"]
+                    * train.micro_batch_size
+                    * model.max_seq_length
+                    * fabric.world_size
+                ),
+                "learning_rate": lr,
+            }
+            if isinstance(val_loss, float):
+                val_loss = f"{val_loss:.3f}"
+            fabric.print(
+                f"Epoch {metrics['epoch']+1} | iter {metrics['iter']} step {metrics['step']} |"
+                f" loss train: {metrics['loss']:.3f},"
+                f" val: {val_loss} |"
+                f" iter time: {metrics['iter_time'] * 1000:.2f} ms"
+                f"{' (step)' if not is_accumulating else ''}"
+                f" remaining time: {timedelta(seconds=int(metrics['remaining_time']))!s}"
+            )
+
+            throughput_metrics = throughput.compute()
+            metrics.update(throughput_metrics)
+            fabric.log_dict(metrics, step=state["iter_num"] - 1)
+
+        if (
+            val_dataloader is not None
+            and not is_accumulating
+            and state["step_count"] % eval.interval == 0
+        ):
+            t0 = time.perf_counter()
+            model.reset_super_network()
+            val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
+            val_loss = val_loss.item()
+            td = time.perf_counter() - t0
+
+            fabric.print(
+                f"iter {state['iter_num']}: val loss {val_loss:.4f}, val time: {td * 1000:.2f} ms"
+            )
+            metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
+            fabric.log_dict(metrics, step=state["iter_num"] - 1)
+            fabric.barrier()
+
+        if (
+            train.save_interval is not None
+            and not is_accumulating
+            and state["step_count"] % train.save_interval == 0
+        ):
+            save_checkpoint(
+                fabric,
+                state,
+                tokenizer_dir,
+                out_dir / f"step-{state['step_count']:08d}" / "lit_model.pth",
+            )
+
+    # Final validation
+    if eval.final_validation:
+        model.reset_super_network()
+        val_loss = validate(fabric, model, val_dataloader, max_iters=eval.max_iters)
+        metrics = {"val_loss": val_loss, "val_ppl": math.exp(val_loss)}
+        fabric.log_dict(metrics, step=state["iter_num"])
+        fabric.print(
+            f"Final evaluation | val loss: {val_loss.item():.3f} | val ppl: {math.exp(val_loss):.3f}"
+        )
+
+
+def main(
+    fabric: L.Fabric,
+    devices: int,
+    seed: int,
+    initial_checkpoint_dir: Optional[Path],
+    resume: Union[bool, Literal["auto"], Path],
+    config: Config,
+    data: DataModule,
+    out_dir: Path,
+    tokenizer_dir: Optional[Path],
+    tokenizer: Optional[Tokenizer],
+    train: TrainArgs,
+    eval: EvalArgs,
+    optimizer: Union[str, dict],
+    training_strategy: str,
+) -> None:
+    validate_args(train, eval, initial_checkpoint_dir, resume)
+
+    if fabric.global_rank == 0:
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+    fabric.seed_everything(seed)  # same seed for every process to init model (FSDP)
+
+    t0 = time.perf_counter()
+    with fabric.init_module(empty_init=True):
+        model = GPT(config)
+
+    initialize_weights(fabric, model, n_layer=config.n_layer, n_embd=config.n_embd)
+
+    if train.tie_embeddings:
+        model.transformer.wte.weight = model.lm_head.weight
+    if train.max_seq_length:
+        model.max_seq_length = train.max_seq_length
+
+    fabric.print(f"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.")
+    fabric.print(f"Total parameters: {num_parameters(model):,}")
+
+    model = torch.compile(model)
+    model = fabric.setup(model)
+
+    extra_kwargs = {"fused": fabric.device.type == "cuda"}
+    optimizer = instantiate_torch_optimizer(
+        optimizer, model.parameters(), **extra_kwargs
+    )
+    optimizer = fabric.setup_optimizers(optimizer)
+
+    sampler = RandomSampler(config_space=get_search_space(config), seed=seed)
+    training_strategy_kwargs = {
+        "loss_function": chunked_cross_entropy,
+        "sampler": sampler,
+        "fabric": fabric,
+    }
+    strategy = training_strategies_cls[training_strategy](**training_strategy_kwargs)
+
+    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,
+        "train_dataloader": train_dataloader,
+        "iter_num": 0,
+        "step_count": 0,
+    }
+
+    resume = find_resume_path(resume, out_dir)
+    if resume:
+        fabric.print(f"Resuming training from {resume}")
+        fabric.load(resume, state)
+
+    train_time = time.perf_counter()
+    fit(
+        fabric,
+        devices,
+        state,
+        train_dataloader,
+        val_dataloader,
+        out_dir,
+        tokenizer_dir,
+        train,
+        eval,
+        strategy,
+    )
+
+    # Save final checkpoint
+    save_checkpoint(fabric, state, tokenizer_dir, out_dir / "final" / "lit_model.pth")
+
+    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")
+
+
+if __name__ == "__main__":
+    from jsonargparse import CLI
+
+    CLI(setup)


Test patch:
diff --git a/test/conftest.py b/test/conftest.py
new file mode 100644
index 0000000..ba7f7fa
--- /dev/null
+++ b/test/conftest.py
@@ -0,0 +1,187 @@
+# Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+
+import os
+import shutil
+from pathlib import Path
+from typing import Optional
+
+import pytest
+import torch
+from lightning.fabric.utilities.testing import _runif_reasons
+from lightning_utilities.core.imports import RequirementCache
+
+
+@pytest.fixture()
+def fake_checkpoint_dir(tmp_path):
+    os.chdir(tmp_path)
+    checkpoint_dir = tmp_path / "checkpoints" / "tmp"
+    checkpoint_dir.mkdir(parents=True)
+    (checkpoint_dir / "lit_model.pth").touch()
+    (checkpoint_dir / "model_config.yaml").touch()
+    (checkpoint_dir / "tokenizer.json").touch()
+    (checkpoint_dir / "tokenizer_config.json").touch()
+    return checkpoint_dir
+
+
+class TensorLike:
+    def __eq__(self, other):
+        return isinstance(other, torch.Tensor)
+
+
+@pytest.fixture()
+def tensor_like():
+    return TensorLike()
+
+
+class FloatLike:
+    def __eq__(self, other):
+        return not isinstance(other, int) and isinstance(other, float)
+
+
+@pytest.fixture()
+def float_like():
+    return FloatLike()
+
+
+@pytest.fixture(autouse=True)
+def restore_default_dtype():
+    # just in case
+    torch.set_default_dtype(torch.float32)
+
+
+@pytest.fixture(autouse=True)
+def destroy_process_group():
+    yield
+
+    import torch.distributed
+
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+
+
+class MockTokenizer:
+    """A dummy tokenizer that encodes each character as its ASCII code."""
+
+    bos_id = 0
+    eos_id = 1
+
+    def encode(
+        self,
+        text: str,
+        bos: Optional[bool] = None,
+        eos: bool = False,
+        max_length: int = -1,
+    ) -> torch.Tensor:
+        output = []
+        if bos:
+            output.append(self.bos_id)
+        output.extend([ord(c) for c in text])
+        if eos:
+            output.append(self.eos_id)
+        output = output[:max_length] if max_length > 0 else output
+        return torch.tensor(output)
+
+    def decode(self, tokens: torch.Tensor) -> str:
+        return "".join(chr(int(t)) for t in tokens.tolist())
+
+
+@pytest.fixture()
+def mock_tokenizer():
+    return MockTokenizer()
+
+
+@pytest.fixture()
+def alpaca_path(tmp_path):
+    file = Path(__file__).parent / "data" / "fixtures" / "alpaca.json"
+    shutil.copyfile(file, tmp_path / "alpaca.json")
+    return tmp_path / "alpaca.json"
+
+
+@pytest.fixture()
+def dolly_path(tmp_path):
+    file = Path(__file__).parent / "data" / "fixtures" / "dolly.json"
+    shutil.copyfile(file, tmp_path / "dolly.json")
+    return tmp_path / "dolly.json"
+
+
+@pytest.fixture()
+def longform_path(tmp_path):
+    path = tmp_path / "longform"
+    path.mkdir()
+    for split in ("train", "val"):
+        file = Path(__file__).parent / "data" / "fixtures" / f"longform_{split}.json"
+        shutil.copyfile(file, path / f"{split}.json")
+    return path
+
+
+def RunIf(thunder: Optional[bool] = None, **kwargs):
+    reasons, marker_kwargs = _runif_reasons(**kwargs)
+
+    if thunder is not None:
+        thunder_available = bool(RequirementCache("lightning-thunder", "thunder"))
+        if thunder and not thunder_available:
+            reasons.append("Thunder")
+        elif not thunder and thunder_available:
+            reasons.append("not Thunder")
+
+    return pytest.mark.skipif(
+        condition=len(reasons) > 0,
+        reason=f"Requires: [{' + '.join(reasons)}]",
+        **marker_kwargs,
+    )
+
+
+# https://github.com/Lightning-AI/lightning/blob/6e517bd55b50166138ce6ab915abd4547702994b/tests/tests_fabric/conftest.py#L140
+def pytest_collection_modifyitems(
+    items: list[pytest.Function], config: pytest.Config
+) -> None:
+    initial_size = len(items)
+    conditions = []
+    filtered, skipped = 0, 0
+
+    options = {
+        "standalone": "PL_RUN_STANDALONE_TESTS",
+        "min_cuda_gpus": "PL_RUN_CUDA_TESTS",
+    }
+    if (
+        os.getenv(options["standalone"], "0") == "1"
+        and os.getenv(options["min_cuda_gpus"], "0") == "1"
+    ):
+        # special case: we don't have a CPU job for standalone tests, so we shouldn't run only cuda tests.
+        # by deleting the key, we avoid filtering out the CPU tests
+        del options["min_cuda_gpus"]
+
+    for kwarg, env_var in options.items():
+        # this will compute the intersection of all tests selected per environment variable
+        if os.getenv(env_var, "0") == "1":
+            conditions.append(env_var)
+            for i, test in reversed(
+                list(enumerate(items))
+            ):  # loop in reverse, since we are going to pop items
+                already_skipped = any(
+                    marker.name == "skip" for marker in test.own_markers
+                )
+                if already_skipped:
+                    # the test was going to be skipped anyway, filter it out
+                    items.pop(i)
+                    skipped += 1
+                    continue
+                has_runif_with_kwarg = any(
+                    marker.name == "skipif" and marker.kwargs.get(kwarg)
+                    for marker in test.own_markers
+                )
+                if not has_runif_with_kwarg:
+                    # the test has `@RunIf(kwarg=True)`, filter it out
+                    items.pop(i)
+                    filtered += 1
+
+    if config.option.verbose >= 0 and (filtered or skipped):
+        writer = config.get_terminal_writer()
+        writer.write(
+            f"\nThe number of tests has been filtered from {initial_size} to {initial_size - filtered} after the"
+            f" filters {conditions}.\n{skipped} tests are marked as unconditional skips.\nIn total,"
+            f" {len(items)} tests will run.\n",
+            flush=True,
+            bold=True,
+            purple=True,  # oh yeah, branded pytest messages
+        )
diff --git a/test/test_pretrain_super_network.py b/test/test_pretrain_super_network.py
new file mode 100644
index 0000000..23403c7
--- /dev/null
+++ b/test/test_pretrain_super_network.py
@@ -0,0 +1,137 @@
+"""
+Adapted from the original LitGPT code.
+"""
+
+import os
+from contextlib import redirect_stdout
+from io import StringIO
+from unittest import mock
+from unittest.mock import ANY, Mock
+
+import torch
+from torch.utils.data import DataLoader
+
+from litgpt.args import EvalArgs, TrainArgs
+from litgpt.config import Config
+
+from whittle import pretrain_super_network
+
+from test.conftest import RunIf
+
+
+@RunIf(min_cuda_gpus=1, standalone=True)
+@mock.patch("litgpt.pretrain.save_hyperparameters")
+def test_training_strategies(_, tmp_path):
+    model_config = Config(
+        block_size=2, n_layer=2, n_embd=4, n_head=2, padded_vocab_size=8
+    )
+
+    dataset = torch.tensor([[0, 1, 2], [3, 4, 5], [0, 1, 2]])
+    dataloader = DataLoader(dataset)
+    pretrain_super_network.get_dataloaders = Mock(return_value=(dataloader, dataloader))
+
+    for strategy in ("standard", "random", "sandwich"):
+        pretrain_super_network.setup(
+            "pythia-14m",
+            devices=1,
+            optimizer="RMSprop",
+            training_strategy=strategy,
+            model_config=model_config,
+            out_dir=tmp_path,
+            train=TrainArgs(
+                global_batch_size=2,
+                max_tokens=16,
+                save_interval=1,
+                micro_batch_size=1,
+                max_norm=1.0,
+            ),
+            eval=EvalArgs(interval=1, max_iters=1, final_validation=False),
+        )
+
+
+@RunIf(min_cuda_gpus=1, 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"})
+# If we were to use `save_hyperparameters()`, we would have to patch `sys.argv` or otherwise
+# the CLI would capture pytest args, but unfortunately patching would mess with subprocess
+# launching, so we need to mock `save_hyperparameters()`
+@mock.patch("litgpt.pretrain.save_hyperparameters")
+def test_pretrain(_, tmp_path):
+    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_super_network.get_dataloaders = Mock(return_value=(dataloader, dataloader))
+
+    out_dir = tmp_path / "out"
+    stdout = StringIO()
+    with redirect_stdout(stdout):
+        pretrain_super_network.setup(
+            "pythia-14m",
+            devices=1,
+            model_config=model_config,
+            out_dir=out_dir,
+            train=TrainArgs(
+                global_batch_size=2,
+                max_tokens=16,
+                save_interval=1,
+                micro_batch_size=1,
+                max_norm=1.0,
+            ),
+            eval=EvalArgs(interval=1, max_iters=1, final_validation=False),
+        )
+
+        # tmp_path is not the same across all ranks, run assert only on rank 0
+    out_dir_contents = set(os.listdir(out_dir))
+    checkpoint_dirs = {
+        "step-00000001",
+        "step-00000002",
+        "step-00000003",
+        "step-00000004",
+        "final",
+    }
+    assert checkpoint_dirs.issubset(out_dir_contents)
+    assert all((out_dir / p).is_dir() for p in checkpoint_dirs)
+    for checkpoint_dir in checkpoint_dirs:
+        # the `tokenizer_dir` is None by default, so only 'lit_model.pth' shows here
+        assert set(os.listdir(out_dir / checkpoint_dir)) == {
+            "lit_model.pth",
+            "model_config.yaml",
+        }
+
+    assert (out_dir / "logs" / "tensorboard" / "version_0").is_dir()
+
+    # logs only appear on rank 0
+    logs = stdout.getvalue()
+    assert logs.count("(step)") == 4
+    assert logs.count("val loss") == 4
+    assert "Total parameters: 1,888" in logs
+
+
+@RunIf(min_cuda_gpus=1, 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")
+# See comment in `test_pretrain` why we need to mock `save_hyperparameters()`
+@mock.patch("litgpt.pretrain.save_hyperparameters")
+def test_initial_checkpoint_dir(_, load_mock, tmp_path):
+    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_super_network.get_dataloaders = Mock(return_value=(dataloader, dataloader))
+    pretrain_super_network.fit = Mock()
+
+    pretrain_super_network.setup(
+        "pythia-14m",
+        initial_checkpoint_dir=tmp_path,
+        devices=1,
+        model_config=model_config,
+        out_dir=tmp_path,
+    )
+
+    load_mock.assert_called_once_with(tmp_path / "lit_model.pth", ANY)