Repository: allenai/OLMo
File: olmo/model.py (Python)

Code before:
    def reset_parameters(self):
        if self.weight is not None:
            torch.nn.init.ones_(self.weight)  # type: ignore
        if self.bias is not None:
            torch.nn.init.zeros_(self.bias)  # type: ignore


class LayerNorm(LayerNormBase):
    """
    The default :class:`LayerNorm` implementation which can optionally run in low precision.
    """

    def __init__(self, config: ModelConfig, size: Optional[int] = None, low_precision: bool = False):
        super().__init__(config)
        self.normalized_shape = (size or config.d_model,)
        self.eps = 1e-05
        self.weight = nn.Parameter(torch.ones(self.normalized_shape, device=config.init_device))
        self.bias = nn.Parameter(torch.zeros(self.normalized_shape, device=config.init_device))

        # We always have weight and bias even if they are turned off/set to 1 and 0, because ROCm has a
        # bug where F.layer_norm() crashes during the backwards pass when no bias was given.
        if not self.config.layer_norm_with_affine:
            self.weight.requires_grad = False
            self.bias.requires_grad = False
        elif not self.config.include_bias:
            self.bias.requires_grad = False
        self.low_precision = low_precision

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.low_precision:
            module_device = x.device
            downcast_x = self._cast_if_autocast_enabled(x)
            downcast_weight = (
                self._cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
            )
            downcast_bias = self._cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
            with torch.autocast(enabled=False, device_type=module_device.type):
                return F.layer_norm(
                    downcast_x, self.normalized_shape, weight=downcast_weight, bias=downcast_bias, eps=self.eps
                )
        else:
            return F.layer_norm(x, self.normalized_shape, weight=self.weight, bias=self.bias, eps=self.eps)


class RMSLayerNorm(LayerNorm):
    """
    RMS layer norm, a simplified :class:`LayerNorm` implementation that can optionally run
    in low-precision.
    """

    def __init__(self, config: ModelConfig, size: Optional[int] = None, low_precision: bool = False):

Review comment (suggestion):
`self.weight` and `self.bias` may need to be buffers when we freeze them. @saurabh111233212 was in a similar situation with ReLoRa and discovered that FSDP has an implicit assumption that all FSDP-managed parameters must have the same value for `requires_grad`, i.e. the entire model must be trainable or frozen, you can't mix and match.

Code after:
    def reset_parameters(self):
        if self.weight is not None:
            torch.nn.init.ones_(self.weight)  # type: ignore
        if self.bias is not None:
            torch.nn.init.zeros_(self.bias)  # type: ignore


class LayerNorm(LayerNormBase):
    """
    The default :class:`LayerNorm` implementation which can optionally run in low precision.
    """

    def __init__(
        self,
        config: ModelConfig,
        size: Optional[int] = None,
        low_precision: bool = False,
        elementwise_affine: Optional[bool] = None,
    ):
        super().__init__(config)
        self.normalized_shape = (size or config.d_model,)
        self.eps = 1e-05
        self.low_precision = low_precision

        # We always have weight and bias even if they are turned off/set to 1 and 0, because ROCm has a
        # bug where F.layer_norm() crashes during the backwards pass when no bias was given.
        # When they are turned off, they need to be buffers, because FSDP can't handle the situation
        # where some parameters don't require gradients.

        if elementwise_affine is None:
            elementwise_affine = self.config.layer_norm_with_affine
        weight = torch.ones(self.normalized_shape, device=config.init_device)
        if elementwise_affine:
            self.register_parameter("weight", nn.Parameter(weight))
        else:
            self.register_buffer("weight", weight, persistent=False)

        needs_bias = elementwise_affine and self.config.include_bias
        bias = torch.zeros(self.normalized_shape, device=config.init_device)
        if needs_bias:
            self.register_parameter("bias", nn.Parameter(bias))
        else:
            self.register_buffer("bias", bias, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.low_precision:
            module_device = x.device
            downcast_x = self._cast_if_autocast_enabled(x)
            downcast_weight = (
                self._cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
            )