[JAX] Accelerate multi-step training loops using on-device jax.lax.scan for steps_per_execution > 1 - #23527
[JAX] Accelerate multi-step training loops using on-device jax.lax.scan for steps_per_execution > 1#23527gaga1313 wants to merge 7 commits into
jax.lax.scan for steps_per_execution > 1#23527Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements support for multi-step execution (steps_per_execution > 1) in the JAX backend, introducing host-side super-batching, JAX-native scanning via jax.lax.scan, and fallback unrolling for partial batches. The review feedback highlights several critical robustness issues where the implementation assumes non-null values in nested batch structures. Specifically, the reviewer pointed out that module-level @jit on _concatenate_outputs bypasses eager execution settings, and multiple utility functions (such as tree.map_structure and layout mapping) will crash with TypeError or AttributeError if optional batch elements (like sample_weight) are None. Adding proper guards for None values is highly recommended to ensure stability.
| @jit | ||
| def _concatenate_outputs(outputs): | ||
| if not outputs: | ||
| return [] | ||
| if len(outputs) == 1: | ||
| return outputs[0] | ||
| return tree.map_structure( | ||
| lambda *args: jax.numpy.concatenate(args, axis=0), | ||
| *outputs, | ||
| ) |
There was a problem hiding this comment.
Unconditionally decorating _concatenate_outputs with @jit at the module level violates run_eagerly=True (or jit_compile=False) and can cause unexpected compilation overhead. Additionally, if any of the outputs are None (e.g., optional outputs or state), jax.numpy.concatenate will raise a TypeError. We should remove the module-level @jit decorator and handle None values safely.
def _concatenate_outputs(outputs):
if not outputs:
return []
if len(outputs) == 1:
return outputs[0]
return tree.map_structure(
lambda *args: jax.numpy.concatenate(args, axis=0) if args[0] is not None else None,
*outputs,
)| leaf = tree.flatten(batch)[0] | ||
| if leaf.shape[0] < self.steps_per_execution: | ||
| sliced_batches = [ | ||
| tree.map_structure(lambda x, i=i: x[i], batch) | ||
| for i in range(leaf.shape[0]) | ||
| ] | ||
| return _unroll_steps(state, sliced_batches) |
There was a problem hiding this comment.
In Case 2, tree.flatten(batch)[0] is fragile because the first element of the flattened batch could be None (e.g., if the batch structure starts with an optional element). Using next(x for x in tree.flatten(batch) if x is not None) is much more robust. Furthermore, tree.map_structure(lambda x, i=i: x[i], batch) will crash with a TypeError if any element in batch is None (such as sample_weight=None or y=None during prediction). We should check if x is not None before indexing.
| leaf = tree.flatten(batch)[0] | |
| if leaf.shape[0] < self.steps_per_execution: | |
| sliced_batches = [ | |
| tree.map_structure(lambda x, i=i: x[i], batch) | |
| for i in range(leaf.shape[0]) | |
| ] | |
| return _unroll_steps(state, sliced_batches) | |
| leaf = next((x for x in tree.flatten(batch) if x is not None), None) | |
| if leaf is not None and leaf.shape[0] < self.steps_per_execution: | |
| sliced_batches = [ | |
| tree.map_structure( | |
| lambda x, i=i: x[i] if x is not None else None, | |
| batch, | |
| ) | |
| for i in range(leaf.shape[0]) | |
| ] | |
| return _unroll_steps(state, sliced_batches) |
| sliced_batches = [ | ||
| tree.map_structure(lambda x, i=i: x[i], batch) | ||
| for i in range(leaf.shape[0]) | ||
| ] | ||
| return _unroll_steps(state, sliced_batches) |
There was a problem hiding this comment.
In Case 4, tree.map_structure(lambda x, i=i: x[i], batch) will crash with a TypeError if any element in batch is None (such as sample_weight=None or y=None during prediction). We should check if x is not None before indexing.
| sliced_batches = [ | |
| tree.map_structure(lambda x, i=i: x[i], batch) | |
| for i in range(leaf.shape[0]) | |
| ] | |
| return _unroll_steps(state, sliced_batches) | |
| if leaf is not None: | |
| sliced_batches = [ | |
| tree.map_structure( | |
| lambda x, i=i: x[i] if x is not None else None, | |
| batch, | |
| ) | |
| for i in range(leaf.shape[0]) | |
| ] | |
| return _unroll_steps(state, sliced_batches) |
| def get_single_layout(d): | ||
| return distribution.get_data_layout(d.shape).backend_layout |
There was a problem hiding this comment.
If batches[0] contains any None elements (e.g., optional sample_weight or y), get_single_layout will crash with an AttributeError when accessing d.shape. We should add a guard to return None if d is None.
def get_single_layout(d):
if d is None:
return None
return distribution.get_data_layout(d.shape).backend_layout| def get_layout(_, d_orig): | ||
| base_layout = distribution.get_data_layout( | ||
| d_orig.shape | ||
| ) |
There was a problem hiding this comment.
If batches[0] contains any None elements, get_layout will crash with an AttributeError when accessing d_orig.shape. We should add a guard to return None if d_orig is None.
| def get_layout(_, d_orig): | |
| base_layout = distribution.get_data_layout( | |
| d_orig.shape | |
| ) | |
| def get_layout(_, d_orig): | |
| if d_orig is None: | |
| return None | |
| base_layout = distribution.get_data_layout( | |
| d_orig.shape | |
| ) |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23527 +/- ##
==========================================
- Coverage 84.29% 84.22% -0.07%
==========================================
Files 468 468
Lines 71096 71191 +95
Branches 11788 11820 +32
==========================================
+ Hits 59931 59962 +31
- Misses 8164 8215 +51
- Partials 3001 3014 +13
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
gaga1313
changed the title
jax.lax.scan for steps_per_execution > 1jax.lax.scan for steps_per_execution > 1
| return tree.map_structure( | ||
| lambda *args: jax.numpy.concatenate(args, axis=0), | ||
| *outputs, | ||
| ) |
There was a problem hiding this comment.
Why not just return ax.numpy.concatenate(outputs, axis=0)?
Line 41 assumes it's a list or tuple already.
There was a problem hiding this comment.
This function was already in the library. The line 41 only checks if the nested output is list of length 1 then return the first outputs[0].
jax.numpy.concatenate will not work with the nested outputs.
| class PartialBatchList(list): | ||
| """Wrapper to distinguish a list of batches from a batch of lists/tuples.""" | ||
|
|
||
| pass |
There was a problem hiding this comment.
Why is this needed?
Doesn't the existing if on the batch size cover this?
There was a problem hiding this comment.
If you are talking about if statements on line 300 and 304, they are different. I have removed the PartialBatchList, and simply replaced the if condition with to check if batch is a list.
| step_function, | ||
| raw_step_function=None, |
There was a problem hiding this comment.
The contract for _make_function was basically to add support for steps_per_execution. The jitting support was added beforehand and step_function was already jitted if applicable.
Now _make_function has this weird contract that takes two functions, one maybe jitted and one not jitted. _make_function jits the loop case but doesn't jit the non-loop case. _make_function uses both step_function and raw_step_function, which is confusing.
Let's clean up the contract and do:
- only
step_functionis passed and it's the non-jitted function _make_functiontakes care of the jitting (if jitting is requested) in all cases, and there is no pre-jitting before calling_make_function
There was a problem hiding this comment.
Thanks for highlighting this. Fixed!
| step_function, | ||
| raw_step_function=None, | ||
| out_shardings=None, | ||
| donate_argnums=0, |
There was a problem hiding this comment.
donate_argnums: do we need this? Isn't it always 0?
There was a problem hiding this comment.
Thanks for highlighting, I will remove it.
| if isinstance(data_batch, list): | ||
| data_batch = data_batch[0] | ||
| else: | ||
| data_batch = tree.map_structure( | ||
| lambda x: x[0] if x is not None else None, data_batch | ||
| ) |
There was a problem hiding this comment.
I think you should always do the tree.map_structure. It will work in the list case too, and it will handle the case of nested lists, which I don't think is handled correctly right now.
There was a problem hiding this comment.
If we always use tree.map_structure then in the case of PartialBatchList [b1, b2], the tree.map_structure will recursively slice into the leaves of each batch, which strips the batch dimension.
| data_batch = next(data_or_iterator) | ||
| break | ||
|
|
||
| if data_batch is not None and self.steps_per_execution > 1: |
There was a problem hiding this comment.
Don't you need to also check if super-batching was enabled?
There was a problem hiding this comment.
If steps_per_execution > 1 then super-batching is enabled by default. Either iterator would return a Super_batch (pytree) or list through get_host_stack_terator or super-batch via tf.data iterator (pytree).
| @parameterized.named_parameters( | ||
| {"testcase_name": "spe_2", "steps_per_execution": 2}, | ||
| {"testcase_name": "spe_4", "steps_per_execution": 4}, | ||
| ) | ||
| def test_steps_per_execution_numeric_equivalence(self, steps_per_execution): |
There was a problem hiding this comment.
There is nothing JAX specific about the tests you added, they should pass with all backends that support steps_per_execution. So, at the very least, they should move to https://github.com/keras-team/keras/blob/master/keras/src/trainers/trainer_test.py
In fact there are already tests for steps_per_execution in https://github.com/keras-team/keras/blob/master/keras/src/trainers/trainer_test.py
So make sure they cover the same and remove these, or augment them if something is missing.
gaga1313
left a comment
gaga1313
left a comment
There was a problem hiding this comment.
Fixed thanks! I have also removed the test from jax/trainer_test.py.
| class PartialBatchList(list): | ||
| """Wrapper to distinguish a list of batches from a batch of lists/tuples.""" | ||
|
|
||
| pass |
There was a problem hiding this comment.
If you are talking about if statements on line 300 and 304, they are different. I have removed the PartialBatchList, and simply replaced the if condition with to check if batch is a list.
| step_function, | ||
| raw_step_function=None, | ||
| out_shardings=None, | ||
| donate_argnums=0, |
There was a problem hiding this comment.
Thanks for highlighting, I will remove it.
| step_function, | ||
| raw_step_function=None, |
There was a problem hiding this comment.
Thanks for highlighting this. Fixed!
| data_batch = next(data_or_iterator) | ||
| break | ||
|
|
||
| if data_batch is not None and self.steps_per_execution > 1: |
There was a problem hiding this comment.
If steps_per_execution > 1 then super-batching is enabled by default. Either iterator would return a Super_batch (pytree) or list through get_host_stack_terator or super-batch via tf.data iterator (pytree).
| if isinstance(data_batch, list): | ||
| data_batch = data_batch[0] | ||
| else: | ||
| data_batch = tree.map_structure( | ||
| lambda x: x[0] if x is not None else None, data_batch | ||
| ) |
There was a problem hiding this comment.
If we always use tree.map_structure then in the case of PartialBatchList [b1, b2], the tree.map_structure will recursively slice into the leaves of each batch, which strips the batch dimension.
4 participants
Description
Context & Problem
In Keras 3 with the JAX backend, setting$N$ dispatches for $N$ steps).
steps_per_execution > 1previously executed a Python host-level loop around single-step compiled executables. On accelerator hardware (GPU/TPU) as well as multi-core CPUs, this incurred Python host dispatch overhead, buffer allocation round-trips, and device-host synchronization barriers on every single step (As shown in the baseline benchmarks below, increasing
steps_per_executiononmasteryielded virtually flat performance (~1.1x to 1.25x on GPU) because host-side dispatch latency was the primary bottleneck.Proposed Solution
This PR introduces true on-device multi-step execution using
jax.lax.scan:jax.lax.scanCompilation: Whensteps_per_execution > 1andjit_compile=True, batches are stacked into super-batches along axis 0 and executed continuously in hardware memory (GPU HBM) with only 1 host dispatch perPartialBatchListwithout shape crashes or dropped samples.tf.data.Dataset: Automatically batches bysteps_per_executionon the host (.batch(SPE).prefetch(AUTOTUNE)) with preserved_super_batchedlifecycle state across epochs._get_host_stacked_iteratorwith clean fallbacks.steps_per_execution=1: The default single-step execution path remains completely untouched.Detailed Performance Benchmarks
1. Infra: CPU
A.
tf.data.Dataset(CPU)SPEB. NumPy Dataset / Generator (CPU)
SPE2. Infra: GPU
A.
tf.data.Dataset(GPU)SPEB. NumPy Dataset / Generator (GPU)
SPEVisualizations & Screenshots
Key Takeaways & Conclusions
master, increasingsteps_per_executionfrom 1 to 64 produced virtually no throughput improvement on GPU (~1.25x ontf.data, ~1.15x on NumPy).tf.data.Dataset:jax.lax.scanand pre-batching.batch(SPE).prefetch(AUTOTUNE)), host dispatch stalls are completely eliminated.np.stack,tf.datais recommended for maximum throughput.steps_per_executionRange:SPE=16toSPE=32, balancing maximum accelerator saturation against memory footprint and compilation time.Testing
Added comprehensive unit tests in
keras/src/backend/jax/trainer_test.py:test_steps_per_execution_numeric_equivalence: Verifies strict numerical equivalence of weights and losses betweensteps_per_execution=1andsteps_per_execution > 1.test_steps_per_execution_remainder_batches: Validates that datasets with non-divisible remainder batches execute seamlessly acrossfit,predict, andevaluate.test_steps_per_execution_tf_dataset_functional_model: Verifies end-to-end multi-step training withtf.data.Datasetpipelines on Functional models across multiple epochs.Ran full test suites:
pytest keras/src/backend/jax/trainer_test.py(12/12 passed)pytest keras/src/trainers/trainer_test.py(162/162 passed)ruff.Contributor Agreement