(model)
| 30 | setattr(layer, name, new_module) |
| 31 | |
| 32 | def get_blocks(model): |
| 33 | if model.__class__.__name__ == 'LlamaForCausalLM': |
| 34 | layers = model.model.layers |
| 35 | elif isinstance(model, OPTForCausalLM): |
| 36 | layers = model.model.decoder.layers |
| 37 | elif isinstance(model, BloomForCausalLM): |
| 38 | layers = model.transformer.h |
| 39 | elif "mpt" in str(model.__class__).lower(): |
| 40 | layers = model.transformer.blocks |
| 41 | elif "falcon" in str(model.__class__).lower(): |
| 42 | layers = model.transformer.h |
| 43 | elif "bigcode" in str(model.__class__).lower(): |
| 44 | layers = model.transformer.h |
| 45 | elif "neox" in str(model.__class__).lower(): |
| 46 | layers = model.gpt_neox.layers |
| 47 | else: |
| 48 | raise NotImplementedError(type(model)) |
| 49 | return layers |
| 50 | |
| 51 | def move_embed(model, device): |
| 52 | if isinstance(model, LlamaForCausalLM): |
no outgoing calls
no test coverage detected