Skip to content

bug: TransformersRunnable.__init__ sets torch default tensor type to torch.cuda.FloatTensor even if gpu is not available #4376

Description

@david921023

Describe the bug

class TransformersRunnable(bentoml.Runnable):
        SUPPORTED_RESOURCES = ("nvidia.com/gpu", "cpu")
        SUPPORTS_CPU_MULTI_THREADING = True

        def __init__(self):
            super().__init__()

            available_gpus = os.getenv("CUDA_VISIBLE_DEVICES", "")
            # assign CPU resources
            kwargs = {}

            if available_gpus not in ("", "-1"):
                # assign GPU resources
                if not available_gpus.isdigit():
                    raise ValueError(
                        f"Expecting numeric value for CUDA_VISIBLE_DEVICES, got {available_gpus}."
                    )
                if "_pretrained_class" not in bento_model.info.metadata:
                    # NOTE: then this is a pipeline. We then pass the device to it.
                    kwargs["device"] = int(available_gpus)
            if "_pretrained_class" not in bento_model.info.metadata:
                self.model = load_model(bento_model, **kwargs)
            else:
                if "_framework" in bento_model.info.metadata:
                    if (
                        "torch" == bento_model.info.metadata["_framework"]
                        or "pt" == bento_model.info.metadata["_framework"]
                    ):
                        self.model = t.cast(
                            transformers.PreTrainedModel,
                            load_model(bento_model, **kwargs),
                        ).to(
                            torch.device(
                                "cuda" if available_gpus not in ("", "-1") else "cpu"
                            )
                        )
                        torch.set_default_tensor_type("torch.cuda.FloatTensor")

Look at the very last line of this code. It tries to set default_tensor_type to torch.cuda.FloatTensor regardless of available_gpus. This is causing the following error when using cpu

File "/opt/homebrew/lib/python3.11/site-packages/bentoml/_internal/runner/runner.py", line 317, in init_local
    raise e
  File "/opt/homebrew/lib/python3.11/site-packages/bentoml/_internal/runner/runner.py", line 307, in init_local
    self._set_handle(LocalRunnerRef)
  File "/opt/homebrew/lib/python3.11/site-packages/bentoml/_internal/runner/runner.py", line 150, in _set_handle
    runner_handle = handle_class(self, *args, **kwargs)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/bentoml/_internal/runner/runner_handle/local.py", line 27, in __init__
    self._runnable = runner.runnable_class(**runner.runnable_init_params)  # type: ignore
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/lib/python3.11/site-packages/bentoml/_internal/frameworks/transformers.py", line 1205, in __init__
    torch.set_default_tensor_type("torch.cuda.FloatTensor")
  File "/opt/homebrew/lib/python3.11/site-packages/torch/__init__.py", line 614, in set_default_tensor_type
    _C._set_default_tensor_type(t)
TypeError: type torch.cuda.FloatTensor not available. Torch not compiled with CUDA enabled.

The issue is mitigated as I manually added line to check available_gpus before calling torch.set_default_tensor_type("torch.cuda.FloatTensor").

Please take a look at this issue.

To reproduce

No response

Expected behavior

No response

Environment

bentoml==1.1.11
torch==2.1.2
python@3.11

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions