change back to warning

RaymondLi0 · RaymondLi0 · commit 0ff574679723 · 2023-03-24T18:36:40.000-04:00
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -604,7 +604,8 @@ def __init__(self, init_method,
                 raise ImportError('einops is not installed, please install with pip install einops')
             
             if self.checkpoint_core_attention:
-                raise NotImplementedError("Using selective recomputation with flash-attn: this is not implemented.")
+                print_rank_0("  Warning, using selective recomputation with flash-attn: this is already handled in the "
+                             "flash-attn library and has no effect.")
             self.core_attention_flash = FlashSelfAttention(
                 causal=True, attention_dropout=args.attention_dropout
             )

Original file line number	Diff line number	Diff line change
`@@ -604,7 +604,8 @@ def __init__(self, init_method,`
`604`	`604`	`raise ImportError('einops is not installed, please install with pip install einops')`
`605`	`605`
`606`	`606`	`if self.checkpoint_core_attention:`
`607`		`- raise NotImplementedError("Using selective recomputation with flash-attn: this is not implemented.")`
	`607`	`+ print_rank_0(" Warning, using selective recomputation with flash-attn: this is already handled in the "`
	`608`	`+ "flash-attn library and has no effect.")`
`608`	`609`	`self.core_attention_flash = FlashSelfAttention(`
`609`	`610`	`causal=True, attention_dropout=args.attention_dropout`
`610`	`611`	`)`