Repository: huggingface/accelerate
Pull request: Add warning for problematic libraries
# What does this PR do?

Since this is a **very** common issue, adds a check in `notebook_launcher` for specific libraries that are *known* to cause issues (such as starting CUDA forks *on import*). 

Fixes https://github.com/huggingface/accelerate/issues/2148


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md#submitting-a-pull-request-pr),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/accelerate/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/accelerate/tree/main/docs#writing-documentation---specification).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@BenjaminBossan @SunMarc 

Problem statement:
notebook_launcher fails when loading model from_pretrained 
### System Info

```Shell
- `Accelerate` version: 0.25.0.dev0
- Platform: Linux-6.2.0-36-generic-x86_64-with-glibc2.35
- Python version: 3.9.18
- Numpy version: 1.25.0
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 2015.52 GB
- GPU type: NVIDIA A40
- `Accelerate` default config:
        - compute_environment: LOCAL_MACHINE
        - distributed_type: MULTI_GPU
        - mixed_precision: no
        - use_cpu: False
        - debug: False
        - num_processes: 2
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: 1,2
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

Run this snippet in a Jupyter Notebook on a multi-gpu machine:

```python
from accelerate import notebook_launcher, Accelerator
from transformers import AutoModelForCausalLM

def a():
    state = Accelerator()
    print(state.process_index, state.device)

model = AutoModelForCausalLM.from_pretrained("gpt2")

notebook_launcher(a, args=(), num_processes=2)
```

**Output**
```python
RuntimeError: CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. This likely stems from an outside import causing issues once the `notebook_launcher()` is called. Please review your imports and test them when running the `notebook_launcher()` to identify which one is problematic and causing CUDA to be initialized.
```

### Expected behavior

The code above without ```model = AutoModelForCausalLM.from_pretrained("gpt2")``` gives:
```terminal
Launching training on 2 GPUs.
0 cuda:0
1 cuda:1

```

Any suggestions?

Hints from issue discussion:
Pretty sure AutoModel will move the model to CUDA. I would check that, and keep that under your a() function if so. 
If that were the case, then the official example for the notebook_launcher would fail too ([link](https://github.com/huggingface/notebooks/blob/main/examples/accelerate_examples/simple_nlp_example.ipynb))

Excerpt:
ln [16]:
```python
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=2)
```

and ln [25]:
```python
from accelerate import notebook_launcher

notebook_launcher(training_function, (model,))
```


@nikrruun can you provide your output for `pip freeze` please?
I narrowed the problem down to bitsandbytes package. Importing it immediately causes cuda init, which is not allowed for notebook_launcher.
Removing the package resolves the issue, but is probably undesirable.
To summarize: AutoModel.from_pretrained automatically imports bitsandbytes (idk for what purpose) if it is available.



pip freeze of the non-working env as a file: 
[pip.txt](https://github.com/huggingface/accelerate/files/13348970/pip.txt)

Yes, that is something unavoidable and what I was going to ask about.

Commit message: Rm comment

Patch to review:
diff --git a/src/accelerate/launchers.py b/src/accelerate/launchers.py
index 310e52c313d..b381469d1f3 100644
--- a/src/accelerate/launchers.py
+++ b/src/accelerate/launchers.py
@@ -19,7 +19,7 @@
 import torch
 
 from .state import AcceleratorState, PartialState
-from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment
+from .utils import PrecisionType, PrepareForLaunch, are_libraries_initialized, is_mps_available, patch_environment
 
 
 def test_launch():
@@ -142,6 +142,18 @@ def train(*args):
                     "inside your training function. Restart your notebook and make sure no cells initializes an "
                     "`Accelerator`."
                 )
+            # Check for specific libraries known to initialize CUDA that users constantly use
+            already_initialized = are_libraries_initialized("bitsandbytes")
+            if any(already_initialized.values()):
+                err = (
+                    "Could not start distributed process. Libraries known to initialize CUDA upon import have been "
+                    "imported already. Please keep these imports inside your `function` to try and help with this:"
+                )
+                for lib_name, initialized in already_initialized.items():
+                    if initialized:
+                        err += f"\n\t* `{lib_name}`"
+                raise AssertionError(err)
+
             # torch.distributed will expect a few environment variable to be here. We set the ones common to each
             # process here (the other ones will be set be the launcher).
             with patch_environment(
diff --git a/src/accelerate/test_utils/scripts/test_notebook.py b/src/accelerate/test_utils/scripts/test_notebook.py
index 8f215d8fd19..999fab34cbe 100644
--- a/src/accelerate/test_utils/scripts/test_notebook.py
+++ b/src/accelerate/test_utils/scripts/test_notebook.py
@@ -1,17 +1,34 @@
 # Test file to ensure that in general certain situational setups for notebooks work.
-import argparse
+import os
+
+from pytest import raises
 
 from accelerate import PartialState, notebook_launcher
+from accelerate.test_utils import require_bnb
+from accelerate.utils import is_bnb_available
+
 
+def basic_function():
+    # Just prints the PartialState
+    print(f"PartialState:\n{PartialState()}")
 
-parser = argparse.ArgumentParser()
-parser.add_argument("--num_processes", type=int, default=1)
-args = parser.parse_args()
 
+NUM_PROCESSES = os.environ.get("ACCELERATE_NUM_PROCESSES", 1)
 
-def function():
-    print(f"PartialState:\n{PartialState()}")
+
+def test_can_initialize():
+    notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES)
+
+
+@require_bnb
+def test_problematic_imports():
+    with raises(AssertionError, match="Please keep these imports"):
+        notebook_launcher(basic_function, (), num_processes=NUM_PROCESSES)
 
 
-if __name__ == "__main__":
-    notebook_launcher(function, num_processes=int(args.num_processes))
+def main():
+    print("Test basic notebook can be ran")
+    test_can_initialize()
+    if is_bnb_available():
+        print("Test problematic imports (bnb)")
+        test_problematic_imports()
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index d6d1e2f2f0a..0b2b6c6d58d 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -423,13 +423,15 @@ class SubprocessCallException(Exception):
     pass
 
 
-def run_command(command: List[str], return_stdout=False):
+def run_command(command: List[str], return_stdout=False, env=None):
     """
     Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture
     if an error occured while running `command`
     """
+    if env is None:
+        env = os.environ.copy()
     try:
-        output = subprocess.check_output(command, stderr=subprocess.STDOUT)
+        output = subprocess.check_output(command, stderr=subprocess.STDOUT, env=env)
         if return_stdout:
             if hasattr(output, "decode"):
                 output = output.decode("utf-8")
diff --git a/src/accelerate/utils/__init__.py b/src/accelerate/utils/__init__.py
index 88cc927f001..f9cab741073 100644
--- a/src/accelerate/utils/__init__.py
+++ b/src/accelerate/utils/__init__.py
@@ -37,7 +37,13 @@
     TensorInformation,
     TorchDynamoPlugin,
 )
-from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env, str_to_bool
+from .environment import (
+    are_libraries_initialized,
+    get_int_from_env,
+    parse_choice_from_env,
+    parse_flag_from_env,
+    str_to_bool,
+)
 from .imports import (
     get_ccl_version,
     is_4bit_bnb_available,
diff --git a/src/accelerate/utils/environment.py b/src/accelerate/utils/environment.py
index 0cd46c2dcaf..5d3c44ffff4 100644
--- a/src/accelerate/utils/environment.py
+++ b/src/accelerate/utils/environment.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import os
+import sys
+from typing import Dict
 
 
 def str_to_bool(value) -> int:
@@ -48,3 +50,12 @@ def parse_flag_from_env(key, default=False):
 def parse_choice_from_env(key, default="no"):
     value = os.environ.get(key, str(default))
     return value
+
+
+def are_libraries_initialized(library_names: list = []) -> Dict[str, bool]:
+    """
+    Checks if any of `library_names` are imported in the environment. Will return results as a `key:bool` pair.
+    """
+    if isinstance(library_names, str):
+        library_names = [library_names]
+    return {lib_name: lib_name in sys.modules for lib_name in library_names}
diff --git a/tests/test_multigpu.py b/tests/test_multigpu.py
index 73ee6367f5c..a479130b74f 100644
--- a/tests/test_multigpu.py
+++ b/tests/test_multigpu.py
@@ -21,7 +21,8 @@
 import accelerate
 from accelerate import Accelerator
 from accelerate.big_modeling import dispatch_model
-from accelerate.test_utils import assert_exception, execute_subprocess_async, require_multi_gpu, skip
+from accelerate.test_utils import assert_exception, execute_subprocess_async, require_multi_gpu
+from accelerate.test_utils.testing import run_command
 from accelerate.utils import patch_environment
 
 
@@ -33,6 +34,9 @@ def setUp(self):
             mod_file.split(os.path.sep)[:-1] + ["scripts", "test_distributed_data_loop.py"]
         )
         self.operation_file_path = os.path.sep.join(mod_file.split(os.path.sep)[:-1] + ["scripts", "test_ops.py"])
+        self.notebook_launcher_path = os.path.sep.join(
+            mod_file.split(os.path.sep)[:-1] + ["scripts", "test_notebook.py"]
+        )
 
     @require_multi_gpu
     def test_multi_gpu(self):
@@ -66,23 +70,16 @@ def test_distributed_data_loop(self):
         with patch_environment(omp_num_threads=1, cuda_visible_devices="0,1"):
             execute_subprocess_async(cmd, env=os.environ.copy())
 
-    # Need to see why this test raises forking issues when ran as a suite
-    @skip
     @require_multi_gpu
     def test_notebook_launcher(self):
         """
-        This test checks that the `notebook_launcher` will be able to intialize
-        a `PartialState` without issue
+        This test checks a variety of situations and scenarios
+        with the `notebook_launcher`
         """
-        cmd = [
-            "python",
-            "-m",
-            "accelerate.test_utils.scripts.test_notebook",
-            "--num_processes",
-            str(torch.cuda.device_count()),
-        ]
+        cmd = ["torchrun", f"--nproc_per_node={torch.cuda.device_count()}", self.notebook_launcher_path]
+        print(f"Running {cmd}")
         with patch_environment(omp_num_threads=1):
-            execute_subprocess_async(cmd, env=os.environ.copy())
+            run_command(cmd, env=os.environ.copy())
 
 
 if __name__ == "__main__":


Code review comments:
diff_hunk: @@ -48,3 +50,12 @@ def parse_flag_from_env(key, default=False):
 def parse_choice_from_env(key, default="no"):
     value = os.environ.get(key, str(default))
     return value
+
+
+def are_libraries_initialized(library_names: list = []) -> Dict[str, bool]:
+    """
+    Checks if any of `library_names` are imported in the environment. Will return results as a `key:bool` pair.
+    """
+    if isinstance(library_names, str):
+        library_names = [library_names]
+    return {lib_name: lib_name in sys.modules for lib_name in library_names}
path: src/accelerate/utils/environment.py
text: @user1:
The code here and at the call site could be simplified if only the list of loaded libraries were returned, instead of returning all libs in a dict with their status as value. But maybe you plan on using this for other purposes later which requires all the results, so not sure.

diff_hunk: @@ -48,3 +50,12 @@ def parse_flag_from_env(key, default=False):
 def parse_choice_from_env(key, default="no"):
     value = os.environ.get(key, str(default))
     return value
+
+
+def are_libraries_initialized(library_names: list = []) -> Dict[str, bool]:
path: src/accelerate/utils/environment.py
text: @user1:
```suggestion
def are_libraries_initialized(library_names: list[str]) -> Dict[str, bool]:
```

Alternatively, to avoid the `isinstance` check below, just do `*library_names`, then we know it's always a tuple.

diff_hunk: @@ -142,6 +142,18 @@ def train(*args):
                     "inside your training function. Restart your notebook and make sure no cells initializes an "
                     "`Accelerator`."
                 )
+            # Check for specific libraries known to initialize CUDA that users constantly use
+            already_initialized = are_libraries_initialized("bitsandbytes")
+            if any(already_initialized.values()):
+                err = (
+                    "Could not start distributed process. Libraries known to initialize CUDA upon import have been "
+                    "imported already. Please keep these imports inside your `function` to try and help with this:"
path: src/accelerate/launchers.py
text: @user1:
Do you think it's clear what "`function`" means in this context?

@author:
Fixed to keep align with the rest of the err's in `notebook_launcher`

diff_hunk: @@ -142,6 +142,18 @@ def train(*args):
                     "inside your training function. Restart your notebook and make sure no cells initializes an "
                     "`Accelerator`."
                 )
+            # Check for specific libraries known to initialize CUDA that users constantly use
+            already_initialized = are_libraries_initialized("bitsandbytes")
+            if any(already_initialized.values()):
+                err = (
+                    "Could not start distributed process. Libraries known to initialize CUDA upon import have been "
+                    "imported already. Please keep these imports inside your `function` to try and help with this:"
+                )
+                for lib_name, initialized in already_initialized.items():
+                    if initialized:
+                        err += f"\n\t* `{lib_name}`"
+                raise AssertionError(err)
path: src/accelerate/launchers.py
text: @user1:
Not sure about `AssertionError`, I thought they are basically reserved for `assert`s. Maybe use `ImportError` or `RuntimeError`?

@author:
`Runtime` would make more sense then I think, will fix