diff --git a/export/orbax/export/obm_configs.py b/export/orbax/export/obm_configs.py index 1dab7808d..0d783d541 100644 --- a/export/orbax/export/obm_configs.py +++ b/export/orbax/export/obm_configs.py @@ -196,6 +196,9 @@ class BatchOptions: all batch sizes no larger than `max_batch_size` are allowed. Otherwise, supplies a list of batch sizes. The entries must increase monotonically. disable_large_batch_splitting: Whether to disable large batch splitting. + pad_variable_length_inputs: Whether to dynamically pad variable-length + non-batch dimensions to the maximum size across requests within each + batch. Requires disable_large_batch_splitting = true. batch_padding_policy: The batch padding policy for the batch scheduler. Default is PAD_UP. low_priority_batch_options: The batch options for low priority inputs. @@ -218,6 +221,7 @@ class BatchOptions: num_batch_threads: int = 1 max_enqueued_batches: int = 250 disable_large_batch_splitting: bool = False + pad_variable_length_inputs: bool = False batch_padding_policy: BatchPaddingPolicy = BatchPaddingPolicy.PAD_UP low_priority_batch_options: LowPriorityBatchOptions | None = None mixed_priority_batching_policy: MixedPriorityBatchingPolicy = ( @@ -317,6 +321,15 @@ def __post_init__(self): if self.batch_component == BatchComponent.NO_BATCHING: return + if ( + self.pad_variable_length_inputs + and not self.disable_large_batch_splitting + ): + raise ValueError( + "disable_large_batch_splitting must be True when" + " pad_variable_length_inputs is True." + ) + if self.max_batch_size is None: if self.allowed_batch_sizes: self.max_batch_size = self.allowed_batch_sizes[-1] diff --git a/export/orbax/export/obm_configs_test.py b/export/orbax/export/obm_configs_test.py index b3b78d540..af4c7cab4 100644 --- a/export/orbax/export/obm_configs_test.py +++ b/export/orbax/export/obm_configs_test.py @@ -26,6 +26,19 @@ def test_batch_options_disable_batch(self): ) self.assertIsNone(batch_options.max_batch_size) + def test_pad_variable_length_inputs_with_large_batch_splitting_fails(self): + with self.assertRaisesRegex( + ValueError, + "disable_large_batch_splitting must be True when" + " pad_variable_length_inputs is True.", + ): + obm_configs.BatchOptions( + batch_component=obm_configs.BatchComponent.MODEL_FUNCTION, + max_batch_size=8, + pad_variable_length_inputs=True, + disable_large_batch_splitting=False, + ) + def test_batch_options_raise_error_without_max_batch_size_and_allowed_batch_sizes( self, ):