Skip to content

refactor(tree_attn): decouple FSDP and Megatron implementations #844

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nuzant merged 1 commit into from
Jan 22, 2026
Merged

refactor(tree_attn): decouple FSDP and Megatron implementations #844

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions areal/models/tree_attn/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
"""Shared constants for tree attention functionality."""

import os

BLOCK_SIZE = int(os.environ.get("AREAL_FLEX_ATTENTION_BLOCK_SIZE", "128"))
318 changes: 25 additions & 293 deletions areal/models/tree_attn/module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,296 +1,28 @@
import os
from contextlib import contextmanager

import torch
from mbridge.core import LLMBridge
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.transformer import TransformerConfig
from megatron.core.transformer.attention import SelfAttention
from megatron.core.transformer.enums import AttnMaskType
from megatron.core.transformer.transformer_block import (
TransformerBlockSubmodules,
)
from megatron.core.transformer.transformer_layer import (
TransformerLayer,
TransformerLayerSubmodules,
)
from torch.nn.attention.flex_attention import (
BlockMask,
create_block_mask,
flex_attention,
)

from areal.utils import logging

logger = logging.getLogger(__name__)

_TORCH_COMPILE_OPTIONS = {
"epilogue_fusion": True,
"max_autotune": True,
"shape_padding": True,
"trace.enabled": False,
"triton.cudagraphs": False,
}
_FLEX_DYNAMIC = not (os.environ.get("AREAL_DISABLE_FLEX_ATTENTION_DYNAMIC", "0") == "1")
logger.info(
"Compiled torch flex attention. Options: %s, dynamic: %s",
str(_TORCH_COMPILE_OPTIONS),
str(_FLEX_DYNAMIC),
from areal.models.tree_attn.constants import BLOCK_SIZE
from areal.models.tree_attn.module_fsdp import (
create_block_mask_from_dense,
patch_fsdp_for_tree_training,
restore_patch_fsdp_for_tree_training,
)
_flex_attention = torch.compile(
flex_attention,
dynamic=_FLEX_DYNAMIC,
options=_TORCH_COMPILE_OPTIONS,
)
BLOCK_SIZE = int(os.environ.get("AREAL_FLEX_ATTENTION_BLOCK_SIZE", "128"))
logger.info("Using block mask in flex attention, block size: %d", BLOCK_SIZE)


def create_block_mask_from_dense(
attention_mask: torch.Tensor,
seq_len: int,
device: torch.device,
) -> BlockMask:
"""Create a flex attention block mask from a dense attention mask.

This function should be called early (during data preparation) to allow
the dense mask to be released and save memory.

Parameters
----------
attention_mask : torch.Tensor
Dense attention mask of shape (seq_len, seq_len).
seq_len : int
Sequence length.
device : torch.device
Device to create the block mask on.

Returns
-------
BlockMask
The created block mask for use with flex_attention.
"""

def arbitrary_mask(
batch: torch.Tensor,
head: torch.Tensor,
q_idx: torch.Tensor,
k_idx: torch.Tensor,
):
return attention_mask[q_idx, k_idx]

block_mask = create_block_mask(
arbitrary_mask,
B=1,
H=1,
Q_LEN=seq_len,
KV_LEN=seq_len,
BLOCK_SIZE=BLOCK_SIZE,
device=device,
_compile=False,
)
return block_mask


class PytorchFlexAttention(torch.nn.Module):
"""Pytorch flex attention implementation that supports arbitrary attention mask type."""

def __init__(
self,
config: TransformerConfig,
layer_number: int,
attn_mask_type: AttnMaskType,
attention_type: str,
attention_dropout: float | None = None,
softmax_scale: float | None = None,
**kwargs,
):
super().__init__()
self.config = config
self.layer_number = layer_number
self.attn_mask_type = attn_mask_type
self.attention_type = attention_type
self.attention_dropout = attention_dropout
self.softmax_scale = softmax_scale
logger.info("Using PytorchFlexAttention for tree training attention")

# PytorchFlexAttention does not support context parallel
if config.context_parallel_size != 1:
raise ValueError(
"PytorchFlexAttention does not support context parallelism."
)

if attention_type != "self":
raise ValueError("PytorchFlexAttention only supports self-attention.")

def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: BlockMask,
attn_mask_type: AttnMaskType,
attention_bias: torch.Tensor = None,
packed_seq_params: PackedSeqParams = None,
):
# query: [S, B, H, D] in which B should be 1 in current tree training implementation
# key: [S, B, H, D]
# value: [S, B, H, D]
# attention_mask: BlockMask (pre-created)
# attention_mask_type: arbitrary

if attention_bias is not None:
raise NotImplementedError(
"PytorchFlexAttention does not support attention_bias yet."
)
if packed_seq_params is not None:
raise NotImplementedError(
"PytorchFlexAttention does not support packed sequences yet."
)
if not isinstance(attention_mask, BlockMask):
raise ValueError(
"PytorchFlexAttention requires a pre-created BlockMask. "
"Use create_block_mask_from_dense() during data preparation."
)

# query, key, value shape: [S, B, H, D] -> [B, H, S, D]
query = query.permute(1, 2, 0, 3)
key = key.permute(1, 2, 0, 3)
value = value.permute(1, 2, 0, 3)
enable_gqa = query.shape[1] != key.shape[1]

output = _flex_attention(
query,
key,
value,
block_mask=attention_mask,
score_mod=None,
scale=self.softmax_scale,
enable_gqa=enable_gqa,
)

# output shape: [B, H, S, D] -> [S, B, H, D] -> [S, B, H*D]
output = (
output.permute(2, 0, 1, 3)
.contiguous()
.view(output.shape[2], output.shape[0], -1)
)
return output


def _tree_attn_fwd_func(
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None = None,
softmax_scale: float | None = None,
*args,
**kwargs,
):
# Require pre-created block_mask
block_mask = kwargs.get("block_mask", None)
if block_mask is None or not isinstance(block_mask, BlockMask):
raise ValueError(
"_tree_attn_fwd_func requires a pre-created BlockMask in kwargs['block_mask']. "
"Use create_block_mask_from_dense() during data preparation."
)

# [B, S, H, D] -> [B, H, S, D]
query = query.permute(0, 2, 1, 3).contiguous()
key = key.permute(0, 2, 1, 3).contiguous()
value = value.permute(0, 2, 1, 3).contiguous()

enable_gqa = query.shape[1] != key.shape[1]

output = _flex_attention(
query,
key,
value,
block_mask=block_mask,
score_mod=None,
scale=softmax_scale,
enable_gqa=enable_gqa,
)
# [B, H, S, D] -> [B, S, H, D]
output = output.permute(0, 2, 1, 3).contiguous()
return output


@contextmanager
def patch_bridge_for_tree_training(enable: bool = True):
"""Context manager to patch LLMBridge for tree training with arbitrary attention mask.

Parameters
----------
enable : bool, default=True
If True, apply the patch. If False, the context manager is a no-op.

Yields
------
None

Examples
--------
>>> with patch_bridge_for_tree_training(enable=True):
... # LLMBridge is patched here
... model = create_model()
... # Patch is reverted after exiting the context
"""
if not enable:
yield
return

# Store original method
original_layer_spec_getter = LLMBridge._get_transformer_layer_spec

def _patched_getter(self, vp_stage: int | None = None):
spec: TransformerBlockSubmodules = original_layer_spec_getter(self, vp_stage)
for layer_spec in spec.layer_specs:
if layer_spec.module is not TransformerLayer:
logger.info(f"Skipping patch module: {layer_spec.module}")
continue
submodules: TransformerLayerSubmodules = layer_spec.submodules
self_attn_spec = submodules.self_attention
if self_attn_spec.module is not SelfAttention:
logger.info(f"Skipping patch module: {self_attn_spec.module}")
continue
self_attn_spec.params["attn_mask_type"] = AttnMaskType.arbitrary
self_attn_spec.submodules.core_attention = PytorchFlexAttention
return spec

# Apply patch
LLMBridge._get_transformer_layer_spec = _patched_getter
try:
yield
finally:
# Revert patch
LLMBridge._get_transformer_layer_spec = original_layer_spec_getter


ORIGINAL_FLASH_ATTENTION_FORWARD = None


def patch_fsdp_for_tree_training(enable: bool = True):
if not enable:
return

from transformers.integrations import flash_attention

global ORIGINAL_FLASH_ATTENTION_FORWARD
ORIGINAL_FLASH_ATTENTION_FORWARD = flash_attention._flash_attention_forward
flash_attention._flash_attention_forward = _tree_attn_fwd_func
logger.info(
"Patched transformers.integrations.flash_attention._flash_attention_forward "
"with tree implementation."
)


# Only used in testing to restore original function
def restore_patch_fsdp_for_tree_training():
from transformers.integrations import flash_attention

flash_attention._flash_attention_forward = ORIGINAL_FLASH_ATTENTION_FORWARD
logger.info(
"Restored transformers.integrations.flash_attention._flash_attention_forward "
"to original implementation."
# Conditionally import Megatron functionality
try:
from areal.models.tree_attn.module_megatron import (
PytorchFlexAttention,
patch_bridge_for_tree_training,
)
except ImportError:
PytorchFlexAttention = None
patch_bridge_for_tree_training = None

__all__ = [
# Shared constants
"BLOCK_SIZE",
# FSDP/common exports
"create_block_mask_from_dense",
"patch_fsdp_for_tree_training",
"restore_patch_fsdp_for_tree_training",
# Megatron exports (may be None if Megatron not installed)
"PytorchFlexAttention",
"patch_bridge_for_tree_training",
]
Loading