diff --git a/3rdparty/NeMo b/3rdparty/NeMo
index e2b0f0ead1..06e6703aaa 160000
--- a/3rdparty/NeMo
+++ b/3rdparty/NeMo
@@ -1 +1 @@
-Subproject commit e2b0f0ead13be29476c047dfb49ad49f85a849bb
+Subproject commit 06e6703aaafe2a7246930b3c1f627a251181bc8f
diff --git a/ci/scripts/run_pytest.sh b/ci/scripts/run_pytest.sh
index 37aa2c1bb8..baff2677ab 100755
--- a/ci/scripts/run_pytest.sh
+++ b/ci/scripts/run_pytest.sh
@@ -25,5 +25,12 @@ if ! set_bionemo_home; then
     exit 1
 fi
 
-echo "Running pytest tests"
-pytest -v --nbval-lax docs/ scripts/ sub-packages/
+python -m coverage erase
+
+for dir in docs/ ./sub-packages/bionemo-*/; do
+    echo "Running pytest in $dir"
+    python -m coverage run --parallel-mode --source sub-packages/ -m pytest -v --nbval-lax $dir
+done
+
+python -m coverage combine
+python -m coverage report
diff --git a/docs/docs/user-guide/examples/bionemo-esm2/pretrain.md b/docs/docs/user-guide/examples/bionemo-esm2/pretrain.md
index 0ab20c1c13..f273673f95 100644
--- a/docs/docs/user-guide/examples/bionemo-esm2/pretrain.md
+++ b/docs/docs/user-guide/examples/bionemo-esm2/pretrain.md
@@ -61,7 +61,7 @@ strategy = nl.MegatronStrategy(
 BioNeMo2 trainer is very similar to PyTorch Lightning trainer. We can configure the training configurations and logging.
 
 ```python
-from pytorch_lightning.callbacks import LearningRateMonitor, RichModelSummary
+from lightning.pytorch.callbacks import LearningRateMonitor, RichModelSummary
 from bionemo.llm.lightning import PerplexityLoggingCallback
 
 num_steps = 20
diff --git a/scripts/gpt-pretrain.py b/scripts/gpt-pretrain.py
index 9c7f13bdfb..e856000a1e 100644
--- a/scripts/gpt-pretrain.py
+++ b/scripts/gpt-pretrain.py
@@ -17,19 +17,19 @@
 from pathlib import Path
 from typing import List, Optional, Sequence, TypedDict
 
+import lightning.pytorch as pl
 import numpy as np
-import pytorch_lightning as pl
 import torch
+
+# In lightning.pytorch 2.0 these are commented as being "any iterable or collection of iterables"
+#  for now we'll use them incase the lightning type becomes something more specific in a future release.
+from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from nemo import lightning as nl
 from nemo.collections import llm
 from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
 from nemo.lightning.megatron_parallel import DataT
 from nemo.lightning.pytorch.plugins import MegatronDataSampler
-
-# In pytorch_lightning 2.0 these are commented as being "any iterable or collection of iterables"
-#  for now we'll use them incase the lightning type becomes something more specific in a future release.
-from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from torch.utils import data
 from torch.utils.data import DataLoader, Dataset
 
diff --git a/sub-packages/bionemo-core/tests/bionemo/core/data/test_load.py b/sub-packages/bionemo-core/tests/bionemo/core/data/test_load.py
index 787de3e288..ee8146b5dd 100644
--- a/sub-packages/bionemo-core/tests/bionemo/core/data/test_load.py
+++ b/sub-packages/bionemo-core/tests/bionemo/core/data/test_load.py
@@ -110,7 +110,7 @@ def test_load_with_file(mocked_s3_download, tmp_path):
     )
 
     mocked_s3_download.side_effect = lambda _1, output_file, _2: Path(output_file).write_text("test")
-    file_path = load("foo/bar", resources=get_all_resources(tmp_path), cache_dir=tmp_path)
+    file_path = load("foo/bar", resources=get_all_resources(tmp_path), cache_dir=tmp_path, source="pbss")
     assert file_path.is_file()
     assert file_path.read_text() == "test"
 
@@ -132,7 +132,7 @@ def write_compressed_text(_1, output_file: str, _2):
 
     mocked_s3_download.side_effect = write_compressed_text
 
-    file_path = load("foo/baz", resources=get_all_resources(tmp_path), cache_dir=tmp_path)
+    file_path = load("foo/baz", resources=get_all_resources(tmp_path), cache_dir=tmp_path, source="pbss")
     assert file_path.is_file()
     assert file_path.read_text() == "test"
 
@@ -155,7 +155,7 @@ def write_compressed_text(_1, output_file: str, _2):
 
     mocked_s3_download.side_effect = write_compressed_text
 
-    file_path = load("foo/baz", resources=get_all_resources(tmp_path), cache_dir=tmp_path)
+    file_path = load("foo/baz", resources=get_all_resources(tmp_path), cache_dir=tmp_path, source="pbss")
 
     # Assert the file remained compressed.
     assert file_path.is_file()
@@ -190,7 +190,7 @@ def write_compressed_dir(_1, output_file: str, _2):
 
     mocked_s3_download.side_effect = write_compressed_dir
 
-    file_path = load("foo/dir", resources=get_all_resources(tmp_path), cache_dir=tmp_path)
+    file_path = load("foo/dir", resources=get_all_resources(tmp_path), cache_dir=tmp_path, source="pbss")
     assert file_path.is_dir()
     assert (file_path / "test_file").read_text() == "test"
 
@@ -223,7 +223,7 @@ def write_tarfile_dir(_1, output_file: str, _2):
 
     mocked_s3_download.side_effect = write_tarfile_dir
 
-    file_path = load("foo/dir", resources=get_all_resources(tmp_path), cache_dir=tmp_path)
+    file_path = load("foo/dir", resources=get_all_resources(tmp_path), cache_dir=tmp_path, source="pbss")
 
     # Assert the file stays as a tarfile.
     assert file_path.is_file()
@@ -259,7 +259,7 @@ def write_compressed_dir(_1, output_file: str, _2):
 
     mocked_s3_download.side_effect = write_compressed_dir
 
-    file_path = load("foo/dir.gz", resources=get_all_resources(tmp_path), cache_dir=tmp_path)
+    file_path = load("foo/dir.gz", resources=get_all_resources(tmp_path), cache_dir=tmp_path, source="pbss")
     assert file_path.is_dir()
     assert (file_path / "test_file").read_text() == "test"
 
diff --git a/sub-packages/bionemo-core/tests/bionemo/core/data/test_load_notebook.ipynb b/sub-packages/bionemo-core/tests/bionemo/core/data/test_load_notebook.ipynb
index 7d61836dfc..9d4dcf5cfe 100644
--- a/sub-packages/bionemo-core/tests/bionemo/core/data/test_load_notebook.ipynb
+++ b/sub-packages/bionemo-core/tests/bionemo/core/data/test_load_notebook.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -21,8 +21,14 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Downloading data from 'nvidia/clara/esm2_pretrain_nemo2_testdata:1.0' to file '/tmp/tmp_v_0_64q/dc23f4aaad387ecc12e53d56b8176430-esm2_pretrain_nemo2_testdata:1.0'.\n",
-      "Untarring contents of '/tmp/tmp_v_0_64q/dc23f4aaad387ecc12e53d56b8176430-esm2_pretrain_nemo2_testdata:1.0' to '/tmp/tmp_v_0_64q/dc23f4aaad387ecc12e53d56b8176430-esm2_pretrain_nemo2_testdata:1.0.untar'\n"
+      "Downloading data from 'nvidia/clara/scdl_sample_test:1.0' to file '/tmp/tmpqif5bfww/7a4237537bf535dfa00301ce8cc7073e0a23d5bc8aa902ad65db9f51b57a6df9-scdl_sample_test.tar.gz'.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Untarring contents of '/tmp/tmpqif5bfww/7a4237537bf535dfa00301ce8cc7073e0a23d5bc8aa902ad65db9f51b57a6df9-scdl_sample_test.tar.gz' to '/tmp/tmpqif5bfww/7a4237537bf535dfa00301ce8cc7073e0a23d5bc8aa902ad65db9f51b57a6df9-scdl_sample_test.tar.gz.untar'\n"
      ]
     },
     {
@@ -30,22 +36,20 @@
      "output_type": "stream",
      "text": [
       "{\n",
-      "    \"download_end\": \"2024-11-07 19:13:48\",\n",
-      "    \"download_start\": \"2024-11-07 19:13:46\",\n",
-      "    \"download_time\": \"2s\",\n",
+      "    \"download_end\": \"2024-12-03 18:39:20\",\n",
+      "    \"download_start\": \"2024-12-03 18:39:03\",\n",
+      "    \"download_time\": \"17s\",\n",
       "    \"files_downloaded\": 1,\n",
-      "    \"local_path\": \"/tmp/tmp_v_0_64q/tmpgadmjb1k/esm2_pretrain_nemo2_testdata_v1.0\",\n",
-      "    \"size_downloaded\": \"69.91 MB\",\n",
+      "    \"local_path\": \"/tmp/tmpqif5bfww/tmprn0ysh0w/scdl_sample_test_v1.0\",\n",
+      "    \"size_downloaded\": \"964.91 KB\",\n",
       "    \"status\": \"COMPLETED\"\n",
       "}\n"
      ]
     }
    ],
    "source": [
-    "# xfail -- we need to file a bug or figure out how to call the ngcsdk from a jupyter notebook\n",
-    "# NBVAL_RAISES_EXCEPTION\n",
     "with tempfile.TemporaryDirectory() as cache_dir:\n",
-    "    load(\"esm2/testdata_esm2_pretrain:2.0\", source=\"ngc\", cache_dir=Path(cache_dir))"
+    "    load(\"scdl/sample\", source=\"ngc\", cache_dir=Path(cache_dir))"
    ]
   },
   {
@@ -57,15 +61,15 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Downloading data from 's3://general-purpose/esm2/pretrain/2024_03_sanity.tar.gz' to file '/tmp/tmpjnvk7m8k/f796e1ca28311606ff7dd62a067508bf-2024_03_sanity.tar.gz'.\n",
-      "s3://general-purpose/esm2/pretrain/2024_03_sanity.tar.gz: 100%|██████████| 73.3M/73.3M [00:01<00:00, 38.7MB/s]\n",
-      "Untarring contents of '/tmp/tmpjnvk7m8k/f796e1ca28311606ff7dd62a067508bf-2024_03_sanity.tar.gz' to '/tmp/tmpjnvk7m8k/f796e1ca28311606ff7dd62a067508bf-2024_03_sanity.tar.gz.untar'\n"
+      "Downloading data from 's3://bionemo-ci/test-data/scdl_sample_test.tar.gz' to file '/tmp/tmpl6cgwhyn/7a4237537bf535dfa00301ce8cc7073e0a23d5bc8aa902ad65db9f51b57a6df9-scdl_sample_test.tar.gz'.\n",
+      "s3://bionemo-ci/test-data/scdl_sample_test.tar.gz: 100%|██████████| 988k/988k [00:00<00:00, 2.70MB/s]\n",
+      "Untarring contents of '/tmp/tmpl6cgwhyn/7a4237537bf535dfa00301ce8cc7073e0a23d5bc8aa902ad65db9f51b57a6df9-scdl_sample_test.tar.gz' to '/tmp/tmpl6cgwhyn/7a4237537bf535dfa00301ce8cc7073e0a23d5bc8aa902ad65db9f51b57a6df9-scdl_sample_test.tar.gz.untar'\n"
      ]
     }
    ],
    "source": [
     "with tempfile.TemporaryDirectory() as cache_dir:\n",
-    "    load(\"esm2/testdata_esm2_pretrain:2.0\", source=\"pbss\", cache_dir=Path(cache_dir))"
+    "    load(\"scdl/sample\", source=\"pbss\", cache_dir=Path(cache_dir))"
    ]
   }
  ],
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py
index e065702d3d..a6fcd73fe4 100644
--- a/sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py
+++ b/sub-packages/bionemo-esm2/src/bionemo/esm2/data/datamodule.py
@@ -18,10 +18,10 @@
 import os
 from typing import Literal
 
+from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from nemo.lightning.data import WrappedDataLoader
 from nemo.lightning.pytorch.plugins import MegatronDataSampler
 from nemo.utils import logging
-from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 
 from bionemo.esm2.data import dataset, tokenizer
 from bionemo.llm.data import collate
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/datamodule.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/datamodule.py
index b7f0e1f823..0b1598da55 100644
--- a/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/datamodule.py
+++ b/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/datamodule.py
@@ -22,10 +22,10 @@
 import pandas as pd
 import torch
 import torch.utils.data
+from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from nemo.lightning.data import WrappedDataLoader
 from nemo.lightning.pytorch.plugins import MegatronDataSampler
 from nemo.utils import logging
-from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from torch import Tensor
 from torch.utils.data import Dataset
 
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/train.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/train.py
index a6d299bf68..638729e3f4 100644
--- a/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/train.py
+++ b/sub-packages/bionemo-esm2/src/bionemo/esm2/model/finetune/train.py
@@ -18,7 +18,9 @@
 from pathlib import Path
 from typing import Sequence, Tuple
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
+from lightning.pytorch.callbacks import Callback, RichModelSummary
+from lightning.pytorch.loggers import TensorBoardLogger
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
 from nemo import lightning as nl
 from nemo.collections import llm as nllm
@@ -28,8 +30,6 @@
 from nemo.lightning.pytorch.callbacks.model_transform import ModelTransform
 from nemo.lightning.pytorch.callbacks.peft import PEFT
 from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule
-from pytorch_lightning.callbacks import Callback, RichModelSummary
-from pytorch_lightning.loggers import TensorBoardLogger
 
 from bionemo.core.data.load import load
 from bionemo.esm2.api import ESM2GenericConfig
diff --git a/sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py b/sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py
index 7fbaf83cd3..ed5e46899b 100644
--- a/sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py
+++ b/sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py
@@ -17,13 +17,13 @@
 from pathlib import Path
 from typing import List, Optional, Sequence, get_args
 
+from lightning.pytorch.callbacks import LearningRateMonitor, RichModelSummary
 from megatron.core.optimizer import OptimizerConfig
 from nemo import lightning as nl
 from nemo.collections import llm
 from nemo.lightning import resume
 from nemo.lightning.pytorch import callbacks as nl_callbacks
 from nemo.lightning.pytorch.optim import MegatronOptimizerModule
-from pytorch_lightning.callbacks import LearningRateMonitor, RichModelSummary
 
 from bionemo.core.utils.dtypes import PrecisionTypes, get_autocast_dtype
 from bionemo.esm2.api import ESM2Config
@@ -170,7 +170,7 @@ def main(
     )
 
     # for wandb integration
-    # Please refer to https://pytorch-lightning.readthedocs.io/en/0.7.6/api/pytorch_lightning.loggers.html"
+    # Please refer to https://pytorch-lightning.readthedocs.io/en/0.7.6/api/lightning.pytorch.loggers.html"
     wandb_config: Optional[WandbConfig] = (
         None
         if wandb_project is None
diff --git a/sub-packages/bionemo-esm2/tests/bionemo/esm2/model/test_stop_and_go.py b/sub-packages/bionemo-esm2/tests/bionemo/esm2/model/test_stop_and_go.py
index 710b456a2a..76e5d8b3c5 100644
--- a/sub-packages/bionemo-esm2/tests/bionemo/esm2/model/test_stop_and_go.py
+++ b/sub-packages/bionemo-esm2/tests/bionemo/esm2/model/test_stop_and_go.py
@@ -17,7 +17,7 @@
 from pathlib import Path
 from typing import Literal
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 from megatron.core.optimizer import OptimizerConfig
 from nemo import lightning as nl
 from nemo.lightning.pytorch.optim import MegatronOptimizerModule
diff --git a/sub-packages/bionemo-example_model/README.md b/sub-packages/bionemo-example_model/README.md
index c24e18e0bb..4abb4f219b 100644
--- a/sub-packages/bionemo-example_model/README.md
+++ b/sub-packages/bionemo-example_model/README.md
@@ -75,7 +75,7 @@ Similarly, `ExampleFineTuneConfig` extends `ExampleGenericConfig` for finetuning
 
 # Training Module
 
-It is helfpul to have a training module that inherits from `pytorch_lightning.LightningModule` which organizes the model architecture, training, validation, and testing logic while abstracting away boilerplate code, enabling easier and more scalable training. This wrapper can be used for all model and loss combinations specified in the config.
+It is helfpul to have a training module that inherits from `lightning.pytorch.LightningModule` which organizes the model architecture, training, validation, and testing logic while abstracting away boilerplate code, enabling easier and more scalable training. This wrapper can be used for all model and loss combinations specified in the config.
 In `bionemo.example_model.lightning.lightning_basic`, we define `BionemoLightningModule`.
 
 In this example, `training_step`, `validation_step`, and `predict_step` define the training, validation, and prediction loops are independent of the forward method. In nemo:
@@ -99,7 +99,7 @@ We specify a training strategy of type `nemo.lightning.MegatronStrategy`. This s
 
 We specify a trainer of type `nemo.lightning.Trainer`, which is an extension of the pytorch lightning trainer. This is where the devices, validation intervals, maximal steps, maximal number of epochs, and how frequently to log are specified.
 
-we specify a nemo-logger. We can set TensorBoard and WandB logging, along with extra loggers. Here, we specify a `CSVLogger` from pytorch_lightning.loggers.
+we specify a nemo-logger. We can set TensorBoard and WandB logging, along with extra loggers. Here, we specify a `CSVLogger` from lightning.pytorch.loggers.
 
 We can now proceed to training. The first pre-training scripts is `bionemo/example_model/training_scripts/pretrain_mnist.py`
 
diff --git a/sub-packages/bionemo-example_model/src/bionemo/example_model/lightning/lightning_basic.py b/sub-packages/bionemo-example_model/src/bionemo/example_model/lightning/lightning_basic.py
index b7795feaa2..396c3ef38e 100644
--- a/sub-packages/bionemo-example_model/src/bionemo/example_model/lightning/lightning_basic.py
+++ b/sub-packages/bionemo-example_model/src/bionemo/example_model/lightning/lightning_basic.py
@@ -19,7 +19,6 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, Generic, List, Optional, Sequence, Tuple, Type, TypedDict, TypeVar
 
-import pytorch_lightning as pl
 import torch
 from megatron.core import ModelParallelConfig
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
@@ -35,6 +34,7 @@
 from torchvision import transforms
 from torchvision.datasets import MNIST
 
+import lightning.pytorch as pl
 from bionemo.core import BIONEMO_CACHE_DIR
 from bionemo.core.data.multi_epoch_dataset import IdentityMultiEpochDatasetWrapper, MultiEpochDatasetResampler
 from bionemo.llm.api import MegatronLossType
diff --git a/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/finetune_mnist.py b/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/finetune_mnist.py
index f03b6856f5..b0e081fd34 100644
--- a/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/finetune_mnist.py
+++ b/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/finetune_mnist.py
@@ -17,11 +17,11 @@
 import argparse
 from pathlib import Path
 
+from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger
 from nemo import lightning as nl
 from nemo.collections import llm
 from nemo.lightning import NeMoLogger, resume
 from nemo.lightning.pytorch import callbacks as nl_callbacks
-from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger
 
 from bionemo.example_model.lightning.lightning_basic import (
     BionemoLightningModule,
diff --git a/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/pretrain_mnist.py b/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/pretrain_mnist.py
index 954af6ceb5..a5adac3cc2 100644
--- a/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/pretrain_mnist.py
+++ b/sub-packages/bionemo-example_model/src/bionemo/example_model/training_scripts/pretrain_mnist.py
@@ -16,10 +16,10 @@
 
 from pathlib import Path
 
+from lightning.pytorch.loggers import CSVLogger, TensorBoardLogger
 from nemo import lightning as nl
 from nemo.collections import llm
 from nemo.lightning import NeMoLogger, resume
-from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger
 
 from bionemo.example_model.lightning.lightning_basic import (
     BionemoLightningModule,
diff --git a/sub-packages/bionemo-example_model/tests/bionemo/example_model/lightning/test_lightning_basic.py b/sub-packages/bionemo-example_model/tests/bionemo/example_model/lightning/test_lightning_basic.py
index 651c87887c..d505cf6247 100644
--- a/sub-packages/bionemo-example_model/tests/bionemo/example_model/lightning/test_lightning_basic.py
+++ b/sub-packages/bionemo-example_model/tests/bionemo/example_model/lightning/test_lightning_basic.py
@@ -20,11 +20,11 @@
 import pytest
 import torch
 from _pytest.compat import LEGACY_PATH
+from lightning.pytorch.loggers import TensorBoardLogger
 from nemo import lightning as nl
 from nemo.collections import llm
 from nemo.lightning import NeMoLogger, io, resume
 from nemo.lightning.pytorch import callbacks as nl_callbacks
-from pytorch_lightning.loggers import TensorBoardLogger
 
 from bionemo.core import BIONEMO_CACHE_DIR
 from bionemo.core.utils.dtypes import PrecisionTypes, get_autocast_dtype
diff --git a/sub-packages/bionemo-geneformer/src/bionemo/geneformer/scripts/train_geneformer.py b/sub-packages/bionemo-geneformer/src/bionemo/geneformer/scripts/train_geneformer.py
index a41dcc7e4f..6507665868 100644
--- a/sub-packages/bionemo-geneformer/src/bionemo/geneformer/scripts/train_geneformer.py
+++ b/sub-packages/bionemo-geneformer/src/bionemo/geneformer/scripts/train_geneformer.py
@@ -26,6 +26,7 @@
 from typing import Dict, List, Optional, Sequence, Type, get_args
 
 import torch
+from lightning.pytorch.callbacks import LearningRateMonitor, RichModelSummary
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.optimizer import OptimizerConfig
 from nemo import lightning as nl
@@ -36,7 +37,6 @@
 from nemo.lightning.pytorch.optim.lr_scheduler import CosineAnnealingScheduler
 from nemo.utils import logging
 from nemo.utils.exp_manager import TimingCallback
-from pytorch_lightning.callbacks import LearningRateMonitor, RichModelSummary
 
 from bionemo.core.utils.dtypes import PrecisionTypes, get_autocast_dtype
 from bionemo.geneformer.api import FineTuneSeqLenBioBertConfig, GeneformerConfig
@@ -195,7 +195,7 @@ def main(
     )
 
     # for wandb integration
-    # Please refer to https://pytorch-lightning.readthedocs.io/en/0.7.6/api/pytorch_lightning.loggers.html"
+    # Please refer to https://pytorch-lightning.readthedocs.io/en/0.7.6/api/lightning.pytorch.loggers.html"
     wandb_options: Optional[WandbConfig] = (
         None
         if wandb_project is None
diff --git a/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py b/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py
index c9338ee731..3252df2ced 100644
--- a/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py
+++ b/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_model.py
@@ -23,6 +23,7 @@
 import pytest
 import torch
 import torch.utils.data
+from lightning.pytorch.loggers import TensorBoardLogger
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
 from megatron.core.transformer.module import Float16Module
 from nemo import lightning as nl
@@ -34,7 +35,6 @@
 from nemo.lightning.pytorch.callbacks.peft import PEFT
 from nemo.lightning.pytorch.optim.lr_scheduler import WarmupPolicyScheduler
 from nemo.lightning.pytorch.optim.megatron import MegatronOptimizerModule
-from pytorch_lightning.loggers import TensorBoardLogger
 from torch.nn import functional as F
 from tqdm import tqdm
 
diff --git a/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py b/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py
index 665af8def9..07f91f6160 100644
--- a/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py
+++ b/sub-packages/bionemo-geneformer/tests/bionemo/geneformer/test_stop_and_go.py
@@ -28,7 +28,7 @@
 import pathlib
 from typing import Literal
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
 from nemo import lightning as nl
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py b/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py
index 2bacb71fac..b86674ce7b 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/data/datamodule.py
@@ -16,7 +16,7 @@
 
 from typing import Any, Dict
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 from nemo.utils import logging
 
 
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py b/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py
index e766e2e559..5219d830f9 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py
@@ -15,7 +15,7 @@
 
 from typing import Any, Callable, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Tuple, TypeVar, Union
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch.distributed
 from megatron.core import parallel_state
 from megatron.core.optimizer.optimizer_config import OptimizerConfig
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/lightning.py b/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/lightning.py
index 8d93adee2a..810ee37a77 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/lightning.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/lightning.py
@@ -16,7 +16,7 @@
 
 from typing import Callable, Dict, Iterable, Optional, Protocol, Sequence, TypedDict, cast
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch.distributed
 from apex.optimizers import FusedAdam
 from megatron.core import parallel_state
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/testing_utils.py b/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/testing_utils.py
index 9e4b46293e..40fd162a38 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/testing_utils.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/model/biobert/testing_utils.py
@@ -14,7 +14,7 @@
 # limitations under the License.
 
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch.nn.functional as F
 
 
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/run/config_models.py b/sub-packages/bionemo-llm/src/bionemo/llm/run/config_models.py
index fce5db3c75..f22e7c3650 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/run/config_models.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/run/config_models.py
@@ -19,7 +19,7 @@
 from dataclasses import field
 from typing import Any, Callable, Dict, Generic, List, Literal, Optional, Type, TypeVar
 
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch
 from pydantic import BaseModel, field_serializer, field_validator, model_validator
 from torch.nn import functional as F
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/train.py b/sub-packages/bionemo-llm/src/bionemo/llm/train.py
index cae6f9c848..442b01957a 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/train.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/train.py
@@ -19,6 +19,7 @@
 from dataclasses import field
 from typing import Optional
 
+from lightning.pytorch.callbacks import LearningRateMonitor, RichModelSummary
 from megatron.core.optimizer import OptimizerConfig
 from nemo import lightning as nl
 from nemo.collections import llm
@@ -28,7 +29,6 @@
 from nemo.lightning.pytorch.optim.lr_scheduler import CosineAnnealingScheduler
 from nemo.utils import logging
 from pydantic import BaseModel
-from pytorch_lightning.callbacks import LearningRateMonitor, RichModelSummary
 
 from bionemo.llm.lightning import BionemoLightningModule, PerplexityLoggingCallback
 from bionemo.llm.model.biobert.lightning import biobert_lightning_module
diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/utils/logger_utils.py b/sub-packages/bionemo-llm/src/bionemo/llm/utils/logger_utils.py
index ebba878c66..912d67bf7b 100644
--- a/sub-packages/bionemo-llm/src/bionemo/llm/utils/logger_utils.py
+++ b/sub-packages/bionemo-llm/src/bionemo/llm/utils/logger_utils.py
@@ -15,11 +15,11 @@
 import pathlib
 from typing import Any, Dict, List, Optional, Sequence
 
+from lightning.pytorch.loggers import TensorBoardLogger, WandbLogger
 from nemo.lightning.nemo_logger import NeMoLogger
 from nemo.lightning.pytorch import callbacks as nemo_callbacks
 from nemo.utils import logging
 from pydantic import BaseModel
-from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger
 
 
 __all__: Sequence[str] = (
diff --git a/sub-packages/bionemo-testing/src/bionemo/testing/callbacks.py b/sub-packages/bionemo-testing/src/bionemo/testing/callbacks.py
index 0c84297039..ed49eeb498 100644
--- a/sub-packages/bionemo-testing/src/bionemo/testing/callbacks.py
+++ b/sub-packages/bionemo-testing/src/bionemo/testing/callbacks.py
@@ -18,7 +18,7 @@
 from typing import Dict, List
 
 import torch
-from pytorch_lightning import Callback
+from lightning.pytorch import Callback
 
 
 class MetricTracker(Callback):  # noqa: D101
diff --git a/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py b/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py
index 85ed5703e0..c3d5b383ef 100644
--- a/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py
+++ b/sub-packages/bionemo-testing/src/bionemo/testing/harnesses/stop_and_go.py
@@ -19,9 +19,9 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Literal, Sequence, Type, TypeVar
 
+import lightning.pytorch as pl
 import nemo.lightning as nl
 import pytest
-import pytorch_lightning as pl
 from nemo.collections import llm
 from nemo.lightning import resume
 from nemo.lightning.nemo_logger import NeMoLogger
@@ -314,6 +314,10 @@ def run_stop_and_go(cls):
         cls.stop()
         cls.resume()
 
+        # Cleanup and reinitialize the temporary directory so we don't conflict with a previous checkpoint.
+        cls.tempdir.cleanup()
+        cls.tempdir = tempfile.TemporaryDirectory()
+
         # Continuous model training.
         cls.continuous()
 
diff --git a/sub-packages/bionemo-testing/src/bionemo/testing/megatron_parallel_state_utils.py b/sub-packages/bionemo-testing/src/bionemo/testing/megatron_parallel_state_utils.py
index e4d2a4f316..1686de309d 100644
--- a/sub-packages/bionemo-testing/src/bionemo/testing/megatron_parallel_state_utils.py
+++ b/sub-packages/bionemo-testing/src/bionemo/testing/megatron_parallel_state_utils.py
@@ -36,8 +36,8 @@ def my_test():
 from unittest import mock
 from unittest.mock import MagicMock
 
+import lightning.pytorch as pl
 import megatron.core.num_microbatches_calculator
-import pytorch_lightning as pl
 import torch
 import torch.distributed
 from megatron.core import parallel_state
diff --git a/sub-packages/bionemo-testing/src/bionemo/testing/testing_callbacks.py b/sub-packages/bionemo-testing/src/bionemo/testing/testing_callbacks.py
index c591d723ff..f8088abd78 100644
--- a/sub-packages/bionemo-testing/src/bionemo/testing/testing_callbacks.py
+++ b/sub-packages/bionemo-testing/src/bionemo/testing/testing_callbacks.py
@@ -19,11 +19,11 @@
 
 import numpy as np
 import torch
+from lightning.pytorch import Callback, LightningModule, Trainer
 from nemo.lightning import io
 from nemo.lightning.data import MegatronPretrainingSampler
 from nemo.lightning.megatron_parallel import CallbackMethods, DataT, MegatronLossReduction, MegatronStep
 from overrides import override
-from pytorch_lightning import Callback, LightningModule, Trainer
 
 from bionemo.testing.harnesses.mode import Mode
 from bionemo.testing.torch import recursive_detach