diff --git a/src/transformers/generation/configuration_utils.py b/src/transformers/generation/configuration_utils.py index 641201ef384f..21c3fa442e49 100644 --- a/src/transformers/generation/configuration_utils.py +++ b/src/transformers/generation/configuration_utils.py @@ -457,7 +457,8 @@ def __init__(self, **kwargs): self.compile_config = kwargs.pop("compile_config", None) self.disable_compile = kwargs.pop("disable_compile", None) - self.continuous_batching_config = kwargs.pop("continuous_batching_config", None) + # Continuous batching config is runtime/hardware-specific and lives beside `GenerationConfig`. + kwargs.pop("continuous_batching_config", None) # Deprecated (moved to the Hub). TODO remove for v5 self.low_memory = kwargs.pop("low_memory", None) @@ -1798,6 +1799,11 @@ class ContinuousBatchingConfig: max_cached_graphs: int | None = None def __post_init__(self): + if isinstance(self.varlen_compile_config, dict): + self.varlen_compile_config = CompileConfig(**self.varlen_compile_config) + if isinstance(self.decode_compile_config, dict): + self.decode_compile_config = CompileConfig(**self.decode_compile_config) + # Only turn off graph mixing support if TP is on graph_mixing_supported = os.environ.get("NCCL_GRAPH_MIXING_SUPPORT", "1") == "1" distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 diff --git a/src/transformers/generation/continuous_batching/continuous_api.py b/src/transformers/generation/continuous_batching/continuous_api.py index 1c405a287a18..882416880345 100644 --- a/src/transformers/generation/continuous_batching/continuous_api.py +++ b/src/transformers/generation/continuous_batching/continuous_api.py @@ -1104,10 +1104,7 @@ def init_continuous_batching( # Retrieve continuous batching config, or create it if none is provided if continuous_batching_config is None: - if isinstance(getattr(gen_config, "continuous_batching_config", None), ContinuousBatchingConfig): - continuous_batching_config = gen_config.continuous_batching_config - else: - continuous_batching_config = ContinuousBatchingConfig() + continuous_batching_config = ContinuousBatchingConfig() # Create and return the manager return ContinuousBatchingManager( diff --git a/tests/generation/test_configuration_utils.py b/tests/generation/test_configuration_utils.py index 36ddf4844d54..ea1f1fc012d4 100644 --- a/tests/generation/test_configuration_utils.py +++ b/tests/generation/test_configuration_utils.py @@ -18,11 +18,19 @@ import tempfile import unittest import warnings +from dataclasses import asdict from huggingface_hub import create_pull_request from parameterized import parameterized -from transformers import AutoConfig, GenerationConfig, WatermarkingConfig, is_torch_available +from transformers import ( + AutoConfig, + CompileConfig, + ContinuousBatchingConfig, + GenerationConfig, + WatermarkingConfig, + is_torch_available, +) from transformers import logging as transformers_logging @@ -755,6 +763,33 @@ def test_serialize_generation_watermarking_config(self): self.assertEqual(watermark.seeding_scheme, seeding_scheme) self.assertEqual(watermark.context_width, context_width) + def test_continuous_batching_config_from_dict(self): + """Tests that ContinuousBatchingConfig round-trips nested CompileConfig parameters.""" + continuous_batching_config = ContinuousBatchingConfig( + block_size=128, + default_compile_level=2, + varlen_compile_config=CompileConfig(dynamic=True), + decode_compile_config=CompileConfig(mode="default"), + ) + new_config = ContinuousBatchingConfig(**asdict(continuous_batching_config)) + + self.assertEqual(new_config.block_size, 128) + self.assertEqual(new_config.default_compile_level, 2) + self.assertIsInstance(new_config.varlen_compile_config, CompileConfig) + self.assertTrue(new_config.varlen_compile_config.dynamic) + self.assertNotIn("_compile_all_devices", new_config.varlen_compile_config.to_dict()) + self.assertIsInstance(new_config.decode_compile_config, CompileConfig) + self.assertEqual(new_config.decode_compile_config.mode, "default") + + def test_generation_config_does_not_serialize_continuous_batching_config(self): + """Tests that ContinuousBatchingConfig lives beside GenerationConfig.""" + continuous_batching_config = ContinuousBatchingConfig(block_size=128) + generation_config = GenerationConfig(continuous_batching_config=continuous_batching_config) + + self.assertFalse(hasattr(generation_config, "continuous_batching_config")) + self.assertNotIn("continuous_batching_config", generation_config.to_dict()) + self.assertNotIn("continuous_batching_config", generation_config.to_json_string()) + @is_staging_test class ConfigPushToHubTester(unittest.TestCase):