Split the graphs to run with flash_attention on 1x #75
Conversation
| @@ -676,6 +677,9 @@ def forward( | |||
| next_decoder_cache = () if not use_new_cache else None | |||
|
|
|||
| for layer_idx, decoder_layer in enumerate(self.layers): | |||
| if torch.distributed.is_initialized() == False: | |||
| htcore.mark_step() | |||
There was a problem hiding this comment.
@kalyanjk what's the impact for input/output not introduced oom? should we add an argument in text-generation from cmd line?
There was a problem hiding this comment.
@kalyanjk ,why only mark_step() for 1x?
There was a problem hiding this comment.
For 8x mark_step will be introduced through a collective call.
There was a problem hiding this comment.
@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.
There was a problem hiding this comment.
Please update as below
if lazy_mode and (torch.distributed.is_initialized() is False or torch.distributed.get_world_size() == 1):
|
Wait we should not put mark step after the start of loop. Will create more graphs and perf is lower. |
|
| @@ -23,6 +23,7 @@ | |||
| _gaudi_prepare_4d_causal_attention_mask, | |||
| ) | |||
|
|
|||
| import habana_frameworks.torch.core as htcore | |||
There was a problem hiding this comment.
If you rebase to latest then this htcore import is not required, as it is part of PR#65
ghost
merged commit 
* Split the graphs to run with flash_attention on 1x * Added lazy_mode check and removed additional htcore import --------- Co-authored-by: Kalyan <kkumar@habana.ai>
* Split the graphs to run with flash_attention on 1x * Added lazy_mode check and removed additional htcore import --------- Co-authored-by: Kalyan <kkumar@habana.ai>
* Split the graphs to run with flash_attention on 1x * Added lazy_mode check and removed additional htcore import --------- Co-authored-by: Kalyan <kkumar@habana.ai>
* Split the graphs to run with flash_attention on 1x * Added lazy_mode check and removed additional htcore import --------- Co-authored-by: Kalyan <kkumar@habana.ai>
* Split the graphs to run with flash_attention on 1x * Added lazy_mode check and removed additional htcore import --------- Co-authored-by: Kalyan <kkumar@habana.ai>
* Split the graphs to run with flash_attention on 1x * Added lazy_mode check and removed additional htcore import --------- Co-authored-by: Kalyan <kkumar@habana.ai>
|
This PR solves the actual issue #126 |
With flash attention enabled for larger batch sizes, recipe arc hbm memory size exceeds QueueComputeScal arc hbm memory. Hence split the graph on 1x.