[add] model

config.json

@@ -0,0 +1,348 @@
+{
+  "_commit_hash": null,
+  "_name_or_path": "/mnt/blob/projects/sca-xiaoke-v3/amlt-results/7301932201.25563-cd1e6021-6ea9-4835-8578-ba26f723a708/checkpoint-100000",
+  "architectures": [
+    "ScaMultitaskV2Model"
+  ],
+  "cache_dir": "/mnt/blob/weights/.model.cache/",
+  "initializer_range": 0.02,
+  "mask_caption_decoder_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "additional_num_hidden_layers": 12,
+    "architectures": null,
+    "attention_downsample_rate": 2,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "relu",
+    "hidden_size": 256,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "iou_head_depth": 3,
+    "iou_head_hidden_dim": 256,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mlp_dim": 2048,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 8,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_caption_heads": 1,
+    "num_caption_tokens": 8,
+    "num_hidden_layers": 2,
+    "num_multimask_outputs": 3,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "model_type": "sca",
+  "num_task_tokens": 6,
+  "prompt_encoder_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_size": 256,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_embedding_size": 64,
+    "image_size": 1024,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "mask_input_channels": 16,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_point_embeddings": 4,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "text_config": {
+    "_name_or_path": "gpt2-large",
+    "activation_function": "gelu_new",
+    "add_cross_attention": false,
+    "architectures": [
+      "GPT2LMHeadModel"
+    ],
+    "attn_pdrop": 0.1,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 50256,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "embd_pdrop": 0.1,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 50256,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_epsilon": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "gpt2",
+    "n_ctx": 1024,
+    "n_embd": 1280,
+    "n_head": 20,
+    "n_inner": null,
+    "n_layer": 36,
+    "n_positions": 1024,
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "reorder_and_upcast_attn": false,
+    "repetition_penalty": 1.0,
+    "resid_pdrop": 0.1,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "scale_attn_by_inverse_layer_idx": false,
+    "scale_attn_weights": true,
+    "sep_token_id": null,
+    "summary_activation": null,
+    "summary_first_dropout": 0.1,
+    "summary_proj_to_labels": true,
+    "summary_type": "cls_index",
+    "summary_use_proj": true,
+    "suppress_tokens": null,
+    "task_specific_params": {
+      "text-generation": {
+        "do_sample": true,
+        "max_length": 50
+      }
+    },
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 50257
+  },
+  "torch_dtype": "float16",
+  "transformers_version": null,
+  "use_decoder_only_language_model": true,
+  "vision_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "global_attn_indexes": [
+      7,
+      15,
+      23,
+      31
+    ],
+    "hidden_act": "gelu",
+    "hidden_size": 1280,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 1024,
+    "initializer_factor": 1.0,
+    "initializer_range": 1e-10,
+    "intermediate_size": 6144,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "mlp_dim": 5120,
+    "mlp_ratio": 4.0,
+    "model_type": "",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 32,
+    "num_pos_feats": 128,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_channels": 256,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "prefix": null,
+    "problem_type": null,
+    "projection_dim": 512,
+    "pruned_heads": {},
+    "qkv_bias": true,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.30.2",
+    "typical_p": 1.0,
+    "use_abs_pos": true,
+    "use_bfloat16": false,
+    "use_rel_pos": true,
+    "window_size": 14
+  }
+}

latest

@@ -0,0 +1 @@
+global_step100000

pytorch_model.bin

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

scheduler.pt

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

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "add_prefix_space": false,
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 20,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

trainer_state.json

@@ -0,0 +1,1114 @@
+{
+  "best_metric": 3.5546064376831055,
+  "best_model_checkpoint": "//amltff6de91c1707b018e950147da959e142/projects/sca-xiaoke-v3/amlt-results/7300230113.98638-a536a63f-7921-4c0f-a350-97817c95c30d/checkpoint-90000",
+  "epoch": 82.64462809917356,
+  "global_step": 100000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "_prepare_inputs_in_ms": 27.127794979605824,
+      "compute_loss_in_ms": 1811.1886270344257,
+      "epoch": 0.0,
+      "learning_rate/full": 0.0,
+      "loss": 7.5531,
+      "step": 1,
+      "training_step_in_ms": 2240.9718980197795
+    },
+    {
+      "epoch": 0.0,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 7.401121139526367,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.029569745772319803,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 83.1742,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 9.618,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.301,
+      "step": 1
+    },
+    {
+      "_prepare_inputs_in_ms": 9.555798228632284,
+      "compute_loss_in_ms": 536.5458341809833,
+      "epoch": 0.83,
+      "learning_rate/full": 0.0003999392508508148,
+      "loss": 4.3937,
+      "step": 1000,
+      "training_step_in_ms": 946.2872614599979
+    },
+    {
+      "_prepare_inputs_in_ms": 9.661823262518737,
+      "compute_loss_in_ms": 537.3916240136023,
+      "epoch": 1.65,
+      "learning_rate/full": 0.00039968541343111914,
+      "loss": 4.0494,
+      "step": 2000,
+      "training_step_in_ms": 955.0557815724169
+    },
+    {
+      "_prepare_inputs_in_ms": 9.627924977685325,
+      "compute_loss_in_ms": 537.9465775624267,
+      "epoch": 2.48,
+      "learning_rate/full": 0.0003992346658781257,
+      "loss": 3.9791,
+      "step": 3000,
+      "training_step_in_ms": 958.343352012278
+    },
+    {
+      "_prepare_inputs_in_ms": 9.754187489685137,
+      "compute_loss_in_ms": 538.0461517369258,
+      "epoch": 3.31,
+      "learning_rate/full": 0.0003985858034191765,
+      "loss": 3.9371,
+      "step": 4000,
+      "training_step_in_ms": 961.1277512954548
+    },
+    {
+      "_prepare_inputs_in_ms": 9.712934053619392,
+      "compute_loss_in_ms": 537.9799637984834,
+      "epoch": 4.13,
+      "learning_rate/full": 0.0003977411181848781,
+      "loss": 3.9065,
+      "step": 5000,
+      "training_step_in_ms": 954.7033457334619
+    },
+    {
+      "epoch": 4.13,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.763737201690674,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.23929773235608354,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 102.1516,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.831,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.245,
+      "step": 5000
+    },
+    {
+      "_prepare_inputs_in_ms": 9.484028106941501,
+      "compute_loss_in_ms": 535.8051792999613,
+      "epoch": 4.96,
+      "learning_rate/full": 0.0003966997561999938,
+      "loss": 3.8844,
+      "step": 6000,
+      "training_step_in_ms": 946.6810292407754
+    },
+    {
+      "_prepare_inputs_in_ms": 9.555567614093889,
+      "compute_loss_in_ms": 537.8309749438195,
+      "epoch": 5.79,
+      "learning_rate/full": 0.000395466162801616,
+      "loss": 3.8646,
+      "step": 7000,
+      "training_step_in_ms": 956.3491519849049
+    },
+    {
+      "_prepare_inputs_in_ms": 9.571906232624315,
+      "compute_loss_in_ms": 535.8670682059601,
+      "epoch": 6.61,
+      "learning_rate/full": 0.0003940366164283571,
+      "loss": 3.8482,
+      "step": 8000,
+      "training_step_in_ms": 955.7397421186324
+    },
+    {
+      "_prepare_inputs_in_ms": 9.494116805610247,
+      "compute_loss_in_ms": 535.7246166499099,
+      "epoch": 7.44,
+      "learning_rate/full": 0.0003924165284571268,
+      "loss": 3.8329,
+      "step": 9000,
+      "training_step_in_ms": 956.2797198233311
+    },
+    {
+      "_prepare_inputs_in_ms": 9.442963820649311,
+      "compute_loss_in_ms": 535.982440788066,
+      "epoch": 8.26,
+      "learning_rate/full": 0.00039060426098193074,
+      "loss": 3.8234,
+      "step": 10000,
+      "training_step_in_ms": 960.6822794212494
+    },
+    {
+      "epoch": 8.26,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.69438099861145,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.24724186480974675,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 99.2464,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 8.061,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.252,
+      "step": 10000
+    },
+    {
+      "_prepare_inputs_in_ms": 9.298451105870942,
+      "compute_loss_in_ms": 533.8574210923398,
+      "epoch": 9.09,
+      "learning_rate/full": 0.00038860732494287416,
+      "loss": 3.8075,
+      "step": 11000,
+      "training_step_in_ms": 953.7077538485173
+    },
+    {
+      "_prepare_inputs_in_ms": 9.43831800215412,
+      "compute_loss_in_ms": 534.8986968573299,
+      "epoch": 9.92,
+      "learning_rate/full": 0.0003864196958913291,
+      "loss": 3.8011,
+      "step": 12000,
+      "training_step_in_ms": 952.7247995454236
+    },
+    {
+      "_prepare_inputs_in_ms": 9.553659793222323,
+      "compute_loss_in_ms": 535.9662136517582,
+      "epoch": 10.74,
+      "learning_rate/full": 0.0003840498188962311,
+      "loss": 3.7873,
+      "step": 13000,
+      "training_step_in_ms": 957.1981301900814
+    },
+    {
+      "_prepare_inputs_in_ms": 9.344144803879317,
+      "compute_loss_in_ms": 535.5388605375774,
+      "epoch": 11.57,
+      "learning_rate/full": 0.0003814979428002389,
+      "loss": 3.7786,
+      "step": 14000,
+      "training_step_in_ms": 959.83031323459
+    },
+    {
+      "_prepare_inputs_in_ms": 9.39019125822233,
+      "compute_loss_in_ms": 536.7152543000993,
+      "epoch": 12.4,
+      "learning_rate/full": 0.00037876659104596076,
+      "loss": 3.7703,
+      "step": 15000,
+      "training_step_in_ms": 961.1087728133425
+    },
+    {
+      "epoch": 12.4,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.6609387397766113,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2538100106775413,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 103.2498,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.748,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.242,
+      "step": 15000
+    },
+    {
+      "_prepare_inputs_in_ms": 9.56299589696999,
+      "compute_loss_in_ms": 535.2517742190976,
+      "epoch": 13.22,
+      "learning_rate/full": 0.00037585846455191155,
+      "loss": 3.7622,
+      "step": 16000,
+      "training_step_in_ms": 952.1916261643055
+    },
+    {
+      "_prepare_inputs_in_ms": 9.559556159831118,
+      "compute_loss_in_ms": 535.5560795084457,
+      "epoch": 14.05,
+      "learning_rate/full": 0.00037277643904168816,
+      "loss": 3.7606,
+      "step": 17000,
+      "training_step_in_ms": 957.2516361214803
+    },
+    {
+      "_prepare_inputs_in_ms": 9.455631462449674,
+      "compute_loss_in_ms": 535.3892399296165,
+      "epoch": 14.88,
+      "learning_rate/full": 0.0003695235622002878,
+      "loss": 3.749,
+      "step": 18000,
+      "training_step_in_ms": 953.2538345798966
+    },
+    {
+      "_prepare_inputs_in_ms": 9.365507787151728,
+      "compute_loss_in_ms": 535.8022946891142,
+      "epoch": 15.7,
+      "learning_rate/full": 0.00036609954385629034,
+      "loss": 3.7464,
+      "step": 19000,
+      "training_step_in_ms": 959.837038420781
+    },
+    {
+      "_prepare_inputs_in_ms": 9.457632376637775,
+      "compute_loss_in_ms": 535.8551492407569,
+      "epoch": 16.53,
+      "learning_rate/full": 0.00036251461734986354,
+      "loss": 3.7357,
+      "step": 20000,
+      "training_step_in_ms": 960.2882608400541
+    },
+    {
+      "epoch": 16.53,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.6332855224609375,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.26106709230289454,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 103.6434,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.719,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.241,
+      "step": 20000
+    },
+    {
+      "_prepare_inputs_in_ms": 9.350896289958278,
+      "compute_loss_in_ms": 534.3410268089501,
+      "epoch": 17.36,
+      "learning_rate/full": 0.0003587689869947874,
+      "loss": 3.7335,
+      "step": 21000,
+      "training_step_in_ms": 951.9007690685103
+    },
+    {
+      "_prepare_inputs_in_ms": 9.444510843255557,
+      "compute_loss_in_ms": 536.0995619180612,
+      "epoch": 18.18,
+      "learning_rate/full": 0.0003548663566868928,
+      "loss": 3.726,
+      "step": 22000,
+      "training_step_in_ms": 959.7129821323324
+    },
+    {
+      "_prepare_inputs_in_ms": 9.432463558274321,
+      "compute_loss_in_ms": 535.7645625439472,
+      "epoch": 19.01,
+      "learning_rate/full": 0.00035081058557264036,
+      "loss": 3.7237,
+      "step": 23000,
+      "training_step_in_ms": 963.4282936626696
+    },
+    {
+      "_prepare_inputs_in_ms": 9.293582127720583,
+      "compute_loss_in_ms": 535.7767343233572,
+      "epoch": 19.83,
+      "learning_rate/full": 0.00034660568423297317,
+      "loss": 3.7142,
+      "step": 24000,
+      "training_step_in_ms": 953.5687257153331
+    },
+    {
+      "_prepare_inputs_in_ms": 9.35004804620985,
+      "compute_loss_in_ms": 537.011823366629,
+      "epoch": 20.66,
+      "learning_rate/full": 0.00034225581071742215,
+      "loss": 3.7071,
+      "step": 25000,
+      "training_step_in_ms": 957.869783453003
+    },
+    {
+      "epoch": 20.66,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.6175005435943604,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.26405073790614886,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 104.9682,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.621,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.238,
+      "step": 25000
+    },
+    {
+      "_prepare_inputs_in_ms": 9.984895624163583,
+      "compute_loss_in_ms": 534.7242227215902,
+      "epoch": 21.49,
+      "learning_rate/full": 0.0003377607023892813,
+      "loss": 3.7083,
+      "step": 26000,
+      "training_step_in_ms": 955.9584859979805
+    },
+    {
+      "_prepare_inputs_in_ms": 10.137336196319666,
+      "compute_loss_in_ms": 535.9911008346826,
+      "epoch": 22.31,
+      "learning_rate/full": 0.00033313379374801615,
+      "loss": 3.7022,
+      "step": 27000,
+      "training_step_in_ms": 965.8537974991486
+    },
+    {
+      "_prepare_inputs_in_ms": 10.160301688476466,
+      "compute_loss_in_ms": 536.8555397295277,
+      "epoch": 23.14,
+      "learning_rate/full": 0.0003283704069974003,
+      "loss": 3.7009,
+      "step": 28000,
+      "training_step_in_ms": 961.128851325775
+    },
+    {
+      "_prepare_inputs_in_ms": 10.110118357348256,
+      "compute_loss_in_ms": 536.1848762420123,
+      "epoch": 23.97,
+      "learning_rate/full": 0.0003234847784387827,
+      "loss": 3.6952,
+      "step": 29000,
+      "training_step_in_ms": 951.6901445918484
+    },
+    {
+      "_prepare_inputs_in_ms": 10.067090163414832,
+      "compute_loss_in_ms": 535.8437895118841,
+      "epoch": 24.79,
+      "learning_rate/full": 0.0003184770409925566,
+      "loss": 3.6897,
+      "step": 30000,
+      "training_step_in_ms": 964.9359193513519
+    },
+    {
+      "epoch": 24.79,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.603607416152954,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.26295626788132803,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 101.4282,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.887,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.246,
+      "step": 30000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.085794753678961,
+      "compute_loss_in_ms": 535.1409671568545,
+      "epoch": 25.62,
+      "learning_rate/full": 0.00031335214659900727,
+      "loss": 3.6866,
+      "step": 31000,
+      "training_step_in_ms": 961.3275795525988
+    },
+    {
+      "_prepare_inputs_in_ms": 10.185186771384906,
+      "compute_loss_in_ms": 536.1119842829066,
+      "epoch": 26.45,
+      "learning_rate/full": 0.0003081098663889108,
+      "loss": 3.6829,
+      "step": 32000,
+      "training_step_in_ms": 963.5977011280484
+    },
+    {
+      "_prepare_inputs_in_ms": 10.249891238170676,
+      "compute_loss_in_ms": 535.0950576688629,
+      "epoch": 27.27,
+      "learning_rate/full": 0.0003027712689775676,
+      "loss": 3.6771,
+      "step": 33000,
+      "training_step_in_ms": 955.397748134099
+    },
+    {
+      "_prepare_inputs_in_ms": 10.17424620629754,
+      "compute_loss_in_ms": 534.6862363539985,
+      "epoch": 28.1,
+      "learning_rate/full": 0.0002973202486383942,
+      "loss": 3.672,
+      "step": 34000,
+      "training_step_in_ms": 958.4871930885711
+    },
+    {
+      "_prepare_inputs_in_ms": 10.059338139777537,
+      "compute_loss_in_ms": 535.8044806576218,
+      "epoch": 28.93,
+      "learning_rate/full": 0.000291778393459607,
+      "loss": 3.6714,
+      "step": 35000,
+      "training_step_in_ms": 956.3396593880607
+    },
+    {
+      "epoch": 28.93,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.590414524078369,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.26513770557999966,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 104.3253,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.668,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.24,
+      "step": 35000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.062245466610099,
+      "compute_loss_in_ms": 534.5042656344594,
+      "epoch": 29.75,
+      "learning_rate/full": 0.0002861457824996332,
+      "loss": 3.6645,
+      "step": 36000,
+      "training_step_in_ms": 959.2977655951399
+    },
+    {
+      "_prepare_inputs_in_ms": 10.064813164470252,
+      "compute_loss_in_ms": 536.1598941528937,
+      "epoch": 30.58,
+      "learning_rate/full": 0.00028042798560981287,
+      "loss": 3.6574,
+      "step": 37000,
+      "training_step_in_ms": 956.0697433989844
+    },
+    {
+      "_prepare_inputs_in_ms": 10.116965677938424,
+      "compute_loss_in_ms": 536.2070164434845,
+      "epoch": 31.4,
+      "learning_rate/full": 0.0002746248158102387,
+      "loss": 3.6578,
+      "step": 38000,
+      "training_step_in_ms": 956.2041009759996
+    },
+    {
+      "_prepare_inputs_in_ms": 10.094019736570772,
+      "compute_loss_in_ms": 535.4710150305182,
+      "epoch": 32.23,
+      "learning_rate/full": 0.0002687536169947349,
+      "loss": 3.6524,
+      "step": 39000,
+      "training_step_in_ms": 959.96617779386
+    },
+    {
+      "_prepare_inputs_in_ms": 10.109433323028497,
+      "compute_loss_in_ms": 535.6941579723498,
+      "epoch": 33.06,
+      "learning_rate/full": 0.0002628144306280816,
+      "loss": 3.6524,
+      "step": 40000,
+      "training_step_in_ms": 961.5450095872511
+    },
+    {
+      "epoch": 33.06,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.580965280532837,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2693198298047019,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 99.8174,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 8.015,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.25,
+      "step": 40000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.042415707621997,
+      "compute_loss_in_ms": 535.0675269728526,
+      "epoch": 33.88,
+      "learning_rate/full": 0.000256813129721104,
+      "loss": 3.6463,
+      "step": 41000,
+      "training_step_in_ms": 951.0752835389576
+    },
+    {
+      "_prepare_inputs_in_ms": 10.117303489823826,
+      "compute_loss_in_ms": 535.5392771296902,
+      "epoch": 34.71,
+      "learning_rate/full": 0.00025076173835033525,
+      "loss": 3.6408,
+      "step": 42000,
+      "training_step_in_ms": 959.7704490462202
+    },
+    {
+      "_prepare_inputs_in_ms": 10.175786619714927,
+      "compute_loss_in_ms": 536.3846810262767,
+      "epoch": 35.54,
+      "learning_rate/full": 0.0002446479775734085,
+      "loss": 3.6381,
+      "step": 43000,
+      "training_step_in_ms": 963.5900993177202
+    },
+    {
+      "_prepare_inputs_in_ms": 10.058058980386704,
+      "compute_loss_in_ms": 535.5972880260088,
+      "epoch": 36.36,
+      "learning_rate/full": 0.00023849615593840492,
+      "loss": 3.6361,
+      "step": 44000,
+      "training_step_in_ms": 966.5910823547165
+    },
+    {
+      "_prepare_inputs_in_ms": 10.140807795512956,
+      "compute_loss_in_ms": 535.745742837491,
+      "epoch": 37.19,
+      "learning_rate/full": 0.00023230626707893625,
+      "loss": 3.6344,
+      "step": 45000,
+      "training_step_in_ms": 960.7837667464628
+    },
+    {
+      "epoch": 37.19,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.575321912765503,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2694013022299941,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 103.7207,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.713,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.241,
+      "step": 45000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.123616369532012,
+      "compute_loss_in_ms": 534.5586088547134,
+      "epoch": 38.02,
+      "learning_rate/full": 0.00022608443191494596,
+      "loss": 3.6259,
+      "step": 46000,
+      "training_step_in_ms": 954.3300566331018
+    },
+    {
+      "_prepare_inputs_in_ms": 10.126525112020317,
+      "compute_loss_in_ms": 536.6205943481764,
+      "epoch": 38.84,
+      "learning_rate/full": 0.0002198305382138328,
+      "loss": 3.6279,
+      "step": 47000,
+      "training_step_in_ms": 953.759260071849
+    },
+    {
+      "_prepare_inputs_in_ms": 10.257417954970151,
+      "compute_loss_in_ms": 535.2681058159797,
+      "epoch": 39.67,
+      "learning_rate/full": 0.00021356327694485794,
+      "loss": 3.6213,
+      "step": 48000,
+      "training_step_in_ms": 962.4967881785124
+    },
+    {
+      "_prepare_inputs_in_ms": 10.271113389520906,
+      "compute_loss_in_ms": 535.4450263003819,
+      "epoch": 40.5,
+      "learning_rate/full": 0.0002072826035235433,
+      "loss": 3.6186,
+      "step": 49000,
+      "training_step_in_ms": 961.3730522751575
+    },
+    {
+      "_prepare_inputs_in_ms": 10.129716445459053,
+      "compute_loss_in_ms": 536.4764073403785,
+      "epoch": 41.32,
+      "learning_rate/full": 0.00020099472864285533,
+      "loss": 3.6137,
+      "step": 50000,
+      "training_step_in_ms": 960.2726457127137
+    },
+    {
+      "epoch": 41.32,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5703847408294678,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2688142864806811,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 100.2164,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.983,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.249,
+      "step": 50000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.502501730693549,
+      "compute_loss_in_ms": 535.0548804986756,
+      "epoch": 42.15,
+      "learning_rate/full": 0.0001947058701169798,
+      "loss": 3.6106,
+      "step": 51000,
+      "training_step_in_ms": 958.0683015501127
+    },
+    {
+      "_prepare_inputs_in_ms": 10.785648202290758,
+      "compute_loss_in_ms": 536.7012275556335,
+      "epoch": 42.98,
+      "learning_rate/full": 0.0001884222467327876,
+      "loss": 3.6073,
+      "step": 52000,
+      "training_step_in_ms": 951.947691895999
+    },
+    {
+      "_prepare_inputs_in_ms": 10.585681669297628,
+      "compute_loss_in_ms": 535.5102476192405,
+      "epoch": 43.8,
+      "learning_rate/full": 0.00018214380145695523,
+      "loss": 3.6024,
+      "step": 53000,
+      "training_step_in_ms": 962.6988452640362
+    },
+    {
+      "_prepare_inputs_in_ms": 10.68288057774771,
+      "compute_loss_in_ms": 535.9786380403675,
+      "epoch": 44.63,
+      "learning_rate/full": 0.00017588929863586686,
+      "loss": 3.6046,
+      "step": 54000,
+      "training_step_in_ms": 966.1844932027161
+    },
+    {
+      "_prepare_inputs_in_ms": 10.645952994469553,
+      "compute_loss_in_ms": 535.4971501872642,
+      "epoch": 45.45,
+      "learning_rate/full": 0.00016965863787013347,
+      "loss": 3.5954,
+      "step": 55000,
+      "training_step_in_ms": 966.1277901912108
+    },
+    {
+      "epoch": 45.45,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5649678707122803,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.27166900209048683,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 102.2773,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.822,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.244,
+      "step": 55000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.577075158859172,
+      "compute_loss_in_ms": 534.4740462127374,
+      "epoch": 46.28,
+      "learning_rate/full": 0.00016345179061642948,
+      "loss": 3.5963,
+      "step": 56000,
+      "training_step_in_ms": 961.5655309080612
+    },
+    {
+      "_prepare_inputs_in_ms": 10.732970108627342,
+      "compute_loss_in_ms": 537.2660472553689,
+      "epoch": 47.11,
+      "learning_rate/full": 0.0001572873072366209,
+      "loss": 3.597,
+      "step": 57000,
+      "training_step_in_ms": 964.5027116436977
+    },
+    {
+      "_prepare_inputs_in_ms": 10.73382615565788,
+      "compute_loss_in_ms": 538.1741794921691,
+      "epoch": 47.93,
+      "learning_rate/full": 0.0001511650606366491,
+      "loss": 3.5895,
+      "step": 58000,
+      "training_step_in_ms": 949.1175830988213
+    },
+    {
+      "_prepare_inputs_in_ms": 10.706503831432201,
+      "compute_loss_in_ms": 536.7453673920827,
+      "epoch": 48.76,
+      "learning_rate/full": 0.0001450911048478807,
+      "loss": 3.5876,
+      "step": 59000,
+      "training_step_in_ms": 968.0035566822626
+    },
+    {
+      "_prepare_inputs_in_ms": 10.903483389178291,
+      "compute_loss_in_ms": 537.1148272417486,
+      "epoch": 49.59,
+      "learning_rate/full": 0.00013906544966141887,
+      "loss": 3.5819,
+      "step": 60000,
+      "training_step_in_ms": 964.3682058413979
+    },
+    {
+      "epoch": 49.59,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5601305961608887,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.27116285100803245,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 101.8706,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.853,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.245,
+      "step": 60000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.826663774233765,
+      "compute_loss_in_ms": 535.5709582263371,
+      "epoch": 50.41,
+      "learning_rate/full": 0.00013310610391163933,
+      "loss": 3.5762,
+      "step": 61000,
+      "training_step_in_ms": 957.2617566076806
+    },
+    {
+      "_prepare_inputs_in_ms": 10.595198371564038,
+      "compute_loss_in_ms": 536.0367681181524,
+      "epoch": 51.24,
+      "learning_rate/full": 0.0001272129067134662,
+      "loss": 3.5751,
+      "step": 62000,
+      "training_step_in_ms": 959.1705903129186
+    },
+    {
+      "_prepare_inputs_in_ms": 10.456255728611723,
+      "compute_loss_in_ms": 536.2733452994144,
+      "epoch": 52.07,
+      "learning_rate/full": 0.00012139747473708569,
+      "loss": 3.5774,
+      "step": 63000,
+      "training_step_in_ms": 955.730251706671
+    },
+    {
+      "_prepare_inputs_in_ms": 10.639551113941707,
+      "compute_loss_in_ms": 535.7200796955731,
+      "epoch": 52.89,
+      "learning_rate/full": 0.00011564819693413189,
+      "loss": 3.568,
+      "step": 64000,
+      "training_step_in_ms": 951.6928666429594
+    },
+    {
+      "_prepare_inputs_in_ms": 10.60645360336639,
+      "compute_loss_in_ms": 535.3036272318568,
+      "epoch": 53.72,
+      "learning_rate/full": 0.00010998812020655949,
+      "loss": 3.5703,
+      "step": 65000,
+      "training_step_in_ms": 963.1167565376963
+    },
+    {
+      "epoch": 53.72,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5590660572052,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2733920686032914,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 99.9831,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 8.001,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.25,
+      "step": 65000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.624523456094831,
+      "compute_loss_in_ms": 534.0653136165347,
+      "epoch": 54.55,
+      "learning_rate/full": 0.00010441152222708468,
+      "loss": 3.5693,
+      "step": 66000,
+      "training_step_in_ms": 955.8435329974163
+    },
+    {
+      "_prepare_inputs_in_ms": 10.65835806389805,
+      "compute_loss_in_ms": 536.2088124424918,
+      "epoch": 55.37,
+      "learning_rate/full": 9.893506973588506e-05,
+      "loss": 3.5599,
+      "step": 67000,
+      "training_step_in_ms": 957.1059285002993
+    },
+    {
+      "_prepare_inputs_in_ms": 10.751230709021911,
+      "compute_loss_in_ms": 536.6291307259817,
+      "epoch": 56.2,
+      "learning_rate/full": 9.355322604944605e-05,
+      "loss": 3.5668,
+      "step": 68000,
+      "training_step_in_ms": 959.1418899303535
+    },
+    {
+      "_prepare_inputs_in_ms": 10.653449523961172,
+      "compute_loss_in_ms": 536.7383008667966,
+      "epoch": 57.02,
+      "learning_rate/full": 8.829252005938348e-05,
+      "loss": 3.5592,
+      "step": 69000,
+      "training_step_in_ms": 955.6575022591278
+    },
+    {
+      "_prepare_inputs_in_ms": 10.65822267276235,
+      "compute_loss_in_ms": 535.9257607464679,
+      "epoch": 57.85,
+      "learning_rate/full": 8.312650768346744e-05,
+      "loss": 3.5575,
+      "step": 70000,
+      "training_step_in_ms": 951.2662082569441
+    },
+    {
+      "epoch": 57.85,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5556719303131104,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2734467694599133,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 104.2955,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.671,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.24,
+      "step": 70000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.608355628053953,
+      "compute_loss_in_ms": 534.1021929110866,
+      "epoch": 58.68,
+      "learning_rate/full": 7.80762980438341e-05,
+      "loss": 3.5556,
+      "step": 71000,
+      "training_step_in_ms": 948.2094520897372
+    },
+    {
+      "_prepare_inputs_in_ms": 10.676257735467516,
+      "compute_loss_in_ms": 536.0682506592711,
+      "epoch": 59.5,
+      "learning_rate/full": 7.315176250595717e-05,
+      "loss": 3.5537,
+      "step": 72000,
+      "training_step_in_ms": 954.8686325427843
+    },
+    {
+      "_prepare_inputs_in_ms": 10.610922348219901,
+      "compute_loss_in_ms": 536.3908133659279,
+      "epoch": 60.33,
+      "learning_rate/full": 6.835266183844516e-05,
+      "loss": 3.5479,
+      "step": 73000,
+      "training_step_in_ms": 954.1712517963024
+    },
+    {
+      "_prepare_inputs_in_ms": 10.579945259494707,
+      "compute_loss_in_ms": 534.4352571795462,
+      "epoch": 61.16,
+      "learning_rate/full": 6.368374166947542e-05,
+      "loss": 3.5467,
+      "step": 74000,
+      "training_step_in_ms": 953.0113322847756
+    },
+    {
+      "_prepare_inputs_in_ms": 10.577161580207758,
+      "compute_loss_in_ms": 536.7000861178385,
+      "epoch": 61.98,
+      "learning_rate/full": 5.914514927911328e-05,
+      "loss": 3.5466,
+      "step": 75000,
+      "training_step_in_ms": 950.8998943779152
+    },
+    {
+      "epoch": 61.98,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5562994480133057,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.27468764903823617,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 102.8368,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.779,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.243,
+      "step": 75000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.397925437655209,
+      "compute_loss_in_ms": 535.1076509790728,
+      "epoch": 62.81,
+      "learning_rate/full": 5.475477712440255e-05,
+      "loss": 3.5454,
+      "step": 76000,
+      "training_step_in_ms": 949.9264260384953
+    },
+    {
+      "_prepare_inputs_in_ms": 10.516429967130534,
+      "compute_loss_in_ms": 537.9902690803865,
+      "epoch": 63.64,
+      "learning_rate/full": 5.0499380128392283e-05,
+      "loss": 3.5418,
+      "step": 77000,
+      "training_step_in_ms": 963.6520826652413
+    },
+    {
+      "_prepare_inputs_in_ms": 10.470940343337134,
+      "compute_loss_in_ms": 536.5116302901879,
+      "epoch": 64.46,
+      "learning_rate/full": 4.639614598504125e-05,
+      "loss": 3.5397,
+      "step": 78000,
+      "training_step_in_ms": 957.950764612644
+    },
+    {
+      "_prepare_inputs_in_ms": 10.59993867285084,
+      "compute_loss_in_ms": 537.1390544432215,
+      "epoch": 65.29,
+      "learning_rate/full": 4.244480421242036e-05,
+      "loss": 3.5388,
+      "step": 79000,
+      "training_step_in_ms": 961.8771279493812
+    },
+    {
+      "_prepare_inputs_in_ms": 10.679984404356219,
+      "compute_loss_in_ms": 536.696684517432,
+      "epoch": 66.12,
+      "learning_rate/full": 3.8649262125702656e-05,
+      "loss": 3.5329,
+      "step": 80000,
+      "training_step_in_ms": 957.4748168167425
+    },
+    {
+      "epoch": 66.12,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.556597948074341,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2742539533100543,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 100.6419,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.949,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.248,
+      "step": 80000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.708052520753771,
+      "compute_loss_in_ms": 535.2645241598366,
+      "epoch": 66.94,
+      "learning_rate/full": 3.500971449864994e-05,
+      "loss": 3.5345,
+      "step": 81000,
+      "training_step_in_ms": 948.9210075238952
+    },
+    {
+      "_prepare_inputs_in_ms": 10.605949840741232,
+      "compute_loss_in_ms": 535.3940441570012,
+      "epoch": 67.77,
+      "learning_rate/full": 3.153703882232173e-05,
+      "loss": 3.5346,
+      "step": 82000,
+      "training_step_in_ms": 957.2232636878034
+    },
+    {
+      "_prepare_inputs_in_ms": 10.527875950676389,
+      "compute_loss_in_ms": 535.3903734084452,
+      "epoch": 68.6,
+      "learning_rate/full": 2.823094906089525e-05,
+      "loss": 3.5362,
+      "step": 83000,
+      "training_step_in_ms": 960.9660885017365
+    },
+    {
+      "_prepare_inputs_in_ms": 10.645452778204344,
+      "compute_loss_in_ms": 535.3771834741347,
+      "epoch": 69.42,
+      "learning_rate/full": 2.509166124130553e-05,
+      "loss": 3.5284,
+      "step": 84000,
+      "training_step_in_ms": 959.7084383748006
+    },
+    {
+      "_prepare_inputs_in_ms": 10.593148670741357,
+      "compute_loss_in_ms": 534.9863913216395,
+      "epoch": 70.25,
+      "learning_rate/full": 2.212855773155269e-05,
+      "loss": 3.5273,
+      "step": 85000,
+      "training_step_in_ms": 970.6172955055954
+    },
+    {
+      "epoch": 70.25,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.555929660797119,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2741459251366367,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 104.7737,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.636,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.239,
+      "step": 85000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.723522963888216,
+      "compute_loss_in_ms": 534.6476732452866,
+      "epoch": 71.07,
+      "learning_rate/full": 1.9341343786710864e-05,
+      "loss": 3.5266,
+      "step": 86000,
+      "training_step_in_ms": 961.9524979006965
+    },
+    {
+      "_prepare_inputs_in_ms": 10.785187994129956,
+      "compute_loss_in_ms": 536.4809199275915,
+      "epoch": 71.9,
+      "learning_rate/full": 1.6732775565058435e-05,
+      "loss": 3.5253,
+      "step": 87000,
+      "training_step_in_ms": 950.8976757206256
+    },
+    {
+      "_prepare_inputs_in_ms": 10.603952513309196,
+      "compute_loss_in_ms": 535.7741147053894,
+      "epoch": 72.73,
+      "learning_rate/full": 1.4305432569654864e-05,
+      "loss": 3.524,
+      "step": 88000,
+      "training_step_in_ms": 959.2451942999614
+    },
+    {
+      "_prepare_inputs_in_ms": 10.706391342449933,
+      "compute_loss_in_ms": 536.408079084591,
+      "epoch": 73.55,
+      "learning_rate/full": 1.2059561887499037e-05,
+      "loss": 3.5265,
+      "step": 89000,
+      "training_step_in_ms": 973.4634857769124
+    },
+    {
+      "_prepare_inputs_in_ms": 10.6856812821934,
+      "compute_loss_in_ms": 537.2633295034757,
+      "epoch": 74.38,
+      "learning_rate/full": 1.0003841866378549e-05,
+      "loss": 3.5218,
+      "step": 90000,
+      "training_step_in_ms": 966.5726283045951
+    },
+    {
+      "epoch": 74.38,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5546064376831055,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.27359656588473774,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 104.0763,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.687,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.24,
+      "step": 90000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.64931857888019,
+      "compute_loss_in_ms": 535.808490979718,
+      "epoch": 75.21,
+      "learning_rate/full": 8.132070715440754e-06,
+      "loss": 3.5219,
+      "step": 91000,
+      "training_step_in_ms": 957.4841072742129
+    },
+    {
+      "_prepare_inputs_in_ms": 10.87383456970565,
+      "compute_loss_in_ms": 537.6603767276974,
+      "epoch": 76.03,
+      "learning_rate/full": 6.453582116869461e-06,
+      "loss": 3.5235,
+      "step": 92000,
+      "training_step_in_ms": 954.7427223158302
+    },
+    {
+      "_prepare_inputs_in_ms": 10.69017239450477,
+      "compute_loss_in_ms": 537.0892604731489,
+      "epoch": 76.86,
+      "learning_rate/full": 4.964706300643118e-06,
+      "loss": 3.5189,
+      "step": 93000,
+      "training_step_in_ms": 950.3190122217638
+    },
+    {
+      "_prepare_inputs_in_ms": 10.618111693882383,
+      "compute_loss_in_ms": 535.7142543839291,
+      "epoch": 77.69,
+      "learning_rate/full": 3.6686926578696213e-06,
+      "loss": 3.5235,
+      "step": 94000,
+      "training_step_in_ms": 958.8555470507126
+    },
+    {
+      "_prepare_inputs_in_ms": 10.73988960427232,
+      "compute_loss_in_ms": 536.3263423398603,
+      "epoch": 78.51,
+      "learning_rate/full": 2.5658174353882404e-06,
+      "loss": 3.5215,
+      "step": 95000,
+      "training_step_in_ms": 964.6057669939473
+    },
+    {
+      "epoch": 78.51,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.5548150539398193,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.27406488207696067,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 104.9846,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.62,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.238,
+      "step": 95000
+    },
+    {
+      "_prepare_inputs_in_ms": 10.761066970824286,
+      "compute_loss_in_ms": 535.5862640386913,
+      "epoch": 79.34,
+      "learning_rate/full": 1.6593767888209988e-06,
+      "loss": 3.5201,
+      "step": 96000,
+      "training_step_in_ms": 953.487945659901
+    },
+    {
+      "_prepare_inputs_in_ms": 10.744887009612285,
+      "compute_loss_in_ms": 536.8042176247109,
+      "epoch": 80.17,
+      "learning_rate/full": 9.484543083217335e-07,
+      "loss": 3.5242,
+      "step": 97000,
+      "training_step_in_ms": 955.935797311482
+    },
+    {
+      "_prepare_inputs_in_ms": 10.609967146418057,
+      "compute_loss_in_ms": 536.5163776981644,
+      "epoch": 80.99,
+      "learning_rate/full": 4.3558991431862373e-07,
+      "loss": 3.5217,
+      "step": 98000,
+      "training_step_in_ms": 951.0243763268227
+    },
+    {
+      "_prepare_inputs_in_ms": 10.550453630159609,
+      "compute_loss_in_ms": 535.8989973879652,
+      "epoch": 81.82,
+      "learning_rate/full": 1.1945383759244255e-07,
+      "loss": 3.523,
+      "step": 99000,
+      "training_step_in_ms": 959.1436755338218
+    },
+    {
+      "_prepare_inputs_in_ms": 10.565008323756047,
+      "compute_loss_in_ms": 536.139583913493,
+      "epoch": 82.64,
+      "learning_rate/full": 9.51678964100644e-10,
+      "loss": 3.523,
+      "step": 100000,
+      "training_step_in_ms": 958.2332628678996
+    },
+    {
+      "epoch": 82.64,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_loss": 3.554699182510376,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_meteor": 0.2739906528866304,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_runtime": 105.8588,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_samples_per_second": 7.557,
+      "eval_visual_genome-region_descriptions_v1.2.0-test_steps_per_second": 0.236,
+      "step": 100000
+    }
+  ],
+  "max_steps": 100000,
+  "num_train_epochs": 83,
+  "total_flos": 1.7305644974792963e+23,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

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

zero_to_fp32.py

@@ -0,0 +1,587 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 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
+from dataclasses import dataclass
+
+# 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.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+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_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        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]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dict = torch.load(f, map_location=device)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    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_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # 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.")
+
+    # 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")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    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 _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # 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
+
+    # 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
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # 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"Trainable params: {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 Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    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("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    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, tag=args.tag)