Skip to content

Split the graphs to run with flash_attention on 1x #75

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
3 commits merged into from
Mar 4, 2024
Merged

Split the graphs to run with flash_attention on 1x #75

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ab65e67
Split the graphs to run with flash_attention on 1x
kalyanjkk Feb 26, 2024
791a644
Merge branch 'HabanaAI:habana-main' into decoder_mark_step
kalyanjk Mar 1, 2024
d4d1b9c
Added lazy_mode check and removed additional htcore import
kalyanjkk Mar 1, 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
3 changes: 3 additions & 0 deletions optimum/habana/transformers/models/llama/modeling_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -692,6 +692,9 @@ def forward(
htcore.mark_step()

for layer_idx, decoder_layer in enumerate(self.layers):
if lazy_mode and torch.distributed.is_initialized() == False:
htcore.mark_step()
Copy link
Copy Markdown
Collaborator

@libinta libinta Feb 26, 2024

Choose a reason for hiding this comment

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

@kalyanjk what's the impact for input/output not introduced oom? should we add an argument in text-generation from cmd line?

Copy link
Copy Markdown

@mandy-li mandy-li Feb 26, 2024

Choose a reason for hiding this comment

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

@kalyanjk ,why only mark_step() for 1x?

Copy link
Copy Markdown
Author

@kalyanjk kalyanjk Feb 27, 2024

Choose a reason for hiding this comment

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

For 8x mark_step will be introduced through a collective call.

Copy link
Copy Markdown
Author

@kalyanjk kalyanjk Feb 27, 2024

Choose a reason for hiding this comment

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

@kalyanjk what's the impact for input/output not introduced oom? should we add an argument in text-generation from cmd line?

The issue is not with oom. The real issue is recipe size being too large and also compilation time is too high.

Copy link
Copy Markdown

@puneeshkhanna puneeshkhanna Mar 5, 2024

Choose a reason for hiding this comment

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

Please update as below
if lazy_mode and (torch.distributed.is_initialized() is False or torch.distributed.get_world_size() == 1):


if output_hidden_states:
all_hidden_states += (hidden_states,)

Expand Down