Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/strategies/launchers/multiprocessing.py (Python)

Code before:
    Lightning users.

    """
    # Use PyTorch's internal check for bad fork state, which is more accurate than just checking if CUDA
    # is initialized. This allows passive CUDA initialization (e.g., from library imports or device queries)
    # while still catching actual problematic cases where CUDA context was created before forking.
    _is_in_bad_fork = getattr(torch.cuda, "_is_in_bad_fork", None)
    if _is_in_bad_fork is not None and _is_in_bad_fork():
        message = (
            "Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, "
            "you must use the 'spawn' start method or avoid CUDA initialization in the main process."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)

    # Fallback to the old check if _is_in_bad_fork is not available (older PyTorch versions)
    if _is_in_bad_fork is None and torch.cuda.is_initialized():
        message = (
            "Lightning can't create new processes if CUDA is already initialized. Did you manually call"
            " `torch.cuda.*` functions, have moved the model to the device, or allocated memory on the GPU any"
            " other way? Please remove any such calls, or change the selected strategy."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)


def _disable_module_memory_sharing(data: Any) -> Any:
    """Disables memory sharing on parameters and buffers of `nn.Module`s contained in the given collection.

    Note: This is only required when running on CPU.

    """
    # PyTorch enables memory sharing automatically on all tensors that are passed through `mp.spawn`.
    # For model weights and buffers, this is undesired and can lead to race conditions between processes.
    # Hence, we copy the tensors in the entire module to ensure it doesn't share memory with other processes.

    @torch.no_grad()
    def unshare(module: Module) -> Module:
        for tensor in itertools.chain(module.parameters(), module.buffers()):
            tensor.data = tensor.data.clone()
        return module

    return apply_to_collection(data, function=unshare, dtype=Module)


def _check_missing_main_guard() -> None:
    """Raises an exception if the ``__name__ == "__main__"`` guard is missing."""
    if not getattr(mp.current_process(), "_inheriting", False):
        return

Review comment (suggestion):
The existing test `test_check_for_bad_cuda_fork` mocks `torch.cuda.is_initialized()` to return `True`, which will only test the fallback path in the new implementation (when `_is_in_bad_fork` is None). This test should be updated to also verify the new behavior when `torch.cuda._is_in_bad_fork` is available and returns `True`.

Code after:
    Lightning users.

    """
    # Use PyTorch's internal check for bad fork state, which is more accurate than just checking if CUDA
    # is initialized. This allows passive CUDA initialization (e.g., from library imports or device queries)
    # while still catching actual problematic cases where CUDA context was created before forking.
    _is_in_bad_fork = getattr(torch.cuda, "_is_in_bad_fork", None)
    if _is_in_bad_fork is not None and callable(_is_in_bad_fork) and _is_in_bad_fork():
        message = (
            "Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, "
            "you must use the 'spawn' start method or avoid CUDA initialization in the main process."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)

    # Fallback to the old check if _is_in_bad_fork is not available (older PyTorch versions)
    if _is_in_bad_fork is None and torch.cuda.is_initialized():
        message = (
            "Lightning can't create new processes if CUDA is already initialized. Did you manually call"
            " `torch.cuda.*` functions, have moved the model to the device, or allocated memory on the GPU any"
            " other way? Please remove any such calls, or change the selected strategy."
        )
        if _IS_INTERACTIVE:
            message += " You will have to restart the Python kernel."
        raise RuntimeError(message)


def _disable_module_memory_sharing(data: Any) -> Any:
    """Disables memory sharing on parameters and buffers of `nn.Module`s contained in the given collection.

    Note: This is only required when running on CPU.

    """
    # PyTorch enables memory sharing automatically on all tensors that are passed through `mp.spawn`.
    # For model weights and buffers, this is undesired and can lead to race conditions between processes.
    # Hence, we copy the tensors in the entire module to ensure it doesn't share memory with other processes.

    @torch.no_grad()
    def unshare(module: Module) -> Module:
        for tensor in itertools.chain(module.parameters(), module.buffers()):
            tensor.data = tensor.data.clone()
        return module

    return apply_to_collection(data, function=unshare, dtype=Module)


def _check_missing_main_guard() -> None:
    """Raises an exception if the ``__name__ == "__main__"`` guard is missing."""
    if not getattr(mp.current_process(), "_inheriting", False):
        return