Replace all nn.LayerNorm modules with optimized wrapper.
(self, repl: KernelReplacement)
| 591 | return count |
| 592 | |
| 593 | def _replace_layernorm_modules(self, repl: KernelReplacement) -> int: |
| 594 | """Replace all nn.LayerNorm modules with optimized wrapper.""" |
| 595 | count = 0 |
| 596 | for name, module in list(self.model.named_modules()): |
| 597 | if isinstance(module, nn.LayerNorm): |
| 598 | self._original_modules[name] = module |
| 599 | wrapper = _LayerNormWrapper(module, repl.module_fn) |
| 600 | parts = name.split(".") |
| 601 | parent = self.model |
| 602 | for p in parts[:-1]: |
| 603 | parent = getattr(parent, p) |
| 604 | setattr(parent, parts[-1], wrapper) |
| 605 | count += 1 |
| 606 | return count |
| 607 | |
| 608 | def _replace_rmsnorm_modules(self, repl: KernelReplacement) -> int: |
| 609 | """ |
no test coverage detected