From 868be332da2e98c342eb8f5b95d2603a6fbdb34f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 3 Dec 2024 11:55:49 -0800 Subject: [PATCH] Bump 3rdparty/NeMo from `e2b0f0e` to `06e6703` (#486) Bumps [3rdparty/NeMo](https://github.com/NVIDIA/NeMo) from `e2b0f0e` to `06e6703`.
Commits

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
--------- Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Peter St. John --- 3rdparty/NeMo | 2 +- ci/scripts/run_pytest.sh | 11 ++++-- .../examples/bionemo-esm2/pretrain.md | 2 +- scripts/gpt-pretrain.py | 10 +++--- .../tests/bionemo/core/data/test_load.py | 12 +++---- .../core/data/test_load_notebook.ipynb | 34 +++++++++++-------- .../src/bionemo/esm2/data/datamodule.py | 2 +- .../bionemo/esm2/model/finetune/datamodule.py | 2 +- .../src/bionemo/esm2/model/finetune/train.py | 6 ++-- .../src/bionemo/esm2/scripts/train_esm2.py | 4 +-- .../bionemo/esm2/model/test_stop_and_go.py | 2 +- sub-packages/bionemo-example_model/README.md | 4 +-- .../lightning/lightning_basic.py | 2 +- .../training_scripts/finetune_mnist.py | 2 +- .../training_scripts/pretrain_mnist.py | 2 +- .../lightning/test_lightning_basic.py | 2 +- .../geneformer/scripts/train_geneformer.py | 4 +-- .../tests/bionemo/geneformer/test_model.py | 2 +- .../bionemo/geneformer/test_stop_and_go.py | 2 +- .../src/bionemo/llm/data/datamodule.py | 2 +- .../bionemo-llm/src/bionemo/llm/lightning.py | 2 +- .../bionemo/llm/model/biobert/lightning.py | 2 +- .../llm/model/biobert/testing_utils.py | 2 +- .../src/bionemo/llm/run/config_models.py | 2 +- .../bionemo-llm/src/bionemo/llm/train.py | 2 +- .../src/bionemo/llm/utils/logger_utils.py | 2 +- .../src/bionemo/testing/callbacks.py | 2 +- .../bionemo/testing/harnesses/stop_and_go.py | 6 +++- .../testing/megatron_parallel_state_utils.py | 2 +- .../src/bionemo/testing/testing_callbacks.py | 2 +- 30 files changed, 74 insertions(+), 59 deletions(-) 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