Add mark step and inplace residual add in llama model code

[puneeshkhanna]

Mark step helping in reducing workspace memory by approx twice of (BS,seq len, hidden dim).

Inplace add helping in reducing persistent tensors by approx twice of (BS, seq len, hidden dim).

Add lazy mode parameter.

---

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[puneeshkhanna]

As an example for the config of BS-172, seq len-2048, hidden dim-8192 (size is ~5.3 GB) for llama-70B on 8x.
Max memory usage (without flash attention) reduced from ~86 Gb to ~66Gb.
Max memory usage (with flash attention) reduced from ~70 GB to ~59GB.

We can further go higher batch sizes because of reduced memory consumptions with the changes in this PR hence very important changes.

@regisss - please review and merge.

[regisss]

LGTM! Just need to run make style.

Quick question, the new mark_step in the forward of GaudiLlamaModel will also benefit training right?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[puneeshkhanna]

@regisss - mark step has no side effect on training. It should help in reducing workspace memory mainly.

[puneeshkhanna]

@regisss - make style fixed. Just required an empty line after the import statement. It should pass now.