Skip to content
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

feat(2D): support selective-checkpoint-offload for 2D attention #396

Open
wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

feat(2D): support selective-checkpoint-offload for 2D attention #396

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dc50325
support 2D offload
yingtongxiong Dec 31, 2024
ef5e8b6
fix lint
yingtongxiong Dec 31, 2024
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
14 changes: 7 additions & 7 deletions internlm/model/ops/ring_flash_attn/zigzag_ring_flash_attn_with_sliding_window.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,10 @@
from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward

from internlm.core.context.parallel_context import global_context as gpc
from internlm.core.parallel.comm import get_offload_manager

from .utils import RingComm, update_out_and_lse

fa_output_mapping = {}


def create_buffer(tensor):
buffer_shape = list(tensor.shape)
Expand Down Expand Up @@ -443,9 +442,10 @@ def forward(
k = k.contiguous()
v = v.contiguous()

if gpc.is_forward is False and gpc.config.selective_checkpoint:
assert layer_idx in fa_output_mapping
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

第11行fa_output_mapping的初始化可以去掉了

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

哦对 好的

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

out, softmax_lse = fa_output_mapping.pop(layer_idx)
_ckpt_block_num = int(gpc.config.model.checkpoint * gpc.config.isp_num_layers)

if gpc.is_forward is False and gpc.config.selective_checkpoint and layer_idx < _ckpt_block_num:
out, softmax_lse = get_offload_manager().get_fa_output_with_layer(layer_idx)
else:
out, softmax_lse = zigzag_double_ring_flash_attn_forward(
context_group,
Expand All @@ -460,8 +460,8 @@ def forward(
)

# store attn forward output to avoid re-computation of attn when activation checkpoint is enabled
if gpc.is_forward and gpc.config.selective_checkpoint:
fa_output_mapping[layer_idx] = (out, softmax_lse)
if gpc.is_forward and gpc.config.selective_checkpoint and layer_idx < _ckpt_block_num:
get_offload_manager().insert_fa_output_with_layer(layer_idx=layer_idx, output=(out, softmax_lse))

# this should be out_padded
ctx.save_for_backward(q, k, v, out, softmax_lse)
Expand Down
Loading