move pp check to train_esm2.py

NVIDIA · Dec 20, 2024 · fb5ad07 · fb5ad07
1 parent dcc84ec
commit fb5ad07
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
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
@@ -255,6 +255,9 @@ def main(
     if scheduler_num_steps is None:
         scheduler_num_steps = num_steps
 
+    if (log_train_ppl or log_val_ppl) and pipeline_model_parallel_size > 1:
+        raise NotImplementedError("Perplexity logging does not support pipeline parallelism yet.")
+
     model = biobert_lightning_module(
         esm2_config,
         tokenizer=tokenizer,

diff --git a/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py b/sub-packages/bionemo-llm/src/bionemo/llm/lightning.py
@@ -293,9 +293,6 @@ def __init__(
         # all scaling on the internal states are cancelled out in the formula "exp(total_log_probs / count)" so we can safely sum across all devices
         self.log_train_ppl = log_train_ppl
         self.log_val_ppl = log_val_ppl
-        if (log_train_ppl or log_val_ppl) and self.trainer.strategy.pipeline_model_parallel_size > 1:
-            raise NotImplementedError("Perplexity logging does not support pipeline parallelism yet.")
-
         if log_train_ppl:
             self.train_ppl = MegatronPerplexityMetric(ignore_index=-100)
         if log_val_ppl: