Core Data Structures & Communication Primitives for Tensor Parallel for Keras

[buildwithsuhana]

This Pull Request introduces the foundational components for a new, backend-agnostic auto-sharding system in Keras, specifically designed for tensor parallelism. It establishes the core data structures and the JAX-specific implementation of communication primitives.

1. Core Backend-Agnostic Abstractions
   The most significant part of this PR is the creation of a generic, backend-agnostic system for defining sharding plans. This logic resides in keras/src/distribution/tensor_parallel/tensor_layout.py.

LayoutAction: An abstract base class that serves as a blueprint for any sharding operation. It defines a contract that requires a forward method (call) to apply the sharding and a reverse method (undo) to reconstruct the original tensor.

Split: The first concrete implementation of LayoutAction. This class contains the logic to manually slice a tensor along a specific dimension. It is robustly designed to handle cases where the tensor size is not perfectly divisible by the number of devices.

Note on Design: This class performs explicit, manual slicing rather than using a declarative approach. This imperative design is intentional and crucial for ensuring multi-backend compatibility. While a backend like JAX could infer this slicing automatically, other frameworks like PyTorch require these explicit instructions to distribute data correctly. This universal approach allows the core Keras logic to remain backend-agnostic.

LayoutMap: A simple data class that acts as a container for the complete sharding plan. It is designed to hold the LayoutAction rules for both the model's weights (state_rules) and its intermediate outputs (output_rules).

2. JAX-Specific Backend Implementation
   This PR provides the first backend-specific implementation of the required distributed communication primitives, targeting JAX. This code is located in keras/src/backend/jax/distributed_backend.py.

get_device_info() and is_multi_device_capable(): Utility functions to query the JAX environment for available devices.

get_communication_ops(): This is the key function that provides the necessary tools for cross-device communication. It returns a dictionary containing backend-agnostic names (all_reduce, all_gather) mapped to their efficient, native JAX implementations (jax.lax.psum and jax.lax.all_gather). This acts as a translation layer, allowing the core Keras logic to use these operations without needing to know it's running on JAX.

Design Document: Autosharding for Keras

The full code of Tensor parallel for Keras has been devided into 4 PRs, this is the first PR for the same.

[gemini-code-assist]

Summary of Changes

Hello @buildwithsuhana, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request lays the fundamental groundwork for introducing backend-agnostic auto-sharding and Tensor Parallelism into Keras 3.0. It establishes a modular and extensible architecture by defining core data structures, abstracting distributed backend functionalities, and providing high-level communication primitives. This initial set of changes is crucial for enabling future capabilities that will allow users to train very large models across multiple devices with significantly simplified code.

Highlights

• Core Distributed Backend Abstraction: Introduced BaseDistributedBackend as an abstract interface for distributed operations and a get_distributed_backend factory function to provide a unified, backend-agnostic way to interact with JAX, TensorFlow, PyTorch, and NumPy distributed environments.
• High-Level Communication Primitives: Defined AllReduceKeras, AllGatherKeras, BroadcastKeras, and ScatterKeras classes, which serve as high-level wrappers for essential collective communication operations required for tensor parallelism.
• Tensor Sharding Actions: Implemented StateActionKeras as an abstract base class for defining how tensors are transformed for distribution. Concrete implementations like SplitKeras handle tensor sharding, while GatherKeras and SumKeras define how to reconstruct original tensors from their distributed parts.
• Sharding Plan Configuration: Introduced the ConfigKeras dataclass to store and manage model-wide sharding rules and output configurations, including a mechanism to dynamically create collective operations based on these rules.
• Tensor Parallel Communicator: Added TensorParallelCommunicator to orchestrate complex communication patterns for tensor parallelism, including specific methods for handling forward and backward passes in column-parallel and row-parallel operations, along with gradient slicing logic.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request lays a solid foundation for tensor parallelism in Keras by introducing backend-agnostic abstractions for distributed operations and core data structures for sharding. The overall design is well-structured, separating concerns between backend-specific implementations, communication primitives, and configuration. However, there are several areas that need attention, particularly regarding the correctness of some backend implementations (especially JAX), placeholder logic, API clarity, and code consistency. Addressing these points will strengthen the foundation and prevent issues in future development.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.08197% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.60%. Comparing base (5ae5503) to head (cd0049f).
⚠️ Report is 14 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/distributed_backend.py	88.23%	1 Missing and 1 partial ⚠️
.../src/distribution/tensor_parallel/tensor_layout.py	97.43%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##           master   #21697    +/-   ##
========================================
  Coverage   82.59%   82.60%            
========================================
  Files         572      574     +2     
  Lines       58327    58596   +269     
  Branches     9131     9165    +34     
========================================
+ Hits        48177    48403   +226     
- Misses       7818     7854    +36     
- Partials     2332     2339     +7     

Flag	Coverage Δ
keras	82.40% <95.08%> (+<0.01%)	⬆️
keras-jax	63.22% <88.52%> (-0.09%)	⬇️
keras-numpy	57.53% <31.14%> (-0.12%)	⬇️
keras-openvino	34.33% <31.14%> (+<0.01%)	⬆️
keras-tensorflow	63.91% <31.14%> (-0.14%)	⬇️
keras-torch	63.45% <31.14%> (-0.18%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[JyotinderSingh]

I've added a few initial comments and questions during my first look.

To make the review more manageable, I propose we split this change up. At almost 1,800 lines, the current change is quite difficult to review properly. What do you think about limiting this PR to just the JAX backend, and introducing the others in subsequent, smaller PRs?

[hertschuh]

Thank you for the PR!

Some high level comments:

• Out of context, it's really hard for me to understand why these abstractions are needed for Tensor Parallel.
  • Why do we need all these primitives?
  • Why do we need 3 layers of abstraction for the same concepts: the communications layer, the state_actions layer and the keras.distributed.get_communication_ops layer? Can we just have one?
• These abstraction look Torch-like and not JAX-like. On JAX you never have to manually split and do an all-gather, you simply shard. You never have to explicitly have to do a "collective sum". You just do a sum, and if the tensors are sharded, it will magically do all the needed collectives for you. So it's unclear to me why any of these are needed for JAX.
• I wouldn't export these symbols that you added to keras.distributed, I don't think they are needed. What we'll expose is the "Tensor Parallel" API.
• For the better or worse, we don't do type annotations in Keras. And unfortunately, mixing code with type annotations with code without type annotation doesn't work well. It's better to not have any type annotations at all.