refactor(tree_attn): decouple FSDP and Megatron implementations

[rchardx]

Description

Split module.py into separate files to allow FSDP engine to import tree attention functionality without requiring Megatron dependencies:

• constants.py: shared BLOCK_SIZE constant (lightweight, no deps)
• module_fsdp.py: FSDP-only functionality (PyTorch only)
• module_megatron.py: Megatron-specific functionality (mbridge, megatron.core)
• module.py: backward-compatible re-exports with conditional Megatron imports

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not
  work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with
  jb build docs
• No critical issues raised by AI reviewers (/gemini review)

[gemini-code-assist]

Summary of Changes

Hello @rchardx, 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 refactors the tree_attn module by separating its components based on their dependencies. The primary goal is to enhance modularity and reduce coupling, enabling the use of tree attention with FSDP without imposing a hard dependency on Megatron. This change improves the flexibility and maintainability of the codebase by clearly delineating responsibilities and allowing for more targeted imports.

Highlights

• Decoupling FSDP and Megatron Implementations: The core tree attention functionality has been refactored to separate FSDP (Fully Sharded Data Parallel) and Megatron-specific code. This allows FSDP-only environments to utilize tree attention without requiring Megatron dependencies.
• New Module Structure: The original module.py has been split into three new files: constants.py for shared constants, module_fsdp.py for FSDP-specific logic, and module_megatron.py for Megatron-specific logic. The original module.py now serves as a backward-compatible entry point, conditionally importing Megatron components.
• Conditional Megatron Imports: Megatron-specific functionalities are now imported conditionally in module.py using a try-except ImportError block. This ensures that if Megatron is not installed, the FSDP-related tree attention features can still be used without errors.
• Updated Import Paths: Files that previously imported directly from module.py, such as tree.py, have been updated to reflect the new modular structure, importing BLOCK_SIZE from constants.py and FSDP-related functions from module_fsdp.py.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 effectively refactors the tree_attn module to decouple the FSDP and Megatron implementations. The changes improve modularity by splitting the code into separate files for constants, FSDP-specific logic, and Megatron-specific logic, while maintaining backward compatibility. My review includes a suggestion to improve the robustness of the monkey-patching mechanism in the FSDP module to prevent potential state management issues.

[Copilot]

Pull request overview

This PR refactors the tree attention module to decouple FSDP and Megatron implementations, allowing the FSDP engine to use tree attention functionality without requiring Megatron dependencies. The refactoring splits the monolithic module.py into separate focused files while maintaining backward compatibility.

Changes:

• Split module.py into constants.py (shared), module_fsdp.py (FSDP-only), and module_megatron.py (Megatron-specific)
• Updated module.py to provide backward-compatible re-exports with conditional Megatron imports
• Updated tree.py to import directly from module_fsdp, eliminating unnecessary Megatron dependencies

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
areal/models/tree_attn/constants.py	New file defining shared BLOCK_SIZE constant with no external dependencies
areal/models/tree_attn/module_fsdp.py	New file with FSDP-specific functionality using PyTorch-only dependencies
areal/models/tree_attn/module_megatron.py	New file with Megatron-specific PytorchFlexAttention and bridge patching
areal/models/tree_attn/module.py	Refactored to provide backward-compatible re-exports with conditional Megatron imports
areal/models/tree_attn/tree.py	Updated imports to use module_fsdp directly, avoiding Megatron dependencies

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.