Refactor: Decouple Core Transformer Blocks

[parambole]

TL;DR

• What: This PR refactors the core DecoderLayer and its related components out of layers/models.py and into a new, foundational file: layers/blocks.py.

• Why: To improve the overall code architecture and break potential circular dependencies. This is a necessary prerequisite for adding new, complex modules that also need access to these core building blocks.

• How: By creating layers/blocks.py to house the fundamental Decoder and DecoderLayer classes. Higher-level files like models.py and other future modules now import these components from a single location.

Detailed Description

This pull request introduces a structural refactoring to improve modularity and maintainability.

The primary change is the creation of MaxText/layers/blocks.py, which now serves as the source for fundamental building blocks of the Transformer architecture, such as:

• DecoderLayer

• Decoder

Previously, these classes were located in layers/models.py, which created tight coupling. As we add more features this tight coupling would lead to circular import dependencies.

By decoupling these core components, we establish a clearer hierarchy in the codebase, where high-level modules can depend on these "building blocks" without depending on each other.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[bvandermoon]

Thanks for refactoring @parambole. Can you explain the circular dependency you mentioned in the description?

[bvandermoon]

Is there a more descriptive file name that can be used here instead of 'blockes.py'? Maybe naming this file decoders.py? Would need to create a separate encoders.py file for the VisionEncoder

[RissyRan]

I agree to name this as decoders.py, and move class VisionEncoder to encoders.py.

[RissyRan]

Could you select one model to run a quick perf test with same configs? main branch vs. your branch. I see functional tests are passing, just ensure some breakdown does not hurt performance.