force spawn start method

EleutherAI · Apr 9, 2023 · ccdc190 · ccdc190
1 parent 158a754
commit ccdc190
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/elk/evaluation/evaluate.py b/elk/evaluation/evaluate.py
@@ -104,7 +104,8 @@ def evaluate_reporters(cfg: EvaluateConfig, out_dir: Optional[Path] = None):
 
     cols = ["layer", "loss", "acc", "cal_acc", "auroc"]
     # Evaluate reporters for each layer in parallel
-    with mp.Pool(num_devices) as pool, open(out_dir / "eval.csv", "w") as f:
+    ctx = mp.get_context("spawn")
+    with ctx.Pool(num_devices) as pool, open(out_dir / "eval.csv", "w") as f:
         fn = partial(
             evaluate_reporter, cfg, ds, devices=devices, world_size=num_devices
         )

diff --git a/elk/extraction/extraction.py b/elk/extraction/extraction.py
@@ -289,6 +289,9 @@ def get_splits() -> SplitDict:
         )
         for (split_name, split_info) in splits.items()
     }
+    import multiprocess as mp
+
+    mp.set_start_method("spawn", force=True)  # type: ignore[attr-defined]
 
     ds = dict()
     for split, builder in builders.items():

diff --git a/elk/training/train.py b/elk/training/train.py
@@ -217,7 +217,8 @@ def train(cfg: RunConfig, out_dir: Optional[Path] = None):
         if feat.startswith("hidden_")
     ]
     # Train reporters for each layer in parallel
-    with mp.Pool(num_devices) as pool, open(out_dir / "eval.csv", "w") as f:
+    ctx = mp.get_context("spawn")
+    with ctx.Pool(num_devices) as pool, open(out_dir / "eval.csv", "w") as f:
         fn = partial(
             train_reporter, cfg, ds, out_dir, devices=devices, world_size=num_devices
         )