Fix version of Transformers fused_bert.py (#2157)

* Fix version of Transformers fused_bert.py

* Update fused_bert.py

* Update fused_bert.py

Author: jianan-gu

intel_extension_for_pytorch/cpu/tpp/fused_bert.py

@@ -703,7 +703,7 @@ def backward(ctx, *grad_outs):
 class BertEmbeddings(BlockedModule):
     """Construct the embeddings from word, position and token_type embeddings."""
 
-    def __init__(self, config):
+    def __init__(self, config, position_ids_persistent=False):
         super().__init__()
         self.word_embeddings = nn.Embedding(
             config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
@@ -725,11 +725,14 @@ def __init__(self, config):
         self.pad_token_id = config.pad_token_id
 
         # position_ids (1, len position emb) is contiguous in memory and exported when serialized
-        self.register_buffer(
-            "position_ids",
-            torch.arange(config.max_position_embeddings).expand((1, -1)),
-            persistent=False,
-        )
+        if not position_ids_persistent:
+            self.register_buffer(
+                "position_ids",
+                torch.arange(config.max_position_embeddings).expand((1, -1)),
+                persistent=False,
+            )
+        else:
+            self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
         self.position_embedding_type = getattr(
             config, "position_embedding_type", "absolute"
         )
@@ -1243,7 +1246,7 @@ def fast_bert(model, dtype=torch.float, optimizer=None, unpad=False):
     # tpp bert optimization depends on the transformers repo to implementate the related module
     installed_pkg = {pkg.key for pkg in pkg_resources.working_set}
     min_version = "4.6.0"
-    max_version = "4.20.0"
+    max_version = "4.31.0"
     if "transformers" not in installed_pkg:
         raise RuntimeError(
             "Please installed the transformers with version: between {} and {}".format(
@@ -1263,6 +1266,9 @@ def fast_bert(model, dtype=torch.float, optimizer=None, unpad=False):
                 min_version, max_version, trans_version
             )
         )
+    position_ids_persistent = False
+    if version.parse(trans_version) < version.parse("4.31.0"):
+        position_ids_persistent = True
     PT_OPTIMIZER_TO_TPP_OPTIMIZER = {
         torch.optim.AdamW: AdamW,
         transformers.optimization.AdamW: AdamW,
@@ -1297,7 +1303,7 @@ def fast_bert(model, dtype=torch.float, optimizer=None, unpad=False):
         assert isinstance(
             new_model.embeddings, transformers.models.bert.modeling_bert.BertEmbeddings
         )
-        new_model.embeddings = BertEmbeddings(model.config)
+        new_model.embeddings = BertEmbeddings(model.config, position_ids_persistent=position_ids_persistent)
         assert isinstance(
             new_model.encoder, transformers.models.bert.modeling_bert.BertEncoder
         )
@@ -1309,7 +1315,7 @@ def fast_bert(model, dtype=torch.float, optimizer=None, unpad=False):
             new_model.bert.embeddings,
             transformers.models.bert.modeling_bert.BertEmbeddings,
         )
-        new_model.bert.embeddings = BertEmbeddings(model.bert.config)
+        new_model.bert.embeddings = BertEmbeddings(model.bert.config, position_ids_persistent=position_ids_persistent)
         assert isinstance(
             new_model.bert.encoder, transformers.models.bert.modeling_bert.BertEncoder
         )