fix jit memory issue

[jiqing-feng]

Fix #604

Hi @echarlaix . As we discussed before, the jit trace memory issue can only workaround for now. I think the best way is to skip the save and load step and just return the cls after jit tracing. WDYT?

Jit trace memory issue: 96726

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

[yao-matrix]

@jiqing-feng , do you observe memory saving with your PR? any before vs after proof data?

[jiqing-feng]

@jiqing-feng , do you observe memory saving with your PR? any before vs after proof data?

Script:

import torch
from optimum.intel.ipex import IPEXModelForCausalLM
model = IPEXModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", torch_dtype=torch.bfloat16, export=True)

Original model size: 13.5G

Memory size before change: 51.7G

Memory size after my change: 13.6G

[echarlaix]

Looks great, thanks @jiqing-feng