pymllm.layers.rms_norm
======================

.. py:module:: pymllm.layers.rms_norm


Classes
-------

.. autoapisummary::

   pymllm.layers.rms_norm.RMSNorm
   pymllm.layers.rms_norm.GemmaRMSNorm


Module Contents
---------------

.. py:class:: RMSNorm(hidden_size, eps = 1e-06)

   Bases: :py:obj:`pymllm.layers.base.MllmBaseLayer`


   RMSNorm layer implemented with FlashInfer kernel.


   .. py:attribute:: hidden_size


   .. py:attribute:: eps
      :value: 1e-06



   .. py:attribute:: weight


   .. py:method:: forward(x, residual = None)


.. py:class:: GemmaRMSNorm(hidden_size, eps = 1e-06)

   Bases: :py:obj:`pymllm.layers.base.MllmBaseLayer`


   Gemma-style RMSNorm layer implemented with FlashInfer kernel.


   .. py:attribute:: hidden_size


   .. py:attribute:: eps
      :value: 1e-06



   .. py:attribute:: weight


   .. py:method:: forward(x, residual = None)