| 52 | lm_head.in_features = dim |
| 53 | |
| 54 | def model_slice(model: nn.Module, dim: int, prune_lm_head=False): |
| 55 | slice_embeddings(model.model.embed_tokens, dim) |
| 56 | |
| 57 | for layer_idx, layer in enumerate(model.model.layers): |
| 58 | slice_attention_inputs(layer.self_attn, dim) |
| 59 | slice_attention_output(layer.self_attn, dim) |
| 60 | slice_mlp_input(layer.mlp, dim) |
| 61 | layer.attn_shortcut_Q.data = layer.attn_shortcut_Q.data[:dim,:dim] |
| 62 | if layer_idx < len(model.model.layers)-1 or prune_lm_head: |
| 63 | slice_mlp_output(layer.mlp, dim) |
| 64 | layer.mlp_shortcut_Q.data = layer.mlp_shortcut_Q.data[:dim,:dim] |
| 65 | else: |
| 66 | layer.mlp_shortcut_Q.data = layer.mlp_shortcut_Q.data[:dim] |
| 67 | |
| 68 | |
| 69 | layer.input_layernorm.weight.data = layer.input_layernorm.weight.data[:dim] * math.sqrt(layer.input_layernorm.weight.data.shape[0]/dim) |
| 70 | layer.post_attention_layernorm.weight.data = layer.post_attention_layernorm.weight.data[:dim] * math.sqrt(layer.post_attention_layernorm.weight.data.shape[0]/dim) |
| 71 | |
| 72 | if prune_lm_head: |
| 73 | slice_lm_head(model.lm_head, dim) |
| 74 | model.model.norm.weight.data = model.model.norm.weight.data[:dim] |