fix bug + style

Author: hyunwoongko

parallelformers/policies/base/auto.py

@@ -639,16 +639,24 @@ def __init__(self):
             ]
 
         with suppress(Exception):
-            from transformers.models.gptj.modeling_gptj import GPTJPreTrainedModel
+            from transformers.models.gptj.modeling_gptj import (
+                GPTJPreTrainedModel,
+            )
+
             from parallelformers.policies.gptj import GPTJPolicy
 
             self.builtin_policies[GPTJPreTrainedModel] = [
                 GPTJPolicy,
             ]
 
         with suppress(Exception):
-            from transformers.models.megatron_bert import MegatronBertPreTrainedModel
-            from parallelformers.policies.megtron_bert import MegatronBertPolicy
+            from transformers.models.megatron_bert import (
+                MegatronBertPreTrainedModel,
+            )
+
+            from parallelformers.policies.megtron_bert import (
+                MegatronBertPolicy,
+            )
 
             self.builtin_policies[MegatronBertPreTrainedModel] = [
                 MegatronBertPolicy,

parallelformers/policies/gptj.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from transformers.models.gptj.modeling_gptj import GPTJBlock
 
 from parallelformers.policies.base import Layer, Policy
 from parallelformers.utils import AllReduceLinear

parallelformers/policies/megtron_bert.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from transformers.models.megatron_bert.modeling_megatron_bert import \
-    MegatronBertLayer
+from transformers.models.megatron_bert.modeling_megatron_bert import (
+    MegatronBertLayer,
+)
 
 from parallelformers.policies.base import Layer, Policy
 from parallelformers.transformers.modeling_bert import BertEmbeddings_