Add Autoconfig and Coordinated_Optimizer implementations for Tensor Parallel Autosharding

[buildwithsuhana]

This PR introduces support for tensor parallelism autosharding in Keras, enabling users to shard large model layers across multiple devices. This is a crucial feature for training models that are too large to fit into the memory of a single accelerator.

The implementation is centered around two new components:

autoconfig.py: This module contains the logic to analyze a Keras model, identify sharding candidates (e.g., Dense, EinsumDense layers), and generate a sharding plan.

coordinated_optimizer.py: This is an optimizer wrapper that consumes the sharding plan. During training, it intercepts gradients for sharded variables and performs a collective AllReduce to ensure weight updates are correctly synchronized across all devices.

[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 significantly enhances Keras's distributed training capabilities by introducing automatic tensor parallelism. It provides the necessary infrastructure to intelligently shard large model layers across multiple devices, enabling the training of models that would otherwise exceed the memory capacity of a single accelerator. The changes include a system for automatically determining sharding strategies for model layers and a specialized optimizer that coordinates state and gradient updates across the sharded components.

Highlights

• Tensor Parallel Autosharding Configuration: Introduced autoconfig.py which provides automatic analysis of Keras models to generate a sharding plan for tensor parallelism. This module includes heuristics to classify Dense layers (up-projection, down-projection) and applies appropriate column-wise or row-wise sharding rules to Dense, EinsumDense, and Embedding layers.
• Coordinated Optimizer for Distributed Training: Added coordinated_optimizer.py which implements a TensorParallelOptimizer wrapper. This optimizer manages sharded optimizer states across multiple devices and synchronizes gradients using collective operations (like AllReduce) based on the tensor parallelism configuration, ensuring correct weight updates in a distributed environment.
• Comprehensive Unit Testing: New unit tests (autoconfig_test.py and coordinated_optimizer_test.py) have been added to validate the functionality of both the autoconfiguration logic and the coordinated optimizer, covering various scenarios including nested models, different layer types, and optimizer state management.

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 introduces significant new functionality for tensor parallelism autosharding. The implementation is split across autoconfig.py for sharding plan generation and coordinated_optimizer.py for synchronized training. My review has identified a few issues ranging from critical to medium severity. Notably, there's a critical bug in TensorParallelOptimizer that prevents the use of learning rate schedules. I've also found some high-severity maintainability concerns in CoordinatedOptimizer due to its reliance on parsing internal variable names, which is fragile. Additionally, there are potential correctness issues in autoconfig.py and some violations of the Keras API design style guide. Addressing these points will improve the robustness and maintainability of this new feature.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 45.20548% with 320 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.22%. Comparing base (5ae5503) to head (5c24951).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
...tribution/tensor_parallel/coordinated_optimizer.py	56.88%	81 Missing and 13 partials ⚠️
...src/distribution/tensor_parallel/communications.py	29.89%	68 Missing ⚠️
keras/src/backend/jax/distributed_backend.py	24.69%	60 Missing and 1 partial ⚠️
...ras/src/distribution/tensor_parallel/autoconfig.py	56.79%	25 Missing and 10 partials ⚠️
keras/src/distribution/tensor_parallel/config.py	35.00%	26 Missing ⚠️
...distribution/tensor_parallel/state_action_keras.py	44.68%	25 Missing and 1 partial ⚠️
keras/src/backend/distributed/backend_resolver.py	33.33%	9 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21703      +/-   ##
==========================================
- Coverage   82.59%   82.22%   -0.38%     
==========================================
  Files         572      580       +8     
  Lines       58327    58906     +579     
  Branches     9131     9232     +101     
==========================================
+ Hits        48177    48437     +260     
- Misses       7818     8112     +294     
- Partials     2332     2357      +25     

Flag	Coverage Δ
keras	82.03% <45.20%> (-0.37%)	⬇️
keras-jax	63.13% <45.20%> (-0.18%)	⬇️
keras-numpy	57.53% <45.20%> (-0.12%)	⬇️
keras-openvino	34.32% <30.82%> (+<0.01%)	⬆️
keras-tensorflow	63.86% <45.20%> (-0.19%)	⬇️
keras-torch	63.47% <45.20%> (-0.16%)	⬇️

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.