minor fix

NVIDIA · Dec 16, 2024 · dcf81f5 · dcf81f5
1 parent 30006df
commit dcf81f5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/nemo/collections/llm/gpt/model/hyena.py b/nemo/collections/llm/gpt/model/hyena.py
@@ -99,12 +99,17 @@ class HyenaConfig(TransformerConfig, io.IOMixin):
     remove_activation_post_first_layer: bool = True
     cross_entropy_loss_fusion: bool = True
     tp_comm_overlap: bool = True
+    bias_activation_fusion: bool = True
+    bias_dropout_add_fusion: bool = True
 
     def __post_init__(self):
         super().__post_init__()
         self.hyena_no_weight_decay_cond_fn = hyena_no_weight_decay_cond if self.hyena_filter_no_wd else None
 
     def configure_model(self, tokenizer) -> "MCoreHyenaModel":
+
+        self.bias_activation_fusion = False if self.remove_activation_post_first_layer else self.bias_activation_fusion
+
         model = MCoreHyenaModel(
             self,
             hyena_stack_spec=hyena_stack_spec,