Fix affine qnorm config (#319)

allenai · Oct 8, 2023 · 5a7d8af · 5a7d8af
1 parent 1d4cbda
commit 5a7d8af
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/olmo/model.py b/olmo/model.py
@@ -340,7 +340,7 @@ def __init__(self, layer_id: int, config: ModelConfig):
                 size=config.d_model // config.n_heads if config.multi_query_attention else None,
                 elementwise_affine=config.attention_layer_norm_with_affine,
             )
-            self.q_norm = LayerNormBase.build(config, elementwise_affine=True)
+            self.q_norm = LayerNormBase.build(config, elementwise_affine=config.attention_layer_norm_with_affine)
 
         # Activation function.
         self.act = Activation.build(config)