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

Code before:

    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)
            self.weight.requires_grad = 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)
            self.bias.requires_grad = 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
            )
            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)



Review comment (suggestion):
This is redundant since buffers default to no grad.
```suggestion
```

Code after:

    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
            )
            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):
    """