def fuse_layer_norms(cls, model_wrapper: ModelWrapper):
        """Fuse layernorms into adjacent linear layers."""
        # TODO(jambayk): should we support models with layernorms? these require:
        # - subtracting mean from embedding
        # - baking mean into output layers
        # - replacing layernorm with RMSNorm
        # Model architecture changes are required
        if isinstance(model_wrapper.get_pre_head_layernorm(False), nn.LayerNorm):
            raise ValueError("Model uses LayerNorm. Only RMSNorm fusion is supported.")

        logger.debug("Fusing layernorms into adjacent linear layers")

        # untie embedding and lm head
        model_wrapper.maybe_untie_word_embeddings()

        # Layers: Fuse layernorms into adjacent linear layers
        for layer_wrapper in model_wrapper.get_layer_wrappers():
            cls.fuse_ln_linear(layer_wrapper.get_first_layer_norm(False), layer_wrapper.get_attention_inputs(False))
            cls.fuse_ln_linear(layer_wrapper.get_second_layer_norm(False), layer_wrapper.get_mlp_inputs(False))

        # LM Head: Fuse layernorm into linear layer
        cls.fuse_ln_linear(model_wrapper.get_pre_head_layernorm(False), [model_wrapper.get_lm_head(False)])
Fuse layernorms into adjacent linear layers.
olive/passes/pytorch/rotate.py