From 5131555f2450f9ee92b991159155d914bcdb6e4e Mon Sep 17 00:00:00 2001 From: Michael Wyatt Date: Mon, 11 Jul 2022 11:19:23 -0700 Subject: [PATCH] use HF NeoX --- deepspeed/module_inject/replace_policy.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepspeed/module_inject/replace_policy.py b/deepspeed/module_inject/replace_policy.py index f4054c633810..624ade6dfc34 100755 --- a/deepspeed/module_inject/replace_policy.py +++ b/deepspeed/module_inject/replace_policy.py @@ -333,9 +333,9 @@ def __init__(self, client_module, inference=True, megatron_v2=True): GPTNEOXLayerPolicy._orig_layer_class = None else: try: - import megatron - from megatron.model.transformer import ParallelTransformerLayerPipe - GPTNEOXLayerPolicy._orig_layer_class = ParallelTransformerLayerPipe + import transformers + from transformers import GPTNeoXLayer + GPTNEOXLayerPolicy._orig_layer_class = GPTNeoXLayer except ImportError: GPTNEOXLayerPolicy._orig_layer_class = None