Fix parallelization of input data in generate() (#2749)

* Fix Distribution

* test

* Address automatic comments

* Addressing comments

* Test fix

* formatting

* Fix test

* Format and polishing

* Format

Author: dvadym

keras_hub/src/models/causal_lm.py

@@ -4,6 +4,7 @@
 import keras
 from keras import ops
 from keras import tree
+from keras.src.distribution import distribution_lib
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.task import Task
@@ -354,6 +355,26 @@ def preprocess(x):
                 x, sequence_length=max_length
             )
 
+        def distribute(x):
+            """Distribute tensors according to the distribution library."""
+            distribution = distribution_lib.distribution()
+            if distribution is None:
+                return x
+
+            def _distribute_tensor(value):
+                if value is None:
+                    return None
+                if not ops.is_tensor(value):
+                    value = ops.convert_to_tensor(value)
+                layout = distribution.get_data_layout(value.shape)
+                return (
+                    distribution_lib.distribute_tensor(value, layout)
+                    if layout
+                    else value
+                )
+
+            return tree.map_structure(_distribute_tensor, x)
+
         def generate(x):
             return generate_function(x, stop_token_ids=stop_token_ids)
 
@@ -394,6 +415,8 @@ def postprocess(x):
         if self.preprocessor is not None:
             inputs = [preprocess(x) for x in inputs]
 
+        inputs = [distribute(x) for x in inputs]
+
         if strip_prompt:
             outputs = [strip_prompt_function(generate(x), x) for x in inputs]
         else:

keras_hub/src/models/gpt2/gpt2_causal_lm_test.py

@@ -244,3 +244,85 @@ def test_get_quantization_layer_structure(self):
         )
 
         self.assertIsNone(causal_lm.get_quantization_layer_structure("int8"))
+
+
+@pytest.mark.skipif(keras.src.backend.backend() != "jax", reason="JAX only")
+@pytest.mark.multi_device
+class GPT2CausalLMDistributionTest(TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.device_count = len(keras.distribution.list_devices())
+
+        # Initialize kwargs for model creation
+        self.merges = ["Ġ a", "Ġ t", "Ġ i", "Ġ b", "a i", "p l", "n e"]
+        self.merges += ["Ġa t", "p o", "r t", "Ġt h", "ai r", "pl a", "po rt"]
+        self.merges += ["Ġai r", "Ġa i", "pla ne"]
+        self.vocab = []
+        for merge in self.merges:
+            a, b = merge.split(" ")
+            self.vocab.extend([a, b, a + b])
+        self.vocab += ["!", "<|endoftext|>"]
+        self.vocab = sorted(set(self.vocab))  # Remove duplicates
+        self.vocab = dict([(token, i) for i, token in enumerate(self.vocab)])
+        self.vocabulary_size = len(self.vocab)
+
+        self.tokenizer = GPT2Tokenizer(
+            vocabulary=self.vocab,
+            merges=self.merges,
+        )
+        self.preprocessor = GPT2CausalLMPreprocessor(
+            tokenizer=self.tokenizer,
+            sequence_length=8,
+        )
+        self.backbone = GPT2Backbone(
+            vocabulary_size=self.vocabulary_size,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=4,
+            intermediate_dim=4,
+            max_sequence_length=8,
+        )
+        self.init_kwargs = {
+            "backbone": self.backbone,
+            "preprocessor": self.preprocessor,
+        }
+
+        self.device_mesh = keras.distribution.DeviceMesh(
+            shape=(self.device_count,),
+            axis_names=["batch"],
+            devices=keras.distribution.list_devices(),
+        )
+
+        self.layout_map = keras.distribution.LayoutMap(self.device_mesh)
+
+        self.distribution = keras.distribution.DataParallel(
+            self.device_mesh,
+            self.layout_map,
+        )
+
+    def test_e2e_data_parallel_generate(self):
+        with self.distribution.scope():
+            causal_lm = GPT2CausalLM(**self.init_kwargs)
+
+            # Pass prompts to match the number of devices.
+            prompts = [" airplane at airport"] * self.device_count
+
+            # This should run without errors and use the distribution.
+            output = causal_lm.generate(prompts)
+            self.assertIsInstance(output, (list, tuple))
+            self.assertLen(output, self.device_count)
+
+    def test_e2e_data_parallel_generate_indivisible_error(self):
+        with self.distribution.scope():
+            if self.device_count < 2:
+                pytest.skip("This test requires at least 2 devices")
+            keras.distribution.set_distribution(self.distribution)
+            causal_lm = GPT2CausalLM(**self.init_kwargs)
+
+            # Pass only 1 prompt, which is not divisible by the number of
+            # devices.
+            prompt = " airplane at airport"
+
+            with self.assertRaises(Exception):
+                causal_lm.generate(prompt)