chore: pinpoint the platform needed (Linux) for ddp examples (#283)

Author: nastaran78

neps_examples/__init__.py

@@ -23,6 +23,13 @@
         "pytorch_native_ddp",
         "pytorch_lightning_ddp",
     ],
+    "async_evaluation": [
+        "submit",
+    ],
+    "experimental": [
+        "ask_and_tell_example",
+        "freeze_thaw",
+    ],
 }
 
 core_examples = [  # Run locally and on github actions

neps_examples/efficiency/pytorch_lightning_ddp.py

@@ -1,4 +1,5 @@
 import logging
+import platform
 
 import lightning as L
 import torch
@@ -7,7 +8,7 @@
 from torch.utils.data import DataLoader, random_split
 import neps
 
-NUM_GPU = 8  # Number of GPUs to use for DDP
+NUM_GPU = 4  # Number of GPUs to use for DDP
 
 
 class ToyModel(nn.Module):
@@ -55,6 +56,14 @@ def configure_optimizers(self):
 
 
 def evaluate_pipeline(lr=0.1, epoch=20):
+    if platform.system() != "Linux":
+        raise RuntimeError(
+            "This example uses torch.distributed via Lightning's DDP "
+            f"strategy, which is only supported on Linux here. Detected "
+            f"platform: {platform.system()}. Run it on a Linux machine with "
+            "GPUs (e.g. a Linux GPU cluster)."
+        )
+
     L.seed_everything(42)
     # Model
     model = LightningModel(lr=lr)
@@ -86,7 +95,7 @@ def evaluate_pipeline(lr=0.1, epoch=20):
 
 
 class HPOSpace(neps.PipelineSpace):
-    lr = neps.Float(lower=0.001, upper=0.1, log=True, prior=0.01)
+    lr = neps.Float(lower=0.001, upper=0.1, log=True, prior=0.01, prior_confidence="high")
     epoch = neps.IntegerFidelity(lower=1, upper=3)
 
 

neps_examples/efficiency/pytorch_lightning_fsdp.py

@@ -2,6 +2,8 @@
 
 Mind that this example does not run on Windows at the moment."""
 
+import platform
+
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -35,6 +37,13 @@ def configure_optimizers(self):
 
 
 def evaluate_pipeline(lr=0.1, epoch=20):
+    if platform.system() != "Linux":
+        raise RuntimeError(
+            "This example uses FSDP with the NCCL backend, which only runs on "
+            f"Linux. Detected platform: {platform.system()}. Run it on a Linux "
+            "machine with GPUs (e.g. a Linux GPU cluster)."
+        )
+
     L.seed_everything(42)
 
     # Data
@@ -45,8 +54,8 @@ def evaluate_pipeline(lr=0.1, epoch=20):
     model = LanguageModel(vocab_size=dataset.vocab_size, lr=lr)
 
     # Trainer
-    trainer = L.Trainer(accelerator="cuda", strategy=FSDPStrategy())
-    trainer.fit(model, train_dataloader, max_epochs=epoch)
+    trainer = L.Trainer(accelerator="cuda", strategy=FSDPStrategy(), max_epochs=epoch)
+    trainer.fit(model, train_dataloader)
     return trainer.logged_metrics["train_loss"].detach().item()
 
 
@@ -57,7 +66,7 @@ def evaluate_pipeline(lr=0.1, epoch=20):
     logging.basicConfig(level=logging.INFO)
 
     class HPOSpace(neps.PipelineSpace):
-        lr = neps.Float(lower=0.001, upper=0.1, log=True, prior=0.01)
+        lr = neps.Float(lower=0.001, upper=0.1, log=True, prior=0.01, prior_confidence="high")
         epoch = neps.IntegerFidelity(lower=1, upper=3)
 
     neps.run(

neps_examples/efficiency/pytorch_native_ddp.py

@@ -3,6 +3,7 @@
 Mind that this example does not run on Windows at the moment."""
 
 import os
+import platform
 import sys
 import tempfile
 import torch
@@ -16,7 +17,7 @@
 import neps
 import logging
 
-NUM_GPU = 8  # Number of GPUs to use for DDP
+NUM_GPU = 4  # Number of GPUs to use for DDP
 
 # On Windows platform, the torch.distributed package only
 # supports Gloo backend, FileStore and TcpStore.
@@ -88,6 +89,13 @@ def demo_basic(rank, world_size, loss_dict, learning_rate, epochs):
 
 
 def evaluate_pipeline(learning_rate, epochs):
+    if platform.system() != "Linux":
+        raise RuntimeError(
+            "This example uses torch.distributed DDP, which is only supported "
+            f"on Linux here. Detected platform: {platform.system()}. Run it on "
+            "a Linux machine with GPUs (e.g. a Linux GPU cluster)."
+        )
+
     from torch.multiprocessing import Manager
 
     world_size = NUM_GPU  # Number of GPUs

neps_examples/efficiency/pytorch_native_fsdp.py

@@ -5,6 +5,7 @@
 import math
 import os
 import functools
+import platform
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -189,6 +190,13 @@ def fsdp_main(rank, world_size, test_loss_tensor, lr, epochs, save_model=False):
 
 
 def evaluate_pipeline(lr=0.1, epoch=20):
+    if platform.system() != "Linux":
+        raise RuntimeError(
+            "This example uses FSDP with the NCCL backend, which only runs on "
+            f"Linux. Detected platform: {platform.system()}. Run it on a Linux "
+            "machine with GPUs (e.g. a Linux GPU cluster)."
+        )
+
     torch.manual_seed(42)
 
     test_loss_tensor = torch.zeros(1)
@@ -209,7 +217,7 @@ def evaluate_pipeline(lr=0.1, epoch=20):
     logging.basicConfig(level=logging.INFO)
 
     class HPOSpace(neps.PipelineSpace):
-        lr = neps.Float(lower=0.0001, upper=0.1, log=True, prior=0.01)
+        lr = neps.Float(lower=0.0001, upper=0.1, log=True, prior=0.01, prior_confidence="high")
         epoch = neps.IntegerFidelity(lower=1, upper=3)
 
     neps.run(