[bug_fix] kv_channels is used when available (NVIDIA#5066)

* fix bug s.t kv_channels is used when available Signed-off-by: arendu <adithya.r@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: arendu <adithya.r@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
titu1994 · Oct 6, 2022 · e307799 · e307799
1 parent 1a689f7
commit e307799
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_t5_adapter_model.py b/nemo/collections/nlp/models/language_modeling/megatron_t5_adapter_model.py
@@ -369,9 +369,16 @@ def _add_adapters_to_component(self, component, layer_cfg, adapter_name_keys):
                             in_features=layer_cfg.ffn_hidden_size // layer_cfg.tensor_model_parallel_size
                         )
                     else:
-                        cfg = InfusedAdapterConfig(
-                            in_features=layer_cfg.hidden_size // layer_cfg.tensor_model_parallel_size
-                        )
+                        if layer_cfg.get('kv_channels', None):
+                            cfg = InfusedAdapterConfig(
+                                in_features=layer_cfg.kv_channels
+                                * layer_cfg.num_attention_heads
+                                // layer_cfg.tensor_model_parallel_size
+                            )
+                        else:
+                            cfg = InfusedAdapterConfig(
+                                in_features=layer_cfg.hidden_size // layer_cfg.tensor_model_parallel_size
+                            )
                     module.add_adapter(name=adapter_key, cfg=cfg)
 
     def _component_state_dict(self, component_name, component, adapter_name_keys):