[Bugfix] Fix transformers import order in megatron scripts (#5)

awslabs · Jan 19, 2023 · f868eb8 · f868eb8
1 parent 353d6a7
commit f868eb8
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 12 deletions.
diff --git a/examples/albert/megatron_hf.py b/examples/albert/megatron_hf.py
@@ -12,6 +12,8 @@
 import torch
 import torch.nn.functional as F
 
+from transformers import AutoConfig, AlbertModel
+
 from megatron import get_args
 from megatron import print_rank_0
 from megatron import get_timers
@@ -35,8 +37,6 @@ def get_model(
     impl="slapo",
     delay_init=True,
 ):
-    from transformers import AutoConfig, AlbertModel
-
     config = AutoConfig.from_pretrained(model_name)
     if padded_vocab_size is not None:
         config.vocab_size = padded_vocab_size

diff --git a/examples/bert/megatron_hf.py b/examples/bert/megatron_hf.py
@@ -12,6 +12,8 @@
 import torch
 import torch.nn.functional as F
 
+from transformers import AutoConfig, BertModel
+
 from megatron import get_args
 from megatron import print_rank_0
 from megatron import get_timers
@@ -35,8 +37,6 @@ def get_model(
     impl="slapo",
     delay_init=True,
 ):
-    from transformers import AutoConfig, BertModel
-
     config = AutoConfig.from_pretrained(model_name)
     if padded_vocab_size is not None:
         config.vocab_size = padded_vocab_size

diff --git a/examples/gpt/megatron_hf.py b/examples/gpt/megatron_hf.py
@@ -7,6 +7,9 @@
 import os
 
 import torch
+
+from transformers import AutoConfig, GPTNeoModel
+
 from functools import partial
 from megatron import get_args
 from megatron import print_rank_0
@@ -30,8 +33,6 @@ def get_model(
     impl="slapo",
     delay_init=True,
 ):
-    from transformers import AutoConfig, GPTNeoModel
-
     config = AutoConfig.from_pretrained(model_name)
     if padded_vocab_size is not None:
         config.vocab_size = padded_vocab_size

diff --git a/examples/opt/megatron_hf.py b/examples/opt/megatron_hf.py
@@ -7,6 +7,9 @@
 import os
 
 import torch
+
+from transformers import AutoConfig, OPTModel
+
 from functools import partial
 from megatron import get_args
 from megatron import print_rank_0
@@ -30,8 +33,6 @@ def get_model(
     impl="slapo",
     delay_init=True,
 ):
-    from transformers import AutoConfig, OPTModel
-
     config = AutoConfig.from_pretrained(model_name)
     if padded_vocab_size is not None:
         config.vocab_size = padded_vocab_size

diff --git a/examples/roberta/megatron_hf.py b/examples/roberta/megatron_hf.py
@@ -9,6 +9,8 @@
 import torch
 import torch.nn.functional as F
 
+from transformers import AutoConfig, RobertaModel
+
 from megatron import get_args
 from megatron import print_rank_0
 from megatron import get_timers
@@ -32,8 +34,6 @@ def get_model(
     impl="slapo",
     delay_init=True,
 ):
-    from transformers import AutoConfig, RobertaModel
-
     config = AutoConfig.from_pretrained(model_name)
     if padded_vocab_size is not None:
         config.vocab_size = padded_vocab_size

diff --git a/examples/t5/megatron_hf.py b/examples/t5/megatron_hf.py
@@ -10,6 +10,8 @@
 
 import torch
 
+from transformers import AutoConfig, T5Model
+
 from megatron import get_args, get_timers, mpu, print_rank_0
 from megatron.data.dataset_utils import build_train_valid_test_datasets
 from megatron.model import ModelType
@@ -58,8 +60,6 @@ def get_model(
     impl="slapo",
     delay_init=True,
 ):
-    from transformers import AutoConfig, T5Model
-
     config = AutoConfig.from_pretrained(model_name)
     config.vocab_size = padded_vocab_size
     config.use_cache = False