Replace all LlamaRMSNorm modules with FTLlamaRMSNorm modules
(model)
| 22 | |
| 23 | |
| 24 | def make_quant_norm(model): |
| 25 | """ |
| 26 | Replace all LlamaRMSNorm modules with FTLlamaRMSNorm modules |
| 27 | """ |
| 28 | |
| 29 | for name, m in model.named_modules(): |
| 30 | if not isinstance(m, LlamaRMSNorm): |
| 31 | continue |
| 32 | |
| 33 | norm = FTLlamaRMSNorm(m.weight, m.variance_epsilon) |
| 34 | |
| 35 | if "." in name: |
| 36 | parent_name = name.rsplit(".", 1)[0] |
| 37 | child_name = name[len(parent_name) + 1 :] |
| 38 | parent = model.get_submodule(parent_name) |
| 39 | else: |
| 40 | parent_name = "" |
| 41 | parent = model |
| 42 | child_name = name |
| 43 | |
| 44 | # print(f"Replacing {name} with quant_attn; parent: {parent_name}, child's name: {child_name}") |
| 45 | |
| 46 | setattr(parent, child_name, norm) |
nothing calls this directly
no test coverage detected