try to simplify checkpointing

modeling_bert.py

@@ -329,9 +329,7 @@ class BertPreTrainedModel(nn.Module):
         """
         # Instantiate model.
         model = cls(config, *inputs, **kwargs)
-        load_return = model.load_state_dict(
-            remap_state_dict(state_dict_from_pretrained(model_name), config), strict=False
-        )
+        load_return = model.load_state_dict(state_dict_from_pretrained(model_name), strict=False)
         logger.info(load_return)
         return model
 
@@ -528,247 +526,4 @@ class BertForPreTraining(BertPreTrainedModel):
             loss=total_loss,
             prediction_logits=prediction_scores,
             seq_relationship_logits=seq_relationship_score,
-        )
-
-
-def remap_state_dict(state_dict, config: PretrainedConfig):
-    """
-    Map the state_dict of a Huggingface BERT model to be flash_attn compatible.
-    """
-
-    # LayerNorm
-    def key_mapping_ln_gamma_beta(key):
-        key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key)
-        key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key)
-        return key
-
-    state_dict = OrderedDict((key_mapping_ln_gamma_beta(k), v) for k, v in state_dict.items())
-
-    # Layers
-    def key_mapping_layers(key):
-        return re.sub(r"^bert.encoder.layer.", "bert.encoder.layers.", key)
-
-    state_dict = OrderedDict((key_mapping_layers(k), v) for k, v in state_dict.items())
-
-    # LayerNorm
-    def key_mapping_ln(key):
-        key = re.sub(r"^bert.embeddings.LayerNorm.", "bert.emb_ln.", key)
-        key = re.sub(
-            r"^bert.encoder.layers.(\d+).attention.output.LayerNorm.(weight|bias)",
-            r"bert.encoder.layers.\1.norm1.\2",
-            key,
-        )
-        key = re.sub(
-            r"^bert.encoder.layers.(\d+).output.LayerNorm.(weight|bias)",
-            r"bert.encoder.layers.\1.norm2.\2",
-            key,
-        )
-        key = re.sub(
-            r"^cls.predictions.transform.LayerNorm.(weight|bias)",
-            r"cls.predictions.transform.layer_norm.\1",
-            key,
-        )
-        return key
-
-    state_dict = OrderedDict((key_mapping_ln(k), v) for k, v in state_dict.items())
-
-    # MLP
-    def key_mapping_mlp(key):
-        key = re.sub(
-            r"^bert.encoder.layers.(\d+).intermediate.dense.(weight|bias)",
-            r"bert.encoder.layers.\1.mlp.fc1.\2",
-            key,
-        )
-        key = re.sub(
-            r"^bert.encoder.layers.(\d+).output.dense.(weight|bias)",
-            r"bert.encoder.layers.\1.mlp.fc2.\2",
-            key,
-        )
-        return key
-
-    state_dict = OrderedDict((key_mapping_mlp(k), v) for k, v in state_dict.items())
-
-    # Attention
-    last_layer_subset = getattr(config, "last_layer_subset", False)
-    for d in range(config.num_hidden_layers):
-        Wq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.weight")
-        Wk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.weight")
-        Wv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.weight")
-        bq = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.query.bias")
-        bk = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.key.bias")
-        bv = state_dict.pop(f"bert.encoder.layers.{d}.attention.self.value.bias")
-        if not (last_layer_subset and d == config.num_hidden_layers - 1):
-            state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.weight"] = torch.cat(
-                [Wq, Wk, Wv], dim=0
-            )
-            state_dict[f"bert.encoder.layers.{d}.mixer.Wqkv.bias"] = torch.cat([bq, bk, bv], dim=0)
-        else:
-            state_dict[f"bert.encoder.layers.{d}.mixer.Wq.weight"] = Wq
-            state_dict[f"bert.encoder.layers.{d}.mixer.Wkv.weight"] = torch.cat([Wk, Wv], dim=0)
-            state_dict[f"bert.encoder.layers.{d}.mixer.Wq.bias"] = bq
-            state_dict[f"bert.encoder.layers.{d}.mixer.Wkv.bias"] = torch.cat([bk, bv], dim=0)
-
-    def key_mapping_attn(key):
-        return re.sub(
-            r"^bert.encoder.layers.(\d+).attention.output.dense.(weight|bias)",
-            r"bert.encoder.layers.\1.mixer.out_proj.\2",
-            key,
-        )
-
-    state_dict = OrderedDict((key_mapping_attn(k), v) for k, v in state_dict.items())
-
-    def key_mapping_decoder_bias(key):
-        return re.sub(r"^cls.predictions.bias", "cls.predictions.decoder.bias", key)
-
-    state_dict = OrderedDict((key_mapping_decoder_bias(k), v) for k, v in state_dict.items())
-
-    # Word embedding
-    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
-    if pad_vocab_size_multiple > 1:
-        word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
-        state_dict["bert.embeddings.word_embeddings.weight"] = F.pad(
-            word_embeddings, (0, 0, 0, config.vocab_size - word_embeddings.shape[0])
-        )
-        decoder_weight = state_dict["cls.predictions.decoder.weight"]
-        state_dict["cls.predictions.decoder.weight"] = F.pad(
-            decoder_weight, (0, 0, 0, config.vocab_size - decoder_weight.shape[0])
-        )
-        # If the vocab was padded, we want to set the decoder bias for those padded indices to be
-        # strongly negative (i.e. the decoder shouldn't predict those indices).
-        # TD [2022-05-09]: I don't think it affects the MLPerf training.
-        decoder_bias = state_dict["cls.predictions.decoder.bias"]
-        state_dict["cls.predictions.decoder.bias"] = F.pad(
-            decoder_bias, (0, config.vocab_size - decoder_bias.shape[0]), value=-100.0
-        )
-
-    return state_dict
-
-
-def inv_remap_state_dict(state_dict, config: PretrainedConfig):
-    """
-    Map the state_dict of a flash_attn model to be Huggingface BERT compatible.
-
-    This function is meant to be the inverse of remap_state_dict.
-    """
-    # Word embedding
-    pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
-    if pad_vocab_size_multiple > 1:
-        word_embeddings = state_dict["bert.embeddings.word_embeddings.weight"]
-        decoder_weight = state_dict["cls.predictions.decoder.weight"]
-        decoder_bias = state_dict["cls.predictions.decoder.bias"]
-        # unpad embeddings
-        state_dict["bert.embeddings.word_embeddings.weight"] = word_embeddings[
-            : config.orig_vocab_size, :
-        ]
-        state_dict["cls.predictions.decoder.weight"] = decoder_weight[: config.orig_vocab_size, :]
-        state_dict["cls.predictions.decoder.bias"] = decoder_bias[: config.orig_vocab_size]
-
-    for d in range(config.num_hidden_layers):
-        last_layer_subset = getattr(config, "last_layer_subset", False)
-        if not last_layer_subset or d != (config.num_hidden_layers - 1):
-            Wqkv_weights = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wqkv.weight")
-            Wqkv_biases = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wqkv.bias")
-            state_dict[f"bert.encoder.layers.{d}.attention.self.query.weight"] = Wqkv_weights[
-                : Wqkv_weights.shape[0] // 3, :
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.key.weight"] = Wqkv_weights[
-                Wqkv_weights.shape[0] // 3 : 2 * Wqkv_weights.shape[0] // 3, :
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.value.weight"] = Wqkv_weights[
-                2 * Wqkv_weights.shape[0] // 3 :, :
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.query.bias"] = Wqkv_biases[
-                : Wqkv_biases.shape[0] // 3
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.key.bias"] = Wqkv_biases[
-                Wqkv_biases.shape[0] // 3 : 2 * Wqkv_biases.shape[0] // 3
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.value.bias"] = Wqkv_biases[
-                2 * Wqkv_biases.shape[0] // 3 :
-            ]
-        else:
-            Wq_weight = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wq.weight")
-            Wkv_weights = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wkv.weight")
-            Wq_bias = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wq.bias")
-            Wkv_biases = state_dict.pop(f"bert.encoder.layers.{d}.mixer.Wkv.bias")
-            state_dict[f"bert.encoder.layers.{d}.attention.self.query.weight"] = Wq_weight
-            state_dict[f"bert.encoder.layers.{d}.attention.self.key.weight"] = Wkv_weights[
-                : Wkv_weights.shape[0] // 2, :
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.value.weight"] = Wkv_weights[
-                Wkv_weights.shape[0] // 2 :, :
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.query.bias"] = Wq_bias
-            state_dict[f"bert.encoder.layers.{d}.attention.self.key.bias"] = Wkv_biases[
-                : Wkv_biases.shape[0] // 2
-            ]
-            state_dict[f"bert.encoder.layers.{d}.attention.self.value.bias"] = Wkv_biases[
-                Wkv_biases.shape[0] // 2 :
-            ]
-
-    def inv_key_mapping_ln(key):
-        key = re.sub(r"bert.emb_ln.", "bert.embeddings.LayerNorm.", key)
-        key = re.sub(
-            r"bert.encoder.layers.(\d+).norm1.(weight|bias)",
-            r"bert.encoder.layers.\1.attention.output.LayerNorm.\2",
-            key,
-        )
-        key = re.sub(
-            r"bert.encoder.layers.(\d+).norm2.(weight|bias)",
-            r"bert.encoder.layers.\1.output.LayerNorm.\2",
-            key,
-        )
-        key = re.sub(
-            r"cls.predictions.transform.layer_norm.(weight|bias)",
-            r"cls.predictions.transform.LayerNorm.\1",
-            key,
-        )
-        return key
-
-    def inv_key_mapping_ln_gamma_beta(key):
-        key = re.sub(r"LayerNorm.weight$", "LayerNorm.gamma", key)
-        key = re.sub(r"LayerNorm.bias$", "LayerNorm.beta", key)
-        return key
-
-    def inv_key_mapping_layers(key):
-        return re.sub(r"bert.encoder.layers.", "bert.encoder.layer.", key)
-
-    def inv_key_mapping_mlp(key):
-        key = re.sub(
-            r"bert.encoder.layer.(\d+).mlp.fc1.(weight|bias)",
-            r"bert.encoder.layer.\1.intermediate.dense.\2",
-            key,
-        )
-        key = re.sub(
-            r"bert.encoder.layer.(\d+).mlp.fc2.(weight|bias)",
-            r"bert.encoder.layer.\1.output.dense.\2",
-            key,
-        )
-        return key
-
-    def inv_key_mapping_attn(key):
-        return re.sub(
-            r"bert.encoder.layer.(\d+).mixer.out_proj.(weight|bias)",
-            r"bert.encoder.layer.\1.attention.output.dense.\2",
-            key,
-        )
-
-    def inv_key_mapping_decoder_bias(key):
-        return re.sub(r"cls.predictions.decoder.bias", "cls.predictions.bias", key)
-
-    state_dict = OrderedDict((inv_key_mapping_ln(key), value) for key, value in state_dict.items())
-    state_dict = OrderedDict(
-        (inv_key_mapping_ln_gamma_beta(key), value) for key, value in state_dict.items()
-    )
-    state_dict = OrderedDict(
-        (inv_key_mapping_layers(key), value) for key, value in state_dict.items()
-    )
-    state_dict = OrderedDict((inv_key_mapping_mlp(key), value) for key, value in state_dict.items())
-    state_dict = OrderedDict(
-        (inv_key_mapping_attn(key), value) for key, value in state_dict.items()
-    )
-    state_dict = OrderedDict(
-        (inv_key_mapping_decoder_bias(key), value) for key, value in state_dict.items()
-    )
-
-    return state_dict
+        )