use the correct available memory API for XPU

[faaany]

What does this PR do?

This PR introduces the new API for get xpu available memory. Below is an example of before and after:

from transformers import AutoModelForSeq2SeqLM

model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp", device_map="auto")
print(model.hf_device_map)

Before

{'': 'cpu'}

After

{'shared': 0, 'decoder.embed_tokens': 0, 'encoder.embed_tokens': 0, 'encoder.block.0': 0, 'encoder.block.1': 0, 'encoder.block.2': 0, 'encoder.block.3': 0, 'encoder.block.4': 0, 'encoder.block.5': 0, 'encoder.block.6': 1, 'encoder.block.7': 1, 'encoder.block.8': 1, 'encoder.block.9': 1, 'encoder.block.10': 1, 'encoder.block.11': 1, 'encoder.block.12': 1, 'encoder.block.13': 1, 'encoder.block.14': 2, 'encoder.block.15': 2, 'encoder.block.16': 2, 'encoder.block.17': 2, 'encoder.block.18': 2, 'encoder.block.19': 2, 'encoder.block.20': 2, 'encoder.block.21': 2, 'encoder.block.22': 3, 'encoder.block.23': 3, 'encoder.final_layer_norm': 3, 'encoder.dropout': 3, 'decoder.block.0': 3, 'decoder.block.1': 3, 'decoder.block.2': 3, 'decoder.block.3': 3, 'decoder.block.4': 4, 'decoder.block.5': 4, 'decoder.block.6': 4, 'decoder.block.7': 4, 'decoder.block.8': 4, 'decoder.block.9': 4, 'decoder.block.10': 5, 'decoder.block.11': 5, 'decoder.block.12': 5, 'decoder.block.13': 5, 'decoder.block.14': 5, 'decoder.block.15': 5, 'decoder.block.16': 6, 'decoder.block.17': 6, 'decoder.block.18': 6, 'decoder.block.19': 6, 'decoder.block.20': 6, 'decoder.block.21': 6, 'decoder.block.22': 7, 'decoder.block.23': 7, 'decoder.final_layer_norm': 7, 'decoder.dropout': 7, 'lm_head': 7}

[muellerzr]

Thanks! Overall this looks good to me, let's use device_index however rather than i (or something similar to avoid single-letter input names)

[faaany]

Thanks! Overall this looks good to me, let's use device_index however rather than i (or something similar to avoid single-letter input names)

thanks for pointing this out! Code Updated.

I also added a version check.

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

[faaany]

Hi @muellerzr , how do you think about the new update?

[muellerzr]

Thought about it a bit more, let's:

1. Move this to utils.memory
2. Use the new implementation I've shown, which takes care of things a bit cleaner.

[muellerzr]

Great! This looks much cleaner :)

[BenjaminBossan]

Thanks, LGTM. I have a clarifying question, but it's not a blocker.

[faaany]

ok, I think we can merge now. 😊