Setting up my custom device map for a LLM
Hello community, greetings!This is my first post on forums, and i have a quick doubt on“”"let’s say i want to load bigscience/bloom-1b7 model, and i have just enough GPU RAM to fit the entire model except the lm_head. Therefore write a custom device_map as follows:device_map = {“transformer.word_embeddings”: 0,“transformer.word_embeddings_layernorm”: 0,“lm_head”: “cpu”,“transformer.h”: 0,“transformer.ln_f”: 0,}“”"This part has been taken from here:QuantizeTransformers models (huggingface.co)How can one know the architecture of that particular model before actually trying to setup a custom device_map.If i want to try help(model) to get the architecture of it, in general i wont be able to load this model in the first hand on  a GPU, if this is the case then how can i know which particular layers i can offload between cpu and gpu to quantize it?Let me know if there is any adequacy in my question.Thanks!
2023-08-10T13:01:06Z
Yeah that’s a question that I’ve been wondering about too - how do you find the layers for each model? I looked upbigscience/bloom-1b7 · Hugging Faceand searched for “lm_head” and wasn’t able to find any results.
2023-08-10T15:22:46Z
One naive solution I found out was to get the device map of that model by running it on a larger gpu machine and store it somewhere and later on use that device map according to the smaller gpu architecture for cpu and gpu offloading.model.REDACTED_TOKEN_map function will help for knowing the device map of a hugging face model