perf(torch): direct F.linear/einsum dispatch in Dense and EinsumDense#23185
perf(torch): direct F.linear/einsum dispatch in Dense and EinsumDense#23185pctablet505 wants to merge 13 commits into
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![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 introduces a PyTorch-specific fast path for the Dense and EinsumDense layers to bypass Keras ops dispatch overhead when LoRA is disabled, along with numeric-equivalence tests. The review feedback correctly identifies that this fast path will fail with a TypeError during symbolic tracing (e.g., when building a Functional model) because the inputs are symbolic KerasTensor objects rather than concrete PyTorch tensors. To resolve this, the reviewer suggests bypassing the fast path when inputs is an instance of backend.KerasTensor and provides corresponding test cases to verify this behavior.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #23185 +/- ##
==========================================
- Coverage 84.70% 84.08% -0.62%
==========================================
Files 465 465
Lines 69335 69770 +435
Branches 11407 11521 +114
==========================================
- Hits 58729 58668 -61
- Misses 7672 8153 +481
- Partials 2934 2949 +15
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:
|
3457205 to
616391d
Compare
616391d to
d0da7f7
Compare
4190f7a to
383afb9
Compare
Dense/EinsumDense's torch fast path aligned inputs (and, for Dense, bias by dtype only) to the kernel's device, ignoring an active keras.device(...) scope. The slow path (ops.matmul/ops.einsum -> convert_to_tensor) always resolves its target through get_device(), which reflects the scope, so the two paths diverged whenever a scope differed from the kernel's device. Resolve the target device through the scope when one is active, falling back to the kernel's device otherwise (unchanged, zero-cost common case). Move kernel, inputs, and bias to that target uniformly.
383afb9 to
3b4233a
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a PyTorch fast-path for the "Dense" and "EinsumDense" layers, bypassing the standard Keras ops dispatch to use "torch.nn.functional.linear" and "torch.einsum" directly when LoRA is disabled. It also includes comprehensive tests verifying numeric equivalence, device alignment, and device scope compatibility. The feedback suggests improving robustness in both layers by checking for the existence of the "compiler" attribute on "self._torch" before accessing "self._torch.compiler.is_compiling()", preventing potential "AttributeError"s in environments where "torch.compiler" is missing or mocked.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
The torch fast paths in Dense.call() and EinsumDense.call() cast the float kernel to the input's dtype on a mismatch instead of promoting, silently truncating the kernel to garbage (near-all-zero output) for integer/uint8/bool inputs. The fast path now only runs when inputs.dtype already matches kernel.dtype; any mismatch falls through to the standard ops path, which has its own correct dtype-promotion rules. Also fixes a ruff import-order violation in both files and strengthens the LoRA fast-path-bypass tests, which previously only checked non-NaN/non-Inf (passing even with a leaked fast path, since lora_kernel_b initializes to zero) to instead assign a nonzero LoRA delta and assert the output differs from a base-kernel-only computation.
Matches what \`pre-commit run --all-files --hook-stage manual\` (the CI format check) actually enforces, verified by running it locally twice to confirm convergence.
… path Rework the torch fast path in Dense and EinsumDense to follow the existing keras house convention for backend-conditional torch usage, as flagged in the gemini-code-assist review of PR keras-team#23185. - Import torch (and torch.nn.functional as F) at module scope under `if backend.backend() == "torch":`, matching the established pattern in keras/src/layers/layer.py and keras/src/layers/preprocessing/ string_lookup.py. Drop the per-instance `import torch as _torch` / `import torch.nn.functional as _torch_F` in build() and the `self._torch` / `self._torch_F` instance attributes, which were invented for this PR and used nowhere else in the codebase. The fast path now references bare `torch` / `F`. - The cached `self._torch_backend` boolean is kept: it is set once at build and read on the hot path, avoiding per-call recomputation, which is the whole point of this PR. Reviewed gemini findings and settled them empirically: - HIGH "symbolic KerasTensor crashes the fast path" (touches inputs.device with no KerasTensor guard): confirmed FALSE POSITIVE. Operation.__call__ routes any_symbolic_tensors(...) to symbolic_call -> compute_output_spec and never reaches Dense/EinsumDense.call(). Built gemini's exact case (Input(shape=(10,)) -> Dense(5)/EinsumDense) on the torch backend: it works and returns (None, 5); an instrumented Dense.call confirms call() never receives a KerasTensor, and KerasTensor has no .device attribute. The previously-removed guard was genuinely dead code; re-adding it would reintroduce per-call overhead this PR exists to remove. The suggested regression test (test_torch_fast_path_symbolic_input) is already present for both layers. - MEDIUM "hasattr(torch.compiler, ...) may AttributeError if torch.compiler is absent": not changed. This mirrors the exact existing pattern in keras/src/backend/torch/rnn.py (4 occurrences); torch is unpinned but keras 3.x effectively requires torch >= 2.1 (see trainer.py), where torch.compiler always exists. Switching to getattr(torch, "compiler", None) would diverge from house style. Behavior is unchanged. Full dense_test.py + einsum_dense_test.py pass on the torch backend (89+109 passed, 2 skipped); tensorflow/jax/numpy import and run unaffected.
…tion test_torch_fast_path_lora_stays_slow (Dense and EinsumDense) computed its base-kernel-only reference directly via torch.nn.functional.linear / torch.einsum using the test's original CPU-resident x_t against layer._kernel.value, which lives on the layer's actual device. On a CUDA machine this raises "Expected all tensors to be on the same device" before the assertion is ever reached, so the test never actually verifies what it's meant to. Mirrors the existing test_torch_fast_path_manual_kernel_check pattern, which already moves x_t to the kernel's device for the same reason. No-op on CPU-only runs (kernel and x_t already share the 'cpu' device), confirmed via KERAS_BACKEND=torch pytest locally.
The fast-path comments in Dense.call/EinsumDense.call read as a stacked
list of short topic-labeled notes ("Fast path: X. Quantization routes...
kernel is [in, out]...") rather than the flowing explanatory prose used
elsewhere in this codebase (see batch_normalization.py, layer.py). Merged
them into single comments that explain the guard as a whole, and reworded
the colon-led "Skip under torch.compile: ..." note into a plain sentence.
No code change.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces a PyTorch-specific fast-path for the Dense and EinsumDense layers. When using the PyTorch backend, the layers bypass Keras's standard ops dispatch during eager execution and directly call torch.nn.functional.linear or torch.einsum. This optimization is bypassed during compilation, when LoRA is enabled, or when there is a dtype mismatch. Comprehensive tests are added to verify numerical equivalence, device alignment, device scopes, and dtype handling. The reviewer feedback suggests improving the robustness of the PyTorch compiler check by verifying that the compiler attribute exists on the torch module before accessing torch.compiler to prevent potential AttributeErrors in environments with older PyTorch versions.
Accessing torch.compiler.is_compiling() without first checking that torch exposes a compiler attribute can raise AttributeError in environments where torch.compiler isn't bound on the torch namespace. Check hasattr(torch, "compiler") before touching the submodule.
…omment Remove the section-divider banner above the torch fast-path test group in dense_test.py and einsum_dense_test.py; the surrounding docstrings already describe what the tests cover. Also reword a comment in einsum_dense_test.py that referenced "this PR" to describe current torch.einsum behavior directly.
… imports Collapse the two gating conditions (backend/lora check, compile check) into a single boolean expression instead of nested ifs, replace the scope-device if/else with a ternary, and check dtype eligibility before any device introspection/moves (dtype is unaffected by device placement, so a mismatch now skips that work entirely instead of doing it first and discarding the result). Move each test file's torch-only fast-path tests into a dedicated skipif-decorated class with one module-level conditional `import torch`, replacing the repeated per-test `if backend.backend() != "torch": self.skipTest(...)` + local `import torch` pattern.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces PyTorch-specific fast-paths for the Dense and EinsumDense layers when running on the PyTorch backend, bypassing Keras's standard ops dispatch to call torch.nn.functional.linear and torch.einsum directly. This optimization is designed for eager execution without LoRA or active torch compilation, and includes handling for device alignment and active device scopes. Comprehensive unit tests are added to verify numeric equivalence, device handling, and dtype promotion rules. The review feedback correctly points out that subclasses overriding the kernel property would have their custom implementations bypassed by the fast-path, and suggests restricting these optimizations strictly to direct instances of the base Dense and EinsumDense classes.
…active The prior flattening moved the kernel-device check out from under the scope-override branch, so it ran an always-false self-comparison (kernel.device != kernel.device) on every call whenever no device_scope was active - the common case. Only compute and compare against an overridden target device when a scope is actually set; kernel is on kernel.device by definition otherwise, so there's nothing to check.
Both fast paths read self._kernel.value directly, bypassing the kernel property entirely. A subclass overriding kernel (a supported extension point per the styleguide - e.g. weight normalization, custom sharing) had that override silently skipped whenever the fast path engaged. Restrict both to type(self) is Dense / type(self) is EinsumDense so any subclass always takes the standard ops path, which correctly calls self.kernel.
A subclass with real logic (overriding the kernel property) is out of place nested inside a single test method - this codebase's convention (see rnn_test.py's OneStateRNNCell) reserves that for trivial pass-only subclasses and puts anything with actual behavior at module level.
3 participants
Closes #23276
Summary
On the torch backend,
Dense.callandEinsumDense.callnow dispatch directly totorch.nn.functional.linear/torch.einsumfor the plain, unquantized, non-LoRA eager forward pass, skipping Keras's genericops.matmul/ops.einsumoverhead.Why it's safe
The fast path is gated to only the case it's proven for:
quantized_call()runs beforecall().enable_loracan flip it on afterbuild().KerasTensorinputs never reachcall()at all (routed tocompute_output_specinstead), so no symbolic-input handling is needed here — pinned bytest_torch_fast_path_symbolic_input.torch.compiletracing is explicitly skipped (torch.compiler.is_compiling()), since meta-device weights makeF.linear/torch.einsumraise under fake tensors.Benchmarks (GPU, 3-way: master / this PR alone / this PR + parents)
Equivalence vs master: forward max|Δ| = 0.0 (CNN), 1.67e-06 (LLM, float32 recompute-order noise from the einsum reformulation — generate-32 token ids identical, the hard-fail criterion).
Part of the torch-backend perf series: #23182 → #23183 → #23184 → #23185 → #23186 → #23187 → #23188 → #23189 → #23190.
torch.compile impact
This PR is compile-neutral: 0 graph-break delta at its stack position under
torch.compile(deterministic dynamo decomposition). TheF.linear/torch.einsumfast path is eager-only and is explicitly skipped undertorch.compiler.is_compiling(), so the compiled graph is unaffected either way. The eager-mode speedup above is real, but the compiled-graph break-count reduction in this series comes later, from #23186 alone.Tests
dense_test.py/einsum_dense_test.py: numeric parity vs. the standard path, LoRA bypass, symbolic-input bypass,torch.compilebypass, device-scope handling, dtype-mismatch fallthrough, andmixed_float16equivalence.Contributor Agreement