perf improvement

Author: root

llmfoundry/models/layers/attention.py

@@ -514,23 +514,29 @@ def _apply_qk_gn(
 ):
     dtype = query.dtype
 
-    w = _expand_params(n_heads, q_ln.weight)
-    b = _expand_params(n_heads, q_ln.bias)
     if isinstance(q_ln, LPLayerNorm):
-        query = low_precision_groupnorm(query, n_heads, w, b,
+        query = low_precision_groupnorm(query,
+                                        n_heads,
+                                        q_ln.weight,
+                                        q_ln.bias,
                                         eps=q_ln.eps).to(dtype)
     elif isinstance(q_ln, nn.LayerNorm):
+        w = _expand_params(n_heads, q_ln.weight)
+        b = _expand_params(n_heads, q_ln.bias)
         query = nn.functional.group_norm(query, n_heads, w, b, eps=q_ln.eps)
     else:
         raise ValueError(
             f'qk_gn not applicable for given q_ln type ({type(q_ln)=}).')
 
-    w = _expand_params(kv_n_heads, k_ln.weight)
-    b = _expand_params(kv_n_heads, k_ln.bias)
     if isinstance(k_ln, LPLayerNorm):
-        key = low_precision_groupnorm(key, kv_n_heads, w, b,
+        key = low_precision_groupnorm(key,
+                                      kv_n_heads,
+                                      k_ln.weight,
+                                      k_ln.bias,
                                       eps=k_ln.eps).to(dtype)
     elif isinstance(k_ln, nn.LayerNorm):
+        w = _expand_params(kv_n_heads, k_ln.weight)
+        b = _expand_params(kv_n_heads, k_ln.bias)
         key = nn.functional.group_norm(key, kv_n_heads, w, b, eps=k_ln.eps)
     else:
         raise ValueError(

llmfoundry/models/layers/norm.py

@@ -53,6 +53,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             )
 
 
+def _expand_params(x: torch.Tensor, param: Optional[torch.Tensor] = None):
+    # repeat param if params are applied per group
+    if param is None:
+        return None
+    if x.shape[-1] == param.shape[-1]:
+        return param
+    return param.repeat(x.shape[-1] // param.shape[-1])
+
+
 def low_precision_groupnorm(
     x: torch.Tensor,
     groups: int,
@@ -62,10 +71,14 @@ def low_precision_groupnorm(
 ):
     device = x.device
     downcast_x = _cast_if_autocast_enabled(x)
-    downcast_weight = _cast_if_autocast_enabled(
-        weight) if weight is not None else weight
-    downcast_bias = _cast_if_autocast_enabled(
-        bias) if bias is not None else bias
+    downcast_weight, downcast_bias = None, None
+    if weight is not None:
+        downcast_weight = _cast_if_autocast_enabled(weight)
+        downcast_weight = _expand_params(x, downcast_weight)
+    if bias is not None:
+        downcast_bias = _cast_if_autocast_enabled(bias)
+        downcast_bias = _expand_params(x, downcast_bias)
+
     with torch.autocast(enabled=False, device_type=device.type):
         return torch.nn.functional.group_norm(
             downcast_x,