| 49 | return layers |
| 50 | |
| 51 | def move_embed(model, device): |
| 52 | if isinstance(model, LlamaForCausalLM): |
| 53 | model.model.embed_tokens = model.model.embed_tokens.to(device) |
| 54 | elif isinstance(model, OPTForCausalLM): |
| 55 | model.model.decoder.embed_tokens = model.model.decoder.embed_tokens.to(device) |
| 56 | model.model.decoder.embed_positions = model.model.decoder.embed_positions.to(device) |
| 57 | elif isinstance(model, BloomForCausalLM): |
| 58 | model.transformer.word_embeddings = model.transformer.word_embeddings.to(device) |
| 59 | model.transformer.word_embeddings_layernorm = model.transformer.word_embeddings_layernorm.to(device) |
| 60 | elif "mpt" in str(model.__class__).lower(): |
| 61 | model.transformer.wte = model.transformer.wte.to(device) |
| 62 | model.transformer.emb_drop = model.transformer.emb_drop.to(device) |
| 63 | elif "falcon" in str(model.__class__).lower(): |
| 64 | model.transformer.word_embeddings = model.transformer.word_embeddings.to(device) |
| 65 | elif "bigcode" in str(model.__class__).lower(): |
| 66 | model.transformer.wte = model.transformer.wte.to(device) |
| 67 | model.transformer.wpe = model.transformer.wpe.to(device) |
| 68 | model.transformer.drop = model.transformer.drop.to(device) |
| 69 | elif "neox" in str(model.__class__).lower(): |
| 70 | model.gpt_neox.embed_in = model.gpt_neox.embed_in.to(device) |
| 71 | model.gpt_neox.emb_dropout = model.gpt_neox.emb_dropout.to(device) |
| 72 | model.embed_out = model.embed_out.to(device) |
| 73 | else: |
| 74 | raise NotImplementedError(type(model)) |
| 75 | |
| 76 | @torch.no_grad() |
| 77 | def run_awq( |