Supporting device_map = 'auto' similar to the one in .from_pretrained method from Huggingface

[Ahmed-Roushdy]

The pruned model is saved using torch.save and torch.load for loading the model. I was wondering if there is a way to use a similar method such as device_map='auto' similar to the one in .from_pretrained method from Huggingface

[horseee]

Hi. It would be challenging to do this since the pruned model does not follow a uniform configuration, like different dimensions for different modules and different head numbers for different layers.

[Ahmed-Roushdy]

Thank you. Do you suggest any similar methods that can enable distributed training while working with the Train method from Huggingface?

[KBhandari11]

I am interested in this as well. Is it possible to implement it on Llama2-70b? Or use some distributed method to prune the model?