load model from drive and convert

README.md

@@ -0,0 +1,3 @@
+# README - long-t5-tglobal-base-16384-booksum-V11-big_patent-V2 
+- this README was added because there wasn't one
+- created 2022-07-31_12-14-50

config.json

@@ -0,0 +1,43 @@
+{
+  "_name_or_path": "pszemraj/long-t5-tglobal-base-16384-booksum-V11-big_patent-V1",
+  "architectures": [
+    "LongT5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 768,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "gelu_new",
+  "dropout_rate": 0.1,
+  "early_stopping": true,
+  "encoder_attention_type": "transient-global",
+  "encoder_no_repeat_ngram_size": 4,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "global_block_size": 16,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "length_penalty": 0.8,
+  "local_radius": 127,
+  "max_length": 512,
+  "min_length": 8,
+  "model_type": "longt5",
+  "n_positions": 4096,
+  "no_repeat_ngram_size": 3,
+  "num_beams": 2,
+  "num_decoder_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "repetition_penalty": 3.5,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.21.0",
+  "use_cache": false,
+  "vocab_size": 32128
+}

latest

@@ -0,0 +1 @@
+global_step325

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1620415284b0de710de1a9af0ea4045a771f6b5f27a79294dddee9b09b068ad2
+size 990388907

rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcbda4b624bb1ed1b6c0fa22fb8bb1555261527dfbe1629fef4820faae0bffa4
+size 14503

special_tokens_map.json

@@ -0,0 +1,107 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

tokenizer_config.json

@@ -0,0 +1,111 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "name_or_path": "pszemraj/long-t5-tglobal-base-16384-booksum-V11-big_patent-V1",
+  "pad_token": "<pad>",
+  "special_tokens_map_file": null,
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

trainer_state.json

@@ -0,0 +1,988 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.15840619002650258,
+  "global_step": 325,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0,
+      "learning_rate": 5.7142857142857135e-05,
+      "loss": 1.9649,
+      "step": 2
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 0.00011428571428571427,
+      "loss": 2.0221,
+      "step": 4
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 0.0001714285714285714,
+      "loss": 1.9956,
+      "step": 6
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 0.00022857142857142854,
+      "loss": 1.9373,
+      "step": 8
+    },
+    {
+      "epoch": 0.0,
+      "learning_rate": 0.0002857142857142857,
+      "loss": 2.0042,
+      "step": 10
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0003428571428571428,
+      "loss": 1.9593,
+      "step": 12
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.00039999999999999996,
+      "loss": 2.031,
+      "step": 14
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0004571428571428571,
+      "loss": 1.9644,
+      "step": 16
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0005142857142857142,
+      "loss": 1.9968,
+      "step": 18
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0005714285714285714,
+      "loss": 1.9694,
+      "step": 20
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0005999996407482917,
+      "loss": 1.9885,
+      "step": 22
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0005999967667397879,
+      "loss": 1.9295,
+      "step": 24
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0005999910187503132,
+      "loss": 1.9646,
+      "step": 26
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0005999823968349338,
+      "loss": 1.9615,
+      "step": 28
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.000599970901076248,
+      "loss": 1.9456,
+      "step": 30
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005999565315843857,
+      "loss": 2.014,
+      "step": 32
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005999392884970068,
+      "loss": 1.9118,
+      "step": 34
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005999191719793011,
+      "loss": 1.944,
+      "step": 36
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005998961822239856,
+      "loss": 1.9475,
+      "step": 38
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.000599870319451303,
+      "loss": 2.038,
+      "step": 40
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005998415839090198,
+      "loss": 1.9438,
+      "step": 42
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005998099758724235,
+      "loss": 2.0804,
+      "step": 44
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005997754956443205,
+      "loss": 1.9767,
+      "step": 46
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005997381435550326,
+      "loss": 1.9322,
+      "step": 48
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0005996979199623944,
+      "loss": 1.9276,
+      "step": 50
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005996548252517495,
+      "loss": 1.9933,
+      "step": 52
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005996088598359469,
+      "loss": 1.8901,
+      "step": 54
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005995600241553371,
+      "loss": 1.9472,
+      "step": 56
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.000599508318677768,
+      "loss": 1.9855,
+      "step": 58
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.00059945374389858,
+      "loss": 1.9887,
+      "step": 60
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005993963003406018,
+      "loss": 1.9798,
+      "step": 62
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005993359885541448,
+      "loss": 1.9956,
+      "step": 64
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005992728091169984,
+      "loss": 1.9411,
+      "step": 66
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.0005992067626344242,
+      "loss": 1.9722,
+      "step": 68
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 0.00059913784973915,
+      "loss": 1.97,
+      "step": 70
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005990660710913641,
+      "loss": 1.9612,
+      "step": 72
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005989914273787089,
+      "loss": 1.8503,
+      "step": 74
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005989139193162741,
+      "loss": 1.992,
+      "step": 76
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.00059883354764659,
+      "loss": 1.9675,
+      "step": 78
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005987503131396204,
+      "loss": 1.9609,
+      "step": 80
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005986642165927551,
+      "loss": 2.0349,
+      "step": 82
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005985752588308026,
+      "loss": 1.9824,
+      "step": 84
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005984834407059817,
+      "loss": 1.9017,
+      "step": 86
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005983887630979137,
+      "loss": 1.903,
+      "step": 88
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.000598291226913614,
+      "loss": 1.9067,
+      "step": 90
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.000598190833087483,
+      "loss": 1.941,
+      "step": 92
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005980875825812974,
+      "loss": 1.9856,
+      "step": 94
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005979814763842014,
+      "loss": 1.9555,
+      "step": 96
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005978725155126967,
+      "loss": 1.9408,
+      "step": 98
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005977607010106324,
+      "loss": 2.0131,
+      "step": 100
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005976460339491963,
+      "loss": 1.9499,
+      "step": 102
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.000597528515426903,
+      "loss": 1.9381,
+      "step": 104
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005974081465695849,
+      "loss": 1.9805,
+      "step": 106
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005972849285303804,
+      "loss": 1.8787,
+      "step": 108
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005971588624897232,
+      "loss": 1.8912,
+      "step": 110
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005970299496553309,
+      "loss": 1.9536,
+      "step": 112
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005968981912621937,
+      "loss": 1.9388,
+      "step": 114
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005967635885725623,
+      "loss": 2.0041,
+      "step": 116
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005966261428759357,
+      "loss": 1.9447,
+      "step": 118
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005964858554890492,
+      "loss": 2.0031,
+      "step": 120
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005963427277558616,
+      "loss": 1.9063,
+      "step": 122
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005961967610475422,
+      "loss": 1.9492,
+      "step": 124
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005960479567624578,
+      "loss": 1.9956,
+      "step": 126
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005958963163261595,
+      "loss": 1.9329,
+      "step": 128
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005957418411913688,
+      "loss": 1.9424,
+      "step": 130
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005955845328379636,
+      "loss": 1.9105,
+      "step": 132
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.000595424392772964,
+      "loss": 1.9439,
+      "step": 134
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005952614225305184,
+      "loss": 1.9586,
+      "step": 136
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005950956236718882,
+      "loss": 1.8851,
+      "step": 138
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005949269977854329,
+      "loss": 1.9031,
+      "step": 140
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005947555464865954,
+      "loss": 1.9294,
+      "step": 142
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.000594581271417886,
+      "loss": 1.9779,
+      "step": 144
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005944041742488665,
+      "loss": 1.9515,
+      "step": 146
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005942242566761351,
+      "loss": 1.9249,
+      "step": 148
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005940415204233092,
+      "loss": 1.9104,
+      "step": 150
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005938559672410093,
+      "loss": 1.9548,
+      "step": 152
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005936675989068425,
+      "loss": 1.9314,
+      "step": 154
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005934764172253849,
+      "loss": 1.9468,
+      "step": 156
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005932824240281645,
+      "loss": 1.9821,
+      "step": 158
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005930856211736438,
+      "loss": 1.9609,
+      "step": 160
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005928860105472022,
+      "loss": 1.9261,
+      "step": 162
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005926835940611172,
+      "loss": 1.9594,
+      "step": 164
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.000592478373654547,
+      "loss": 1.8914,
+      "step": 166
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005922703512935113,
+      "loss": 1.9509,
+      "step": 168
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005920595289708723,
+      "loss": 1.9988,
+      "step": 170
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005918459087063165,
+      "loss": 1.9886,
+      "step": 172
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005916294925463346,
+      "loss": 2.0024,
+      "step": 174
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005914102825642018,
+      "loss": 1.859,
+      "step": 176
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005911882808599586,
+      "loss": 1.9439,
+      "step": 178
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005909634895603902,
+      "loss": 1.9823,
+      "step": 180
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.000590735910819006,
+      "loss": 1.9308,
+      "step": 182
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005905055468160197,
+      "loss": 1.9459,
+      "step": 184
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005902723997583274,
+      "loss": 1.9146,
+      "step": 186
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005900364718794873,
+      "loss": 1.9036,
+      "step": 188
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005897977654396977,
+      "loss": 1.9035,
+      "step": 190
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.000589556282725776,
+      "loss": 1.9607,
+      "step": 192
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 0.0005893120260511362,
+      "loss": 2.0468,
+      "step": 194
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005890649977557668,
+      "loss": 1.9687,
+      "step": 196
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005888152002062089,
+      "loss": 1.9958,
+      "step": 198
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005885626357955329,
+      "loss": 1.9025,
+      "step": 200
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005883073069433159,
+      "loss": 1.9077,
+      "step": 202
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005880492160956185,
+      "loss": 1.9494,
+      "step": 204
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005877883657249612,
+      "loss": 1.8716,
+      "step": 206
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.000587524758330301,
+      "loss": 1.9075,
+      "step": 208
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005872583964370073,
+      "loss": 1.9406,
+      "step": 210
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005869892825968375,
+      "loss": 1.9179,
+      "step": 212
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005867174193879131,
+      "loss": 1.9702,
+      "step": 214
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005864428094146943,
+      "loss": 1.9297,
+      "step": 216
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005861654553079557,
+      "loss": 1.8467,
+      "step": 218
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005858853597247606,
+      "loss": 1.9145,
+      "step": 220
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005856025253484358,
+      "loss": 1.944,
+      "step": 222
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005853169548885461,
+      "loss": 1.9321,
+      "step": 224
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005850286510808675,
+      "loss": 1.9838,
+      "step": 226
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005847376166873624,
+      "loss": 1.9891,
+      "step": 228
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005844438544961515,
+      "loss": 1.9384,
+      "step": 230
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005841473673214886,
+      "loss": 1.8826,
+      "step": 232
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005838481580037324,
+      "loss": 1.8983,
+      "step": 234
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005835462294093202,
+      "loss": 1.8804,
+      "step": 236
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.00058324158443074,
+      "loss": 1.8997,
+      "step": 238
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005829342259865026,
+      "loss": 1.9478,
+      "step": 240
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005826241570211144,
+      "loss": 1.9727,
+      "step": 242
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005823113805050482,
+      "loss": 1.9216,
+      "step": 244
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005819958994347157,
+      "loss": 1.9208,
+      "step": 246
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.000581677716832438,
+      "loss": 1.9201,
+      "step": 248
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005813568357464172,
+      "loss": 1.869,
+      "step": 250
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005810332592507066,
+      "loss": 1.9111,
+      "step": 252
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005807069904451822,
+      "loss": 1.8696,
+      "step": 254
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005803780324555121,
+      "loss": 1.8946,
+      "step": 256
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005800463884331269,
+      "loss": 1.9641,
+      "step": 258
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005797120615551896,
+      "loss": 1.8923,
+      "step": 260
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005793750550245648,
+      "loss": 1.8612,
+      "step": 262
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005790353720697887,
+      "loss": 1.927,
+      "step": 264
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005786930159450374,
+      "loss": 1.9709,
+      "step": 266
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005783479899300962,
+      "loss": 1.9665,
+      "step": 268
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005780002973303283,
+      "loss": 1.8657,
+      "step": 270
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005776499414766424,
+      "loss": 2.0055,
+      "step": 272
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005772969257254615,
+      "loss": 1.9147,
+      "step": 274
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005769412534586908,
+      "loss": 1.9383,
+      "step": 276
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005765829280836846,
+      "loss": 1.9575,
+      "step": 278
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005762219530332142,
+      "loss": 1.9192,
+      "step": 280
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005758583317654352,
+      "loss": 1.8842,
+      "step": 282
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005754920677638535,
+      "loss": 1.9905,
+      "step": 284
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.000575123164537293,
+      "loss": 1.9686,
+      "step": 286
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005747516256198616,
+      "loss": 2.0003,
+      "step": 288
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005743774545709163,
+      "loss": 1.9195,
+      "step": 290
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.000574000654975031,
+      "loss": 1.8899,
+      "step": 292
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.0005736212304419609,
+      "loss": 1.9143,
+      "step": 294
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 0.000573239184606608,
+      "loss": 1.8431,
+      "step": 296
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005728545211289866,
+      "loss": 1.8978,
+      "step": 298
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005724672436941882,
+      "loss": 1.9017,
+      "step": 300
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005720773560123461,
+      "loss": 1.8912,
+      "step": 302
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005716848618185996,
+      "loss": 1.9412,
+      "step": 304
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.000571289764873059,
+      "loss": 1.8843,
+      "step": 306
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005708920689607684,
+      "loss": 1.8971,
+      "step": 308
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005704917778916709,
+      "loss": 1.9243,
+      "step": 310
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005700888955005706,
+      "loss": 1.8342,
+      "step": 312
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.000569683425647097,
+      "loss": 1.8725,
+      "step": 314
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005692753722156673,
+      "loss": 1.8597,
+      "step": 316
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005688647391154496,
+      "loss": 1.954,
+      "step": 318
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 0.0005684515302803256,
+      "loss": 1.9454,
+      "step": 320
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 0.000568035749668852,
+      "loss": 1.9336,
+      "step": 322
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 0.000567617401264224,
+      "loss": 1.96,
+      "step": 324
+    }
+  ],
+  "max_steps": 2051,
+  "num_train_epochs": 1,
+  "total_flos": 4.557933379584e+17,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5163391d1b5a17f8e9a16e3191f7ee53c40429fe20ca10e4a048ce91e34dc316
+size 4591

zero_to_fp32.py

@@ -0,0 +1,484 @@
+#!/usr/bin/env python
+
+# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+import deepspeed
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION,
+                                            OPTIMIZER_STATE_DICT,
+                                            PARAM_SHAPES,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS,
+                                            ZERO_STAGE,
+                                            PARTITION_COUNT,
+                                            PARAM_SHAPES,
+                                            BUFFER_NAMES)
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage == 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_optim_files(checkpoint_dir):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    optim_files = sorted(glob.glob(os.path.join(checkpoint_dir,
+                                                "*_optim_states.pt")),
+                         key=natural_keys)
+
+    if len(optim_files) == 0:
+        raise FileNotFoundError(
+            f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")
+
+    return optim_files
+
+
+def parse_model_state(file):
+    state_dict = torch.load(file, map_location=device)
+
+    if BUFFER_NAMES not in state_dict:
+        raise ValueError(f"{file} is not a model state checkpoint")
+    buffer_names = state_dict[BUFFER_NAMES]
+    if debug:
+        print("Found buffers:", buffer_names)
+
+    # recover just the buffers while restoring them to fp32 if they were saved in fp16
+    buffers = {
+        k: v.float()
+        for k,
+        v in state_dict["module"].items() if k in buffer_names
+    }
+    param_shapes = state_dict[PARAM_SHAPES]
+
+    ds_version = state_dict.get(DS_VERSION, None)
+
+    return buffers, param_shapes, ds_version
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dicts.append(torch.load(f, map_location=device))
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage == 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage == 2:
+        fp32_flat_groups = [
+            state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key]
+            for i in range(len(state_dicts))
+        ]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key],
+                      0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(
+        f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)
+    buffers, param_shapes, ds_version = parse_model_state(model_file)
+    print(f'Parsing checkpoint created by deepspeed=={ds_version}')
+
+    if zero_stage == 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(
+                    f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum([
+        full_single_fp32_vector.numel()
+        for full_single_fp32_vector in merged_single_partition_of_fp32_groups
+    ])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum(
+            [sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(
+                    f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} "
+                )
+            state_dict[name] = full_single_fp32_vector.narrow(
+                0,
+                offset,
+                unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(
+                f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0,
+                                             offset,
+                                             partitioned_numel)
+                  for i in range(world_size)),
+            0).narrow(0,
+                      0,
+                      unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(
+            f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "checkpoint_dir",
+        type=str,
+        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help=
+        "path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)"
+    )
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)