Lightning Fabric fails in AWS sagemaker jupyter notebook

[lrnilingy]

Dear All - I am a new user of fabric, and I tried to make it work in AWS Sagemaker Jupyter notebook.

The instance type is ml.g5.12xlarge which has 4 GPUs (NVIDIA A10G). It has a preinstalled torch conda environment with CUDA drivers & torch 2.2, etc.

However, after I made the changes following the documentation. I ran all the cells of the jupyter notebook and got errors from fabric.launch(train):

---------------------------------------------------------------------------
ProcessRaisedException                    Traceback (most recent call last)
Cell In[33], line 3
      1 # Notebook Cell
      2 fabric = Fabric(strategy="auto", accelerator="gpu", devices=4)
----> 3 fabric.launch(train)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/fabric.py:843, in Fabric.launch(self, function, *args, **kwargs)
    838 elif isinstance(self.strategy.launcher, (_MultiProcessingLauncher, _XLALauncher)):
    839     raise TypeError(
    840         f"To spawn processes with the `{type(self.strategy).__name__}` strategy, `.launch()` needs to be called"
    841         " with a function that contains the code to launch in processes."
    842     )
--> 843 return self._wrap_and_launch(function, self, *args, **kwargs)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/fabric.py:928, in Fabric._wrap_and_launch(self, to_run, *args, **kwargs)
    926 to_run = partial(self._wrap_with_setup, to_run)
    927 if (launcher := self._strategy.launcher) is not None:
--> 928     return launcher.launch(to_run, *args, **kwargs)
    929 return to_run(*args, **kwargs)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/strategies/launchers/multiprocessing.py:117, in _MultiProcessingLauncher.launch(self, function, *args, **kwargs)
    114 else:
    115     process_args = [function, args, kwargs, return_queue]
--> 117 mp.start_processes(
    118     self._wrapping_function,
    119     args=process_args,
    120     nprocs=self._strategy.num_processes,
    121     start_method=self._start_method,
    122 )
    123 return return_queue.get()

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:197, in start_processes(fn, args, nprocs, join, daemon, start_method)
    194     return context
    196 # Loop on join until it returns True or raises an exception.
--> 197 while not context.join():
    198     pass

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/multiprocessing/spawn.py:158, in ProcessContext.join(self, timeout)
    156 msg = "\n\n-- Process %d terminated with the following error:\n" % error_index
    157 msg += original_trace
--> 158 raise ProcessRaisedException(msg, error_index, failed_process.pid)

ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 68, in _wrap
    fn(i, *args)
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/strategies/launchers/multiprocessing.py", line 142, in _wrapping_function
    results = function(*args, **kwargs)
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 932, in _wrap_with_setup
    self._strategy.setup_environment()
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/strategies/ddp.py", line 119, in setup_environment
    super().setup_environment()
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/strategies/strategy.py", line 109, in setup_environment
    self.accelerator.setup_device(self.root_device)
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/accelerators/cuda.py", line 37, in setup_device
    _check_cuda_matmul_precision(device)
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/accelerators/cuda.py", line 161, in _check_cuda_matmul_precision
    if not torch.cuda.is_available() or not _is_ampere_or_later(device):
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/lightning/fabric/accelerators/cuda.py", line 155, in _is_ampere_or_later
    major, _ = torch.cuda.get_device_capability(device)
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/cuda/__init__.py", line 444, in get_device_capability
    prop = get_device_properties(device)
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/cuda/__init__.py", line 458, in get_device_properties
    _lazy_init()  # will define _get_device_properties
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/cuda/__init__.py", line 307, in _lazy_init
    raise e
  File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/cuda/__init__.py", line 303, in _lazy_init
    torch._C._cuda_init()
RuntimeError: ('CUDA driver initialization failed, you might not have a CUDA gpu.', "This image doesn't seem to support current EC2 instance type, please check release notes for supported EC2 instance type")

I am confused since cuda/gpu/instance should not be a problem. In fact, I was using Huggingface Accelerate and that package worked on this instance.

Do I miss something? Thank you very much.

Some environment info:

Last updated: 2024-08-07T21:20:45.887449+00:00

Python implementation: CPython
Python version       : 3.10.14
IPython version      : 8.25.0

Compiler    : GCC 12.3.0
OS          : Linux
Release     : 5.10.217-205.860.amzn2.x86_64
Machine     : x86_64
Processor   : x86_64
CPU cores   : 48
Architecture: 64bit

torch    : 2.2.0
lightning: 2.4.0

Yesterday I was using lightning v2.3.3 and it failed as well...