Replace all nn.Linear modules with optimized matmul wrapper.
(self, repl: KernelReplacement)
| 573 | return count |
| 574 | |
| 575 | def _replace_linear_modules(self, repl: KernelReplacement) -> int: |
| 576 | """Replace all nn.Linear modules with optimized matmul wrapper.""" |
| 577 | count = 0 |
| 578 | for name, module in list(self.model.named_modules()): |
| 579 | if isinstance(module, nn.Linear): |
| 580 | # Save original |
| 581 | self._original_modules[name] = module |
| 582 | # Create wrapper |
| 583 | wrapper = _LinearWrapper(module, repl.module_fn) |
| 584 | # Install wrapper |
| 585 | parts = name.split(".") |
| 586 | parent = self.model |
| 587 | for p in parts[:-1]: |
| 588 | parent = getattr(parent, p) |
| 589 | setattr(parent, parts[-1], wrapper) |
| 590 | count += 1 |
| 591 | return count |
| 592 | |
| 593 | def _replace_layernorm_modules(self, repl: KernelReplacement) -> int: |
| 594 | """Replace all nn.LayerNorm modules with optimized wrapper.""" |
no test coverage detected