Dmitry

I'm trying to launch MipNerf traing but like in https://github.com/kakaobrain/nerf-factory/issues/12 I have huge CPU memory consumption while GPU idles. I guess it's because data loader and sampler working on CPU....

# What does this PR do? This PR fixes HPU Graphs usage and Flash Attention for Gemma model. Changes are based on Starcoder 2 and Qwen 2 implementations. ## Before...