From fa9fe945f59b7aedbdbd70d4f36edd757a58f4d6 Mon Sep 17 00:00:00 2001 From: Logan Adams Date: Wed, 13 Sep 2023 16:27:50 -0700 Subject: [PATCH] Replace with versions from transformers and torch --- tests/unit/compression/test_compression.py | 6 +++--- .../unit/ops/accelerators/test_accelerator_backward.py | 10 ++++++---- .../unit/ops/accelerators/test_accelerator_forward.py | 10 ++++++---- 3 files changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/unit/compression/test_compression.py b/tests/unit/compression/test_compression.py index c6e5031349cb..e36b7fce7c9b 100644 --- a/tests/unit/compression/test_compression.py +++ b/tests/unit/compression/test_compression.py @@ -9,8 +9,8 @@ import numpy as np from unit.megatron_model import get_gpt2_model from deepspeed.compression.compress import init_compression -from unit.modeling import BertConfig -from unit.modelingpreln import BertEncoder as BertEncoderPreln +from transformers.models.bert.configuration_bert import BertConfig +from transformers.models.bert.modeling_bert import BertEncoder as BertEncoder from deepspeed.compression.basic_layer import LinearLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress from deepspeed.compression.helper import convert_conv1d_to_linear from deepspeed.accelerator import get_accelerator @@ -63,7 +63,7 @@ def create_bert_model(): biases.append(torch.nn.Parameter(torch.Tensor(hidden_size))) biases.append(torch.nn.Parameter(torch.Tensor(hidden_size))) - return BertEncoderPreln(bert_config, weights, biases) + return BertEncoder(bert_config, weights, biases) class Conv1D(torch.nn.Module): diff --git a/tests/unit/ops/accelerators/test_accelerator_backward.py b/tests/unit/ops/accelerators/test_accelerator_backward.py index 43f7b471e2ae..57a6d0f9061e 100644 --- a/tests/unit/ops/accelerators/test_accelerator_backward.py +++ b/tests/unit/ops/accelerators/test_accelerator_backward.py @@ -12,10 +12,12 @@ from torch import nn from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig from deepspeed.accelerator import get_accelerator -from unit.modeling import BertConfig, BertLayerNorm, BertEncoder as BertEncoderPostln -from unit.modelingpreln import BertEncoder as BertEncoderPreln +from transformers.models.bert.configuration_bert import BertConfig +from transformers.models.bert.modeling_bert import BertEncoder from unit.common import DistributedTest, is_rocm_pytorch +BertLayerNorm = torch.nn.LayerNorm + #if not deepspeed.ops.__installed_ops__['transformer']: #pytest.skip( # "transformer kernels are temporarily disabled because of unexplained failures", @@ -194,9 +196,9 @@ def create_models(ds_config): biases[7].data.zero_() if (ds_config.pre_layer_norm): - bert_encoder = BertEncoderPreln(bert_config, weights, biases) + bert_encoder = BertEncoder(bert_config, weights, biases) else: - bert_encoder = BertEncoderPostln(bert_config, weights, biases) + bert_encoder = BertEncoder(bert_config, weights, biases) ds_encoder = DSEncoder(ds_config, weights, biases) if ds_config.fp16: diff --git a/tests/unit/ops/accelerators/test_accelerator_forward.py b/tests/unit/ops/accelerators/test_accelerator_forward.py index ee9464f63aa1..f1f9a573a545 100644 --- a/tests/unit/ops/accelerators/test_accelerator_forward.py +++ b/tests/unit/ops/accelerators/test_accelerator_forward.py @@ -9,12 +9,14 @@ import random import copy from torch import nn -from unit.modelingpreln import BertEncoder as BertEncoderPreln -from unit.modeling import BertLayerNorm, BertConfig, BertEncoder as BertEncoderPostln +from transformers.models.bert.configuration_bert import BertConfig +from transformers.models.bert.modeling_bert import BertEncoder from deepspeed import DeepSpeedTransformerLayer, DeepSpeedTransformerConfig from deepspeed.accelerator import get_accelerator from unit.common import DistributedTest +BertLayerNorm = torch.nn.LayerNorm + if torch.half not in get_accelerator().supported_dtypes(): pytest.skip(f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}", allow_module_level=True) @@ -139,9 +141,9 @@ def create_models(ds_config): biases[7].data.zero_() if (ds_config.pre_layer_norm): - bert_encoder = BertEncoderPreln(bert_config, weights, biases) + bert_encoder = BertEncoder(bert_config, weights, biases) else: - bert_encoder = BertEncoderPostln(bert_config, weights, biases) + bert_encoder = BertEncoder(bert_config, weights, biases) ds_encoder = DSEncoder(ds_config, weights, biases) if ds_config.fp16: