From 09bf3f574d779ab62721d45810e49ff0dca22a57 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Tue, 5 Mar 2024 21:29:49 -0800 Subject: [PATCH] fix ia3 mlp infused adapter Signed-off-by: Chen Cui --- .../nlp/modules/common/megatron/adapters/mcore_mixins.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py b/nemo/collections/nlp/modules/common/megatron/adapters/mcore_mixins.py index 59b09b643f89..f0d05b48e543 100644 --- 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)