Fix attention mask inference (NVIDIA#7213)

* Fix attention mask inference Signed-off-by: Cheng-Ping Hsieh <[email protected]> * Fix config Signed-off-by: Cheng-Ping Hsieh <[email protected]> --------- Signed-off-by: Cheng-Ping Hsieh <[email protected]>
guyueh1 · Aug 14, 2023 · dede2b4 · dede2b4
1 parent 63c45cf
commit dede2b4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/examples/nlp/language_modeling/tuning/megatron_gpt_peft_eval.py b/examples/nlp/language_modeling/tuning/megatron_gpt_peft_eval.py
@@ -172,6 +172,8 @@ def main(cfg) -> None:
     )
 
     model.freeze()
+    if not cfg.model.get('use_flash_attention', False):
+        cfg.inference.compute_attention_mask = True
     config = OmegaConf.to_container(cfg.inference, resolve=True)
     model.set_inference_config(config)
 

diff --git a/examples/nlp/language_modeling/tuning/megatron_gpt_sft.py b/examples/nlp/language_modeling/tuning/megatron_gpt_sft.py
@@ -204,6 +204,8 @@ def main(cfg) -> None:
         model = load_from_checkpoint_dir(MegatronGPTSFTModel, cfg, trainer, modify_confg_fn=_modify_config)
 
     if 'inference' in cfg:
+        if not cfg.model.use_flash_attention:
+            cfg.inference.compute_attention_mask = True
         config = OmegaConf.to_container(cfg.inference, resolve=True)
         model.set_inference_config(config)