From 9ddd6ab0e0d61cdadb5da336a475d9bfda183512 Mon Sep 17 00:00:00 2001 From: "meizhiyu.mzy" Date: Sun, 1 Feb 2026 17:25:47 +0800 Subject: [PATCH 1/3] lazy attn mask for tree training --- areal/engine/fsdp_engine.py | 19 +++++++++- areal/engine/megatron_engine.py | 21 ++++++++++- areal/models/tree_attn/module.py | 2 ++ areal/models/tree_attn/tree.py | 61 +++++++++++++++++++++++++++----- areal/utils/data.py | 7 ++++ 5 files changed, 99 insertions(+), 11 deletions(-) diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index b03e353d10..330cb9aafa 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -70,7 +70,11 @@ gather_packed_tree_logprobs_entropy, merge_packed_tree_results, ) -from areal.models.tree_attn.module import BLOCK_SIZE, patch_fsdp_for_tree_training +from areal.models.tree_attn.module import ( + BLOCK_SIZE, + build_block_mask_from_trie, + patch_fsdp_for_tree_training, +) from areal.models.tree_attn.tree import TrieNode, build_packed_tree_batch from areal.platforms import current_platform from areal.utils import ( @@ -541,10 +545,23 @@ def forward_backward_batch( for mb_item in mb_list: inputs, ctx = self._prepare_mb_inputs(mb_item) + # Lazily create block mask for tree training just before forward + if self.enable_tree_training and ctx.trie_node is not None: + padded_size = mb_item.padded_to_length + if padded_size is not None: + block_mask = build_block_mask_from_trie( + ctx.trie_node, padded_size, self.device + ) + inputs["block_mask"] = block_mask + with trace_scope("fsdp_engine.forward"): outputs = self.model(**inputs) logits = outputs.logits.squeeze(0) + # Release block mask memory after forward pass + if self.enable_tree_training and "block_mask" in inputs: + del inputs["block_mask"] + ctx_dict = ctx.to_dict() loss = process_output_fn(logits, ctx_dict) diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index e4ab4312e2..acdc574bf0 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -57,7 +57,11 @@ gather_packed_tree_logprobs_entropy, merge_packed_tree_results, ) -from areal.models.tree_attn.module import BLOCK_SIZE, patch_bridge_for_tree_training +from areal.models.tree_attn.module import ( + BLOCK_SIZE, + build_block_mask_from_trie, + patch_bridge_for_tree_training, +) from areal.models.tree_attn.tree import build_packed_tree_batch from areal.platforms import current_platform from areal.utils import logging, name_resolve, names, perf_tracer, stats_tracker @@ -565,8 +569,23 @@ def forward_step(batch_iter, model): mb_input: MicroBatchItem = next(batch_iter) cu_seqlens = mb_input.padded_mb.get("cu_seqlens", None) + + # Lazily create block mask for tree training just before forward + if self.enable_tree_training: + trie_node = mb_input.padded_mb.get("trie_node", None) + padded_size = mb_input.padded_to_length + if trie_node is not None and padded_size is not None: + block_mask = build_block_mask_from_trie( + trie_node, padded_size, self.device + ) + mb_input.padded_mb["block_mask"] = block_mask + output = packed_context_parallel_forward(model, mb_input.padded_mb) + # Release block mask memory after forward pass + if self.enable_tree_training and "block_mask" in mb_input.padded_mb: + del mb_input.padded_mb["block_mask"] + def _process_output(input_, output_): loss = process_output_fn(output_, input_) if loss is None: diff --git a/areal/models/tree_attn/module.py b/areal/models/tree_attn/module.py index 424994fadc..21f5df3793 100644 --- a/areal/models/tree_attn/module.py +++ b/areal/models/tree_attn/module.py @@ -4,6 +4,7 @@ patch_fsdp_for_tree_training, restore_patch_fsdp_for_tree_training, ) +from areal.models.tree_attn.tree import build_block_mask_from_trie # Conditionally import Megatron functionality try: @@ -22,6 +23,7 @@ "create_block_mask_from_dense", "patch_fsdp_for_tree_training", "restore_patch_fsdp_for_tree_training", + "build_block_mask_from_trie", # Megatron exports (may be None if Megatron not installed) "PytorchFlexAttention", "patch_bridge_for_tree_training", diff --git a/areal/models/tree_attn/tree.py b/areal/models/tree_attn/tree.py index 54f36090ed..9b819f75c5 100644 --- a/areal/models/tree_attn/tree.py +++ b/areal/models/tree_attn/tree.py @@ -8,7 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.distributed as dist @@ -20,6 +20,9 @@ from areal.utils.data import MicroBatchList from areal.utils.perf_tracer import trace_perf, trace_scope +if TYPE_CHECKING: + from torch.nn.attention.flex_attention import BlockMask + logger = logging.getLogger(__name__) @@ -370,12 +373,6 @@ def build_packed_tree_batch( mask_template.device, ) - # Create block mask from dense mask - with trace_scope("tree_attn.create_block_mask"): - block_mask = create_block_mask_from_dense( - attention_mask, padded_size, mask_template.device - ) - # Compute position_ids (needs dense attention_mask) with trace_scope("tree_attn.get_position_ids"): position_ids = get_packed_tree_position_ids( @@ -384,6 +381,7 @@ def build_packed_tree_batch( ) # Release dense attention mask memory after position_ids are computed + # Block mask will be lazily created in forward_backward_batch del attention_mask # Pack extra data @@ -396,10 +394,9 @@ def build_packed_tree_batch( non_packable_keys, ) - # Build micro-batch dict with block_mask + # Build micro-batch dict without block_mask (will be created lazily in forward) mb = { "input_ids": input_ids, - "block_mask": block_mask, "position_ids": position_ids, "trie_node": trie, **extra_data, @@ -611,3 +608,49 @@ def get_packed_tree_position_ids( position_ids = torch.clamp_min(ancestor_counts - 1, 0) return position_ids.unsqueeze(0) + + +@trace_perf("tree_attn.build_block_mask_from_trie") +def build_block_mask_from_trie( + trie: TrieNode, + padded_size: int, + device: torch.device, +) -> BlockMask: + """Lazily build a block mask from a trie node. + + This function builds the dense attention mask from the trie structure and + converts it to a block mask for use with flex attention. It should be called + just before the forward pass to minimize memory usage. + + Parameters + ---------- + trie : TrieNode + The root trie node containing the tree structure. + padded_size : int + The padded sequence length. + device : torch.device + Device to create the block mask on. + + Returns + ------- + BlockMask + The created block mask for use with flex_attention. + """ + # Handle dummy trie (empty tree for DP synchronization) + if not trie.all_sequence_ids: + # Create a minimal valid block mask for empty trees + dummy_mask = torch.zeros( + (padded_size, padded_size), dtype=torch.bool, device=device + ) + return create_block_mask_from_dense(dummy_mask, padded_size, device) + + with trace_scope("tree_attn.build_attention_mask"): + attention_mask = _build_attention_mask(trie, padded_size, device) + + with trace_scope("tree_attn.create_block_mask"): + block_mask = create_block_mask_from_dense(attention_mask, padded_size, device) + + # Release dense attention mask memory + del attention_mask + + return block_mask diff --git a/areal/utils/data.py b/areal/utils/data.py index 68b5eccc17..73eb28d1e5 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -372,12 +372,14 @@ class MicroBatchItem(NamedTuple): padded_mb: Padded micro-batch dict (for model forward) padding_length: Batch-level padding added to this micro-batch old_cu_seqlens: Original cu_seqlens before sequence alignment (or None) + padded_to_length: The padded sequence length for this micro-batch (or None) """ orig_mb: dict[str, Any] padded_mb: dict[str, Any] padding_length: int old_cu_seqlens: torch.Tensor | None + padded_to_length: int | None = None @dataclass @@ -421,6 +423,7 @@ def __iter__(self) -> Iterator[MicroBatchItem]: - padded_mb: Padded micro-batch dict (for model forward) - padding_length: Batch-level padding added to this micro-batch - old_cu_seqlens: Original cu_seqlens before sequence alignment (or None) + - padded_to_length: The padded sequence length for this micro-batch (or None) """ if self.padded_mbs is None: raise ValueError("padded_mbs is None. Call pad_mb_list first.") @@ -428,11 +431,15 @@ def __iter__(self) -> Iterator[MicroBatchItem]: old_cu_seqlens = ( self.old_cu_seqlens_list[i] if self.old_cu_seqlens_list else None ) + padded_to_length = ( + self.padded_to_lengths[i] if self.padded_to_lengths else None + ) yield MicroBatchItem( orig_mb=self.mbs[i], padded_mb=self.padded_mbs[i], padding_length=self.padding_lengths[i], old_cu_seqlens=old_cu_seqlens, + padded_to_length=padded_to_length, ) def to(self, *args, **kwargs): From 38144de6e61720073696aa5377c67413cbb8b429 Mon Sep 17 00:00:00 2001 From: nuzant Date: Sun, 1 Feb 2026 18:29:12 +0800 Subject: [PATCH 2/3] Raise error for padded_size. --- areal/engine/fsdp_engine.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 330cb9aafa..2adc5f768f 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -548,11 +548,14 @@ def forward_backward_batch( # Lazily create block mask for tree training just before forward if self.enable_tree_training and ctx.trie_node is not None: padded_size = mb_item.padded_to_length - if padded_size is not None: - block_mask = build_block_mask_from_trie( - ctx.trie_node, padded_size, self.device + if padded_size is None: + raise ValueError( + "padded_size must be set for tree training with FSDP." ) - inputs["block_mask"] = block_mask + block_mask = build_block_mask_from_trie( + ctx.trie_node, padded_size, self.device + ) + inputs["block_mask"] = block_mask with trace_scope("fsdp_engine.forward"): outputs = self.model(**inputs) From aef9c009eb817e1fac85e323fdabf8143bfd7315 Mon Sep 17 00:00:00 2001 From: nuzant Date: Mon, 2 Feb 2026 13:58:10 +0800 Subject: [PATCH 3/3] fix import --- areal/models/tree_attn/tree.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/areal/models/tree_attn/tree.py b/areal/models/tree_attn/tree.py index 9b819f75c5..e9499e23a1 100644 --- a/areal/models/tree_attn/tree.py +++ b/areal/models/tree_attn/tree.py @@ -8,10 +8,11 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any +from typing import Any import torch import torch.distributed as dist +from torch.nn.attention.flex_attention import BlockMask from areal.api.cli_args import MicroBatchSpec from areal.models.tree_attn.constants import BLOCK_SIZE @@ -20,9 +21,6 @@ from areal.utils.data import MicroBatchList from areal.utils.perf_tracer import trace_perf, trace_scope -if TYPE_CHECKING: - from torch.nn.attention.flex_attention import BlockMask - logger = logging.getLogger(__name__)