Repository: huggingface/transformers
Pull request: [CB] [Bug] Fix crashes when running without cuda

This PR fixes a bug in continuous batching where non-CUDA devices cannot use the feature because some CUDA-exclusive objects are always instantiated.
It also adds a test to make sure this will not break again in the future.

Changed files:
filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -17,7 +17,7 @@
 import threading
 from abc import abstractmethod
 from collections.abc import Generator
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from math import ceil
 from time import perf_counter
 
@@ -161,7 +161,7 @@ def __init__(
             )
         # Set up the graph pool. This allows all graphs to share the same memory pool, which is fine because they never
         # run concurrently. This greatly saves memory.
-        self.graph_pool = torch.cuda.graph_pool_handle()
+        self.graph_pool = torch.cuda.graph_pool_handle() if self.use_cuda_graph else None
 
     def __repr__(self) -> str:
         return (
@@ -374,7 +374,9 @@ def update_batch(self) -> None:
         if copy_source:
             # FIXME: this will avoid any race condition, but it can cause issue when using async batching with a sliding
             # window model. Fix will be fixed in a PR in the near future (tempfix, v5.3)
-            with torch.cuda.stream(self.inputs_and_outputs.compute_stream):
+            compute_stream = self.inputs_and_outputs.compute_stream
+            maybe_stream = torch.cuda.stream(compute_stream) if compute_stream is not None else nullcontext()
+            with maybe_stream:
                 self.cache.copy_cache(copy_source, copy_destination)
 
     @traced
@@ -434,10 +436,11 @@ def _generation_step(self, model: nn.Module, logit_processor: LogitsProcessorLis
 
         # If we are not using cuda graphs, we perform the generation step and return
         if not self.use_cuda_graph:
-            with torch.cuda.stream(compute_stream):
+            maybe_stream = torch.cuda.stream(compute_stream) if compute_stream is not None else nullcontext()
+            with maybe_stream:
                 self._forward_process_and_sample(model, batch_data, logit_processor, do_sample)
 
-        # Otherwise, we use create or replay the graph
+        # Otherwise, we use create or replay the graph (cuda is available in this path)
         else:
             graph = self.inputs_and_outputs.get_graph()
             # Case: the graph already exists, so we replay it

filename: src/transformers/generation/continuous_batching/input_outputs.py
patch: @@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from contextlib import nullcontext
 from dataclasses import dataclass
 from functools import partial
 from itertools import count
@@ -221,7 +222,8 @@ def _transfer_inputs(
         other.max_seqlen_q = self.max_seqlen_q
         other.max_seqlen_k = dict(self.max_seqlen_k.items())
         # Transfer static tensors
-        with torch.cuda.stream(stream):
+        maybe_stream = torch.cuda.stream(stream) if stream is not None else nullcontext()
+        with maybe_stream:
             other._bulk_input_tensor.copy_(self._bulk_input_tensor, non_blocking=non_blocking)  # fast bulk transfer
             # Only transfer block_table for decode-only batches (when it's actually used)
             if self.use_block_table:
@@ -496,15 +498,17 @@ def __init__(
         # The host IO has automatic pinned memory because it is created on the CPU
         self.host_io = ContinuousBatchingIOs(cache, config, torch.device("cpu"), model_dtype, max_graphs)
         self.device_io = ContinuousBatchingIOs(cache, config, device, model_dtype, max_graphs)
-        self.h2d_over = torch.cuda.Event()
-        self.compute_over = torch.cuda.Event()
-        self.d2h_over = torch.cuda.Event()
+        # Create events only on CUDA devices
+        self.h2d_over = torch.cuda.Event() if torch.cuda.is_available() else None
+        self.compute_over = torch.cuda.Event() if torch.cuda.is_available() else None
+        self.d2h_over = torch.cuda.Event() if torch.cuda.is_available() else None
 
     def transfer_inputs_h2d(self, stream: torch.cuda.Stream) -> None:
         self.host_io._transfer_inputs(self.device_io, stream=stream, non_blocking=True)
 
-    def transfer_outputs_d2h(self, stream: torch.cuda.Stream) -> None:
-        with torch.cuda.stream(stream):
+    def transfer_outputs_d2h(self, stream: torch.cuda.Stream | None) -> None:
+        maybe_stream = torch.cuda.stream(stream) if stream is not None else nullcontext()
+        with maybe_stream:
             self.host_io.output_ids.copy_(self.device_io.output_ids, non_blocking=True)
 
 
@@ -562,6 +566,9 @@ def __init__(
         model_dtype: torch.dtype,
         max_graphs: int = 32,
     ) -> None:
+        # Async batching needs streams to function, so check is CUDA is available
+        if not torch.cuda.is_available():
+            raise RuntimeError(f"Async batching requires CUDA, but {torch.cuda.is_available() = }")
         # IO pairs used to avoid race conditions
         self.current_pair = 0
         self.io_pairs = [HostDeviceIOPair(cache, config, device, model_dtype, max_graphs) for _ in range(2)]
@@ -668,5 +675,5 @@ def retrieve_device_outputs(self) -> None:
     # This method is called after the switch and not during the first batch
     def prepare_batch_update(self) -> tuple[list[FutureRequestState], list[int]]:
         io_pair = self.io_pairs[self.current_pair]
-        io_pair.d2h_over.synchronize()
+        io_pair.d2h_over.synchronize()  # ty:ignore[unresolved-attribute]  <- this is always a CUDA event
         return io_pair.host_io.prepare_batch_update()

filename: tests/generation/test_continuous_batching.py
patch: @@ -46,7 +46,7 @@
     slow,
     torch_device,
 )
-from transformers.utils import is_flash_attn_2_available, is_kernels_available
+from transformers.utils import is_flash_attn_2_available, is_kernels_available, is_torch_xpu_available
 
 
 def flush_memory(flush_compile: bool = True) -> None:
@@ -77,7 +77,7 @@ def flush_memory(flush_compile: bool = True) -> None:
     gc.collect()
 
 
-class ContinuousBatchingNonGenerationTest(unittest.TestCase):
+class ContinuousBatchingNoAcceleratorTest(unittest.TestCase):
     @parameterized.expand(
         [
             (None, None, "0"),
@@ -341,9 +341,57 @@ def test_full_attention_get_write_indices(
             f"Failed for {block_size=}, {block_table=}, {past_length=}, {query_length=}",
         )
 
+    @slow
+    def test_continuous_batching_no_accelerators(self) -> None:
+        """Test continuous batching generation when no accelerator is available. It uses a simulated CPU-only PyTorch
+        environment by mocking all acceleratoravailability checks to return False"""
+        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+        # Mock all accelerator availability checks to simulate CPU-only PyTorch
+        with (
+            patch("torch.cuda.is_available", return_value=False),
+            patch("transformers.utils.is_torch_xpu_available", return_value=False),
+            patch("torch.backends.mps.is_available", return_value=False),
+        ):
+            # Verify patches work
+            self.assertFalse(torch.cuda.is_available())
+            self.assertFalse(is_torch_xpu_available())
+            self.assertFalse(torch.backends.mps.is_available())
+
+            tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side="left")
+            tokenizer.pad_token = tokenizer.eos_token
+
+            user_messages = [
+                "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?"
+            ]
+            chats = [[{"role": "user", "content": user_message}] for user_message in user_messages]
+            tokenized = [tokenizer.apply_chat_template(chat, add_generation_prompt=True) for chat in chats]
+            input_ids = [(x if isinstance(x, list) else x["input_ids"]) for x in tokenized]
+
+            # Load model on CPU
+            model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="sdpa")
+            model = model.to("cpu").eval()
+            model.generation_config.max_new_tokens = 10
+            model.generation_config.do_sample = False
+
+            continuous_batching_config = ContinuousBatchingConfig(use_cuda_graph=False, use_async_batching=False)
+
+            # This should not crash even with all accelerators unavailable
+            outputs = model.generate_batch(
+                inputs=input_ids,
+                generation_config=model.generation_config,
+                continuous_batching_config=continuous_batching_config,
+            )
+
+            # Verify we got outputs
+            self.assertEqual(len(outputs), len(input_ids))
+            for output in outputs.values():
+                self.assertIsNotNone(output.generated_tokens)
+                self.assertGreater(len(output.generated_tokens), 0)
+
 
 @require_torch_accelerator
-class ContinuousBatchingGenerationTest(unittest.TestCase):
+class ContinuousBatchingWithAcceleratorTest(unittest.TestCase):
     # -----------------------------------------------Parity tests----------------------------------------------- #
     #         Ensure continuous batching and non-continuous batching generation produce the same outputs         #
     # ---------------------------------------------------------------------------------------------------------- #