fix ia3 mlp infused adapter (NVIDIA#8597)

Signed-off-by: Chen Cui <chcui@nvidia.com>
rohitrango · Mar 6, 2024 · a2cee24 · a2cee24
1 parent e3f8282
commit a2cee24
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py b/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py
@@ -277,6 +277,10 @@ def glu(x):
             else:
                 intermediate_parallel = self.activation_func(intermediate_parallel)
 
+        infused_adapter = self.get_adapter_module(AdapterName.MLP_INFUSED)
+        if infused_adapter:
+            intermediate_parallel = infused_adapter(intermediate_parallel)
+
         # [s, b, h]
         output, output_bias = self.linear_fc2(intermediate_parallel)