Method _replace_rmsnorm_modules

verify.py:608–637

Replace RMSNorm modules. Since there is no standard nn.RMSNorm, we look for common class names and attributes.

(self, repl: KernelReplacement)

        return count

    def _replace_rmsnorm_modules(self, repl: KernelReplacement) -> int:
        """
        Replace RMSNorm modules. Since there is no standard nn.RMSNorm,
        we look for common class names and attributes.
        """
        count = 0
        rmsnorm_names = {"RMSNorm", "LlamaRMSNorm", "T5LayerNorm", "GemmaRMSNorm"}

        for name, module in list(self.model.named_modules()):
            cls_name = type(module).__name__
            # Match by class name or by having 'weight' but no 'bias' and a norm-like name
            is_rmsnorm = (
                cls_name in rmsnorm_names
                or (hasattr(module, "weight")
                    and hasattr(module, "eps")
                    and not hasattr(module, "bias")
                    and cls_name.lower().endswith("norm")
                    and not isinstance(module, nn.LayerNorm))
            )

            if is_rmsnorm:
                self._original_modules[name] = module
                wrapper = _RMSNormWrapper(module, repl.module_fn)
                parts = name.split(".")
                parent = self.model
                for p in parts[:-1]:
                    parent = getattr(parent, p)
                setattr(parent, parts[-1], wrapper)
                count += 1
        return count

    @property
    def applied_summary(self) -> List[str]: