Run Llama2 with torch.compile on Gaudi2

[kausikmaiti]

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?

[vivekgoe]

@kausikmaiti looks good to me. @regisss please review and help merge this if it looks ok to you. Thanks.

[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.

[libinta]

@kausikmaiti I have following questions

1. what's the default cmd for llama2? torch compile or use_lazy_mode? if the default is torch.compile, please also update the readme
2. what's the performance comparison with use_lazy_mode and use_hpu_graph?
3. any depenency to the to-be-released docker?

[kausikmaiti]

@kausikmaiti I have following questions

1. what's the default cmd for llama2? torch compile or use_lazy_mode? if the default is torch.compile, please also update the readme
2. what's the performance comparison with use_lazy_mode and use_hpu_graph?
3. any depenency to the to-be-released docker?

@libinta, Please find my answers below.

1. Here, I am just adding support for torch.compile. But, default mode is still unchanged. So, we don't need to update readme here.
2. This is the very first step of (functional) enablement. Performance data is not yet available.
3. Yes. It has dependency on the content of release 1.14.

[bhargaveede]

We could avoid this by using args.torch_compile directly in 312 and 492 with lazy_mode = not args.torch_compile.

[kausikmaiti]

Yes, that can be done of course. But, I wanted to have the decision making about use_lazy_mode in one place for all. A secondary reason is that I plan to include additional checks in future PR to control use_lazy_mode.

[regisss]

Yes. It has dependency on the content of release 1.14.

So we should not merge it before 1.14 is released @libinta right?

[bhargaveede]

@kausikmaiti Can we add model specific check as generation using torch.compile isn't verified on models

[kausikmaiti]

Added model specific check in separate commit. Please review.

[regisss]

@kausikmaiti Let's also add a test to: https://github.com/huggingface/optimum-habana/blob/main/tests/test_text_generation_example.py

You can define a new test at the end of this file:

@pytest.mark.parametrize("model_name, baseline", MODELS_TO_TEST["torch_compile"])
def test_text_generation_torch_compile(model_name: str, baseline: float, token: str):
    _test_text_generation(model_name, baseline, token, torch_compile=True)

and adding a new torch_compile entry in the dict here:

optimum-habana/tests/test_text_generation_example.py

Line 15 in 979c132

MODELS_TO_TEST = {

[kalyanjk]

os.environ["WORLD_SIZE"] = "0" ? WORLD_SIZE should be set to 1 for 1x runs

[kausikmaiti]

As you mentioned offline, WORLD_SIZE setting does not matter, as I'm not using deepspeed / gaudi_spawn.py script.
Also as per my observation, if I don't set WORLD_SIZE=0, due to the logic like "use_deepspeed = args.world_size > 0", setup_distributed_model() gets called and the test fails at very early stage while importing deepspeed. This is not the expectation.