codingutf8 2023 google llc and huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis 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 convert t5x checkpoint to pytorch steps install gsutil according to https cloud google comstoragedocsgsutilinstall get a t5x checkpoint at https github comgoogleresearcht5xblobmaindocsmodels mdt511checkpoints example gsutil m cp r gs t5datapretrainedmodelst5xt511small home create or download a corresponding config for the downloaded model e g for t5 v1 1 small you can use https huggingface cogooglet5v11smallblobmainconfig json convert python3 convertt5xcheckpointtopytorch py t5xcheckpointpathhomet511small configfileconfig json pytorchdumppathhomet511smallpt returns the relative position bias parameters of a layer does not transpose return paramsfprefixprefixrelposbiasrelembedding i def t5xattentionlookupparams i prefix layernameattention returns the mlp parameters of a layer does not transpose if splitmlpwi wi0 paramsfprefixprefixmlpwi0kernel i wi1 paramsfprefixprefixmlpwi1kernel i wi wi0 wi1 else wi paramsfprefixprefixmlpwikernel i wo paramsfprefixprefixmlpwokernel i return wi wo def t5xlayernormlookupparams i prefix layername converts the parameters from t5xflax to transformerspytorch old traverseutil flattendictvariablestarget old joink v for k v in old items v1 1 models have a gated gelu with wi0 and wi1 instead of wi splitmlpwi encoderencodermlpwi0kernel in old printsplit mlp splitmlpwi new collections ordereddict shared embeddings newshared weight oldtokenembedderembedding encoder for i in rangenumlayers block i layer 0 self attention layernorm t5xlayernormlookupold i encoder preattentionlayernorm k o q v t5xattentionlookupold i encoder attention newfencoder block i layer 0 layernorm weight layernorm newfencoder block i layer 0 selfattention k weight k t newfencoder block i layer 0 selfattention o weight o t newfencoder block i layer 0 selfattention q weight q t newfencoder block i layer 0 selfattention v weight v t block i layer 1 mlp layernorm t5xlayernormlookupold i encoder premlplayernorm wi wo t5xmlplookupold i encoder splitmlpwi newfencoder block i layer 1 layernorm weight layernorm if splitmlpwi newfencoder block i layer 1 densereludense wi0 weight wi0 t newfencoder block i layer 1 densereludense wi1 weight wi1 t else newfencoder block i layer 1 densereludense wi weight wi t newfencoder block i layer 1 densereludense wo weight wo t if scalableattention convert the relembedding of each layer newfencoder block i layer 0 selfattention relativeattentionbias weight t5xrelposbiaslookup old i encoder t newencoder finallayernorm weight oldencoderencodernormscale if not scalableattention newencoder block 0 layer 0 selfattention relativeattentionbias weight t5xrelposbiaslookup old 0 encoder t newdecoder block 0 layer 0 selfattention relativeattentionbias weight t5xrelposbiaslookup old 0 decoder t if not isencoderonly decoder for i in rangenumlayers block i layer 0 self attention layernorm t5xlayernormlookupold i decoder preselfattentionlayernorm k o q v t5xattentionlookupold i decoder selfattention newfdecoder block i layer 0 layernorm weight layernorm newfdecoder block i layer 0 selfattention k weight k t newfdecoder block i layer 0 selfattention o weight o t newfdecoder block i layer 0 selfattention q weight q t newfdecoder block i layer 0 selfattention v weight v t block i layer 1 cross attention layernorm t5xlayernormlookupold i decoder precrossattentionlayernorm k o q v t5xattentionlookupold i decoder encoderdecoderattention newfdecoder block i layer 1 layernorm weight layernorm newfdecoder block i layer 1 encdecattention k weight k t newfdecoder block i layer 1 encdecattention o weight o t newfdecoder block i layer 1 encdecattention q weight q t newfdecoder block i layer 1 encdecattention v weight v t block i layer 2 mlp layernorm t5xlayernormlookupold i decoder premlplayernorm wi wo t5xmlplookupold i decoder splitmlpwi newfdecoder block i layer 2 layernorm weight layernorm if splitmlpwi newfdecoder block i layer 2 densereludense wi0 weight wi0 t newfdecoder block i layer 2 densereludense wi1 weight wi1 t else newfencoder block i layer 2 densereludense wi weight wi t newfdecoder block i layer 2 densereludense wo weight wo t if scalableattention convert the relembedding of each layer new fdecoder block i layer 0 selfattention relativeattentionbias weight t5xrelposbiaslookupold i decoder t newdecoder finallayernorm weight olddecoderdecodernormscale lm head only in v1 1 checkpoints in v1 0 embeddings are used instead if decoderlogitsdensekernel in old newlmhead weight olddecoderlogitsdensekernel t return new def makestatedictconvertedparams isencoderonly bool make a state dict with torch tensors add what is missing replaces the params in model witht the t5x converted params variables checkpoints loadt5xcheckpointt5xcheckpointpath converted convertt5xtopytorch variables numlayersconfig numlayers isencoderonlyisencoderonly scalableattentionscalableattention statedict makestatedictconverted isencoderonly model loadstatedictstatedict stricttrue def convertt5xcheckpointtopytorch t5xcheckpointpath configfile pytorchdumppath isencoderonly bool false scalableattention bool false initialise pytorch model nonv1 1 checkpoints could also use t5model but this works for all the v1 0 checkpoints will simply have an lm head that is the word embeddings load weights from tf checkpoint save pytorchmodel verify that we can load the checkpoint required parameters coding utf 8 2023 google llc and huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis 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 convert t5x checkpoint to pytorch steps install gsutil according to https cloud google com storage docs gsutil_install get a t5x checkpoint at https github com google research t5x blob main docs models md t5 11 checkpoints example gsutil m cp r gs t5 data pretrained_models t5x t5_1_1_small home create or download a corresponding config for the downloaded model e g for t5 v1 1 small you can use https huggingface co google t5 v1_1 small blob main config json convert python3 convert_t5x_checkpoint_to_pytorch py t5x_checkpoint_path home t5_1_1_small config_file config json pytorch_dump_path home t5_1_1_small_pt returns the relative position bias parameters of a layer does not transpose returns the koqv parameters of self attention does not transpose returns the mlp parameters of a layer does not transpose returns the layer norm param of a layer converts the parameters from t5x flax to transformers pytorch v1 1 models have a gated gelu with wi_0 and wi_1 instead of wi shared embeddings encoder block i layer 0 self attention block i layer 1 mlp convert the rel_embedding of each layer decoder block i layer 0 self attention block i layer 1 cross attention block i layer 2 mlp convert the rel_embedding of each layer lm head only in v1 1 checkpoints in v1 0 embeddings are used instead prepares a state dict for the pytorch model make a state dict with torch tensors add what is missing for old 1 0 models replaces the params in model witht the t5x converted params loads the config and model converts the t5x checkpoint and saves a pytorch checkpoint initialise pytorch model non v1 1 checkpoints could also use t5model but this works for all the v1 0 checkpoints will simply have an lm head that is the word embeddings load weights from tf checkpoint save pytorch model verify that we can load the checkpoint required parameters
import argparse
import collections

import numpy as np
import torch
from flax import traverse_util
from t5x import checkpoints

from transformers import MT5Config, UMT5EncoderModel, UMT5ForConditionalGeneration
from transformers.utils import logging


logging.set_verbosity_info()


def t5x_relpos_bias_lookup(params, i, prefix):
    
    return params[f"{prefix}/{prefix}/relpos_bias/rel_embedding"][:, i, :]


def t5x_attention_lookup(params, i, prefix, layer_name="attention"):
    
    k_tmp = k_tmp = np.ascontiguousarray(params[f"{prefix}/{prefix}/{layer_name}/key/kernel"][:, i, :, :])
    k = k_tmp.reshape(k_tmp.shape[0], k_tmp.shape[1] * k_tmp.shape[2])
    o_tmp = np.ascontiguousarray(params[f"{prefix}/{prefix}/{layer_name}/out/kernel"][:, i, :, :])
    o = o_tmp.reshape(o_tmp.shape[0] * o_tmp.shape[1], o_tmp.shape[2])
    q_tmp = np.ascontiguousarray(params[f"{prefix}/{prefix}/{layer_name}/query/kernel"][:, i, :, :])
    q = q_tmp.reshape(q_tmp.shape[0], q_tmp.shape[1] * q_tmp.shape[2])
    v_tmp = np.ascontiguousarray(params[f"{prefix}/{prefix}/{layer_name}/value/kernel"][:, i, :, :])
    v = v_tmp.reshape(v_tmp.shape[0], v_tmp.shape[1] * v_tmp.shape[2])
    return k, o, q, v


def t5x_mlp_lookup(params, i, prefix, split_mlp_wi=False):
    
    if split_mlp_wi:
        wi_0 = params[f"{prefix}/{prefix}/mlp/wi_0/kernel"][:, i, :]
        wi_1 = params[f"{prefix}/{prefix}/mlp/wi_1/kernel"][:, i, :]
        wi = (wi_0, wi_1)
    else:
        wi = params[f"{prefix}/{prefix}/mlp/wi/kernel"][:, i, :]

    wo = params[f"{prefix}/{prefix}/mlp/wo/kernel"][:, i, :]
    return wi, wo


def t5x_layer_norm_lookup(params, i, prefix, layer_name):
    
    return params[f"{prefix}/{prefix}/{layer_name}/scale"][:, i]


def convert_t5x_to_pytorch(
    variables: dict, *, num_layers: int, is_encoder_only: bool, scalable_attention: bool = False
):
    
    old = traverse_util.flatten_dict(variables["target"])
    old = {"/".join(k): v for k, v in old.items()}

    
    split_mlp_wi = "encoder/encoder/mlp/wi_0/kernel" in old
    print("Split MLP:", split_mlp_wi)

    new = collections.OrderedDict()

    
    new["shared.weight"] = old["token_embedder/embedding"]

    
    for i in range(num_layers):
        
        layer_norm = t5x_layer_norm_lookup(old, i, "encoder", "pre_attention_layer_norm")
        k, o, q, v = t5x_attention_lookup(old, i, "encoder", "attention")
        new[f"encoder.block.{i}.layer.0.layer_norm.weight"] = layer_norm
        new[f"encoder.block.{i}.layer.0.SelfAttention.k.weight"] = k.T
        new[f"encoder.block.{i}.layer.0.SelfAttention.o.weight"] = o.T
        new[f"encoder.block.{i}.layer.0.SelfAttention.q.weight"] = q.T
        new[f"encoder.block.{i}.layer.0.SelfAttention.v.weight"] = v.T

        
        layer_norm = t5x_layer_norm_lookup(old, i, "encoder", "pre_mlp_layer_norm")
        wi, wo = t5x_mlp_lookup(old, i, "encoder", split_mlp_wi)
        new[f"encoder.block.{i}.layer.1.layer_norm.weight"] = layer_norm
        if split_mlp_wi:
            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi_0.weight"] = wi[0].T
            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi_1.weight"] = wi[1].T
        else:
            new[f"encoder.block.{i}.layer.1.DenseReluDense.wi.weight"] = wi.T
        new[f"encoder.block.{i}.layer.1.DenseReluDense.wo.weight"] = wo.T
        if scalable_attention:
            
            new[f"encoder.block.{i}.layer.0.SelfAttention.relative_attention_bias.weight"] = t5x_relpos_bias_lookup(
                old, i, "encoder"
            ).T

    new["encoder.final_layer_norm.weight"] = old["encoder/encoder_norm/scale"]

    if not scalable_attention:
        new["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = t5x_relpos_bias_lookup(
            old, 0, "encoder"
        ).T
        new["decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"] = t5x_relpos_bias_lookup(
            old, 0, "decoder"
        ).T

    if not is_encoder_only:
        
        for i in range(num_layers):
            
            layer_norm = t5x_layer_norm_lookup(old, i, "decoder", "pre_self_attention_layer_norm")
            k, o, q, v = t5x_attention_lookup(old, i, "decoder", "self_attention")
            new[f"decoder.block.{i}.layer.0.layer_norm.weight"] = layer_norm
            new[f"decoder.block.{i}.layer.0.SelfAttention.k.weight"] = k.T
            new[f"decoder.block.{i}.layer.0.SelfAttention.o.weight"] = o.T
            new[f"decoder.block.{i}.layer.0.SelfAttention.q.weight"] = q.T
            new[f"decoder.block.{i}.layer.0.SelfAttention.v.weight"] = v.T

            
            layer_norm = t5x_layer_norm_lookup(old, i, "decoder", "pre_cross_attention_layer_norm")
            k, o, q, v = t5x_attention_lookup(old, i, "decoder", "encoder_decoder_attention")
            new[f"decoder.block.{i}.layer.1.layer_norm.weight"] = layer_norm
            new[f"decoder.block.{i}.layer.1.EncDecAttention.k.weight"] = k.T
            new[f"decoder.block.{i}.layer.1.EncDecAttention.o.weight"] = o.T
            new[f"decoder.block.{i}.layer.1.EncDecAttention.q.weight"] = q.T
            new[f"decoder.block.{i}.layer.1.EncDecAttention.v.weight"] = v.T

            
            layer_norm = t5x_layer_norm_lookup(old, i, "decoder", "pre_mlp_layer_norm")
            wi, wo = t5x_mlp_lookup(old, i, "decoder", split_mlp_wi)
            new[f"decoder.block.{i}.layer.2.layer_norm.weight"] = layer_norm
            if split_mlp_wi:
                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi_0.weight"] = wi[0].T
                new[f"decoder.block.{i}.layer.2.DenseReluDense.wi_1.weight"] = wi[1].T
            else:
                new[f"encoder.block.{i}.layer.2.DenseReluDense.wi.weight"] = wi.T
            new[f"decoder.block.{i}.layer.2.DenseReluDense.wo.weight"] = wo.T

            if scalable_attention:
                
                new[
                    f"decoder.block.{i}.layer.0.SelfAttention.relative_attention_bias.weight"
                ] = t5x_relpos_bias_lookup(old, i, "decoder").T

        new["decoder.final_layer_norm.weight"] = old["decoder/decoder_norm/scale"]

        
        if "decoder/logits_dense/kernel" in old:
            new["lm_head.weight"] = old["decoder/logits_dense/kernel"].T

    return new


def make_state_dict(converted_params, is_encoder_only: bool):
    
    
    state_dict = collections.OrderedDict([(k, torch.from_numpy(v.copy())) for (k, v) in converted_params.items()])

    
    if "encoder.embed_tokens.weight" not in state_dict:
        state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"]

    if not is_encoder_only:
        if "decoder.embed_tokens.weight" not in state_dict:
            state_dict["decoder.embed_tokens.weight"] = state_dict["shared.weight"]

        if "lm_head.weight" not in state_dict:  
            print("Using shared word embeddings as lm_head.")
            state_dict["lm_head.weight"] = state_dict["shared.weight"]

    return state_dict


def load_t5x_weights_in_t5(model, config, t5x_checkpoint_path, is_encoder_only, scalable_attention):
    
    variables = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)
    converted = convert_t5x_to_pytorch(
        variables, num_layers=config.num_layers, is_encoder_only=is_encoder_only, scalable_attention=scalable_attention
    )
    state_dict = make_state_dict(converted, is_encoder_only)
    model.load_state_dict(state_dict, strict=True)


def convert_t5x_checkpoint_to_pytorch(
    t5x_checkpoint_path,
    config_file,
    pytorch_dump_path,
    is_encoder_only: bool = False,
    scalable_attention: bool = False,
):
    
    
    config = MT5Config.from_json_file(config_file)
    print(f"Building PyTorch model from configuration: {config}")
    
    
    if is_encoder_only:
        model = UMT5EncoderModel(config)
    else:
        model = UMT5ForConditionalGeneration(config)

    
    load_t5x_weights_in_t5(model, config, t5x_checkpoint_path, is_encoder_only, scalable_attention)

    
    print(f"Save PyTorch model to {pytorch_dump_path}")
    model.save_pretrained(pytorch_dump_path)

    
    model.from_pretrained(pytorch_dump_path)
    print("Done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Converts a native T5X checkpoint into a PyTorch checkpoint.")
    
    parser.add_argument(
        "--t5x_checkpoint_path", default=None, type=str, required=True, help="Path to the T5X checkpoint."
    )
    parser.add_argument(
        "--config_file",
        default=None,
        type=str,
        required=True,
        help="The config json file corresponding to the pre-trained T5 model.\nThis specifies the model architecture.",
    )
    parser.add_argument(
        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--is_encoder_only", action="store_true", help="Check if the model is encoder-decoder model", default=False
    )
    parser.add_argument(
        "--scalable_attention",
        action="store_true",
        help="Whether the model uses scaled attention (umt5 model)",
        default=False,
    )
    args = parser.parse_args()
    convert_t5x_checkpoint_to_pytorch(
        args.t5x_checkpoint_path,
        args.config_file,
        args.pytorch_dump_path,
        args.is_encoder_only,
        args.scalable_attention,
    )