Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoint-1506/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+checkpoint-3012/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+checkpoint-4518/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+checkpoint-6024/tokenizer.json filter=lfs diff=lfs merge=lfs -text

checkpoint-1506/1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+  "word_embedding_dimension": 1024,
+  "pooling_mode_cls_token": true,
+  "pooling_mode_mean_tokens": false,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false,
+  "pooling_mode_weightedmean_tokens": false,
+  "pooling_mode_lasttoken": false,
+  "include_prompt": true
+}

checkpoint-1506/README.md

@@ -0,0 +1,57 @@
+---
+library_name: sentence-transformers
+pipeline_tag: sentence-similarity
+tags:
+- sentence-transformers
+- feature-extraction
+- sentence-similarity
+
+---
+
+# {MODEL_NAME}
+
+This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 1024 dimensional dense vector space and can be used for tasks like clustering or semantic search.
+
+<!--- Describe your model here -->
+
+## Usage (Sentence-Transformers)
+
+Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
+
+```
+pip install -U sentence-transformers
+```
+
+Then you can use the model like this:
+
+```python
+from sentence_transformers import SentenceTransformer
+sentences = ["This is an example sentence", "Each sentence is converted"]
+
+model = SentenceTransformer('{MODEL_NAME}')
+embeddings = model.encode(sentences)
+print(embeddings)
+```
+
+
+
+## Evaluation Results
+
+<!--- Describe how your model was evaluated -->
+
+For an automated evaluation of this model, see the *Sentence Embeddings Benchmark*: [https://seb.sbert.net](https://seb.sbert.net?model_name={MODEL_NAME})
+
+
+
+## Full Model Architecture
+```
+SentenceTransformer(
+  (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False}) with Transformer model: XLMRobertaModel 
+  (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
+  (2): Normalize()
+)
+```
+
+## Citing & Authors
+
+<!--- Describe where people can find more information -->

checkpoint-1506/colbert_linear.pt

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

checkpoint-1506/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "saved_models/bgem3_unified_finetune_20240330/checkpoint-1506",
+  "architectures": [
+    "XLMRobertaModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 8194,
+  "model_type": "xlm-roberta",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.39.1",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 250002
+}

checkpoint-1506/config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.6.0",
+    "transformers": "4.39.1",
+    "pytorch": "2.0.1+cu117"
+  },
+  "prompts": {},
+  "default_prompt_name": null
+}

checkpoint-1506/global_step1506/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt

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

checkpoint-1506/global_step1506/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt

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

checkpoint-1506/global_step1506/mp_rank_00_model_states.pt

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

checkpoint-1506/latest

@@ -0,0 +1 @@
+global_step1506

checkpoint-1506/model.safetensors

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

checkpoint-1506/modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

checkpoint-1506/rng_state_0.pth

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

checkpoint-1506/rng_state_1.pth

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

checkpoint-1506/sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 8192,
+  "do_lower_case": false
+}

checkpoint-1506/sentencepiece.bpe.model

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

checkpoint-1506/sparse_linear.pt

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

checkpoint-1506/special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint-1506/tokenizer.json

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

checkpoint-1506/tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "250001": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_max_length": 8192,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}

checkpoint-1506/trainer_state.json

@@ -0,0 +1,1071 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 500,
+  "global_step": 1506,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.01,
+      "grad_norm": 41.35739895402257,
+      "learning_rate": 7.193423539345941e-06,
+      "loss": 0.5141,
+      "step": 10
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 10.443694874625423,
+      "learning_rate": 9.358859796204429e-06,
+      "loss": 0.4195,
+      "step": 20
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 5.720021591145773,
+      "learning_rate": 1.0625558804168632e-05,
+      "loss": 0.2851,
+      "step": 30
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 7.59250257002781,
+      "learning_rate": 1.1524296053062918e-05,
+      "loss": 0.2174,
+      "step": 40
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 11.521550428114184,
+      "learning_rate": 1.2221410821833392e-05,
+      "loss": 0.1817,
+      "step": 50
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 6.027998448017936,
+      "learning_rate": 1.2790995061027121e-05,
+      "loss": 0.1886,
+      "step": 60
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 4.929816146115093,
+      "learning_rate": 1.3272571673439616e-05,
+      "loss": 0.1553,
+      "step": 70
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 2.3881515971871345,
+      "learning_rate": 1.3689732309921406e-05,
+      "loss": 0.129,
+      "step": 80
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 5.441118636680569,
+      "learning_rate": 1.4057694068991321e-05,
+      "loss": 0.1433,
+      "step": 90
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 6.085213496930297,
+      "learning_rate": 1.4386847078691883e-05,
+      "loss": 0.1092,
+      "step": 100
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 4.713361075579101,
+      "learning_rate": 1.4684602194465794e-05,
+      "loss": 0.1231,
+      "step": 110
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 3.8259500358417924,
+      "learning_rate": 1.495643131788561e-05,
+      "loss": 0.0697,
+      "step": 120
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 2.67920682727699,
+      "learning_rate": 1.5206489871327869e-05,
+      "loss": 0.084,
+      "step": 130
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 8.189491553913532,
+      "learning_rate": 1.54380079302981e-05,
+      "loss": 0.1023,
+      "step": 140
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 12.000369384166694,
+      "learning_rate": 1.5653546086656083e-05,
+      "loss": 0.0972,
+      "step": 150
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 2.3726942012239953,
+      "learning_rate": 1.5855168566779895e-05,
+      "loss": 0.1036,
+      "step": 160
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 2.0242117206977044,
+      "learning_rate": 1.604456377435124e-05,
+      "loss": 0.1081,
+      "step": 170
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 2.2261934744219967,
+      "learning_rate": 1.6223130325849813e-05,
+      "loss": 0.082,
+      "step": 180
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 4.70421073268978,
+      "learning_rate": 1.6392039793463407e-05,
+      "loss": 0.0744,
+      "step": 190
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 12.064244277626111,
+      "learning_rate": 1.6552283335550368e-05,
+      "loss": 0.0934,
+      "step": 200
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 2.150990470911233,
+      "learning_rate": 1.67047069382623e-05,
+      "loss": 0.0737,
+      "step": 210
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 4.074258308547766,
+      "learning_rate": 1.6850038451324284e-05,
+      "loss": 0.0841,
+      "step": 220
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 2.6764310841274237,
+      "learning_rate": 1.6988908609137504e-05,
+      "loss": 0.0821,
+      "step": 230
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 3.1391283710987063,
+      "learning_rate": 1.71218675747441e-05,
+      "loss": 0.0747,
+      "step": 240
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 3.1103728720325754,
+      "learning_rate": 1.7249398104320845e-05,
+      "loss": 0.0907,
+      "step": 250
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 5.354173276825422,
+      "learning_rate": 1.7371926128186358e-05,
+      "loss": 0.0765,
+      "step": 260
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 2.8677808267225005,
+      "learning_rate": 1.7489829333814013e-05,
+      "loss": 0.0777,
+      "step": 270
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 2.6124301773304874,
+      "learning_rate": 1.760344418715659e-05,
+      "loss": 0.0897,
+      "step": 280
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 2.915238130503912,
+      "learning_rate": 1.7713071721324668e-05,
+      "loss": 0.0866,
+      "step": 290
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 4.061775729833908,
+      "learning_rate": 1.781898234351457e-05,
+      "loss": 0.0804,
+      "step": 300
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.5648704857906242,
+      "learning_rate": 1.7921419853452233e-05,
+      "loss": 0.0834,
+      "step": 310
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 2.225837991191409,
+      "learning_rate": 1.8020604823638384e-05,
+      "loss": 0.0646,
+      "step": 320
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 4.041622096074851,
+      "learning_rate": 1.811673745928848e-05,
+      "loss": 0.0636,
+      "step": 330
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.21100008000919218,
+      "learning_rate": 1.821000003120973e-05,
+      "loss": 0.0656,
+      "step": 340
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 2.2923304434569443,
+      "learning_rate": 1.8300558955927067e-05,
+      "loss": 0.0756,
+      "step": 350
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 3.6524729986322013,
+      "learning_rate": 1.83885665827083e-05,
+      "loss": 0.0466,
+      "step": 360
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 2.1793660025907284,
+      "learning_rate": 1.847416273569235e-05,
+      "loss": 0.0704,
+      "step": 370
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 4.363503091034068,
+      "learning_rate": 1.8557476050321896e-05,
+      "loss": 0.0712,
+      "step": 380
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 3.0942807753734476,
+      "learning_rate": 1.863862513615056e-05,
+      "loss": 0.0831,
+      "step": 390
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 1.789230871814937,
+      "learning_rate": 1.8717719592408857e-05,
+      "loss": 0.078,
+      "step": 400
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 3.4361028432853677,
+      "learning_rate": 1.879486089815082e-05,
+      "loss": 0.0663,
+      "step": 410
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 3.690439378170989,
+      "learning_rate": 1.8870143195120794e-05,
+      "loss": 0.0738,
+      "step": 420
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 2.159404934400987,
+      "learning_rate": 1.8943653978491198e-05,
+      "loss": 0.0768,
+      "step": 430
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 0.8042070009727055,
+      "learning_rate": 1.901547470818277e-05,
+      "loss": 0.0777,
+      "step": 440
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 1.70758347995601,
+      "learning_rate": 1.9085681351478775e-05,
+      "loss": 0.05,
+      "step": 450
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.2113560599237037,
+      "learning_rate": 1.9154344865995993e-05,
+      "loss": 0.0715,
+      "step": 460
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 2.5816071085486527,
+      "learning_rate": 1.9221531630710657e-05,
+      "loss": 0.0688,
+      "step": 470
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.4447841193953872,
+      "learning_rate": 1.9287303831602588e-05,
+      "loss": 0.0659,
+      "step": 480
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 1.2668557011358101,
+      "learning_rate": 1.9351719807533285e-05,
+      "loss": 0.0515,
+      "step": 490
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 2.0353893651728865,
+      "learning_rate": 1.9414834361179333e-05,
+      "loss": 0.0687,
+      "step": 500
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 2.005863891715338,
+      "learning_rate": 1.947669903917393e-05,
+      "loss": 0.0459,
+      "step": 510
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.4330303529877615,
+      "learning_rate": 1.9537362385044847e-05,
+      "loss": 0.0557,
+      "step": 520
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.8683959370360874,
+      "learning_rate": 1.959687016805845e-05,
+      "loss": 0.0656,
+      "step": 530
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 1.5963373312962135,
+      "learning_rate": 1.9655265590672502e-05,
+      "loss": 0.0519,
+      "step": 540
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 1.824110595633705,
+      "learning_rate": 1.9712589476953243e-05,
+      "loss": 0.0557,
+      "step": 550
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 2.024299349665376,
+      "learning_rate": 1.976888044401508e-05,
+      "loss": 0.0617,
+      "step": 560
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 1.0490238284158413,
+      "learning_rate": 1.98241750582861e-05,
+      "loss": 0.0529,
+      "step": 570
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.2602675746455845,
+      "learning_rate": 1.9878507978183157e-05,
+      "loss": 0.0617,
+      "step": 580
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.7353304044327778,
+      "learning_rate": 1.9931912084590654e-05,
+      "loss": 0.0464,
+      "step": 590
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 2.043608410331547,
+      "learning_rate": 1.998441860037306e-05,
+      "loss": 0.0716,
+      "step": 600
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 0.9522735357022469,
+      "learning_rate": 1.997786386275595e-05,
+      "loss": 0.0496,
+      "step": 610
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.3740017908881452,
+      "learning_rate": 1.9940970300682533e-05,
+      "loss": 0.0686,
+      "step": 620
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.1533525092583308,
+      "learning_rate": 1.9904076738609114e-05,
+      "loss": 0.0562,
+      "step": 630
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.564671647601224,
+      "learning_rate": 1.9867183176535695e-05,
+      "loss": 0.0553,
+      "step": 640
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 0.8996693997352794,
+      "learning_rate": 1.9830289614462276e-05,
+      "loss": 0.0524,
+      "step": 650
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.9340857910489522,
+      "learning_rate": 1.979339605238886e-05,
+      "loss": 0.0534,
+      "step": 660
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.892682744806536,
+      "learning_rate": 1.9756502490315442e-05,
+      "loss": 0.0625,
+      "step": 670
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.2063853192894494,
+      "learning_rate": 1.9719608928242023e-05,
+      "loss": 0.0519,
+      "step": 680
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.984312822464147,
+      "learning_rate": 1.9682715366168604e-05,
+      "loss": 0.0494,
+      "step": 690
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.6705446030280595,
+      "learning_rate": 1.9645821804095185e-05,
+      "loss": 0.059,
+      "step": 700
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.7422620817223426,
+      "learning_rate": 1.960892824202177e-05,
+      "loss": 0.0454,
+      "step": 710
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.3182961229762868,
+      "learning_rate": 1.957203467994835e-05,
+      "loss": 0.0581,
+      "step": 720
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.330360819481426,
+      "learning_rate": 1.9535141117874932e-05,
+      "loss": 0.0575,
+      "step": 730
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 1.9840741868184866,
+      "learning_rate": 1.9498247555801517e-05,
+      "loss": 0.063,
+      "step": 740
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.6018064731760029,
+      "learning_rate": 1.9461353993728094e-05,
+      "loss": 0.0519,
+      "step": 750
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.391561342963203,
+      "learning_rate": 1.9424460431654675e-05,
+      "loss": 0.0572,
+      "step": 760
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 1.0319629863193043,
+      "learning_rate": 1.938756686958126e-05,
+      "loss": 0.0619,
+      "step": 770
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.7398111190442345,
+      "learning_rate": 1.935067330750784e-05,
+      "loss": 0.05,
+      "step": 780
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.675180749962208,
+      "learning_rate": 1.9313779745434422e-05,
+      "loss": 0.0654,
+      "step": 790
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.486943722740635,
+      "learning_rate": 1.9276886183361007e-05,
+      "loss": 0.0463,
+      "step": 800
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1444156835752686,
+      "learning_rate": 1.9239992621287588e-05,
+      "loss": 0.0598,
+      "step": 810
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1260777714975718,
+      "learning_rate": 1.920309905921417e-05,
+      "loss": 0.039,
+      "step": 820
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.328313971146618,
+      "learning_rate": 1.916620549714075e-05,
+      "loss": 0.0591,
+      "step": 830
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 0.8805517771854091,
+      "learning_rate": 1.912931193506733e-05,
+      "loss": 0.0392,
+      "step": 840
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 0.9704226543415952,
+      "learning_rate": 1.9092418372993916e-05,
+      "loss": 0.056,
+      "step": 850
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 1.4328262810995938,
+      "learning_rate": 1.9055524810920497e-05,
+      "loss": 0.0602,
+      "step": 860
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 2.6133255890647167,
+      "learning_rate": 1.9018631248847078e-05,
+      "loss": 0.0507,
+      "step": 870
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 1.4574352859865667,
+      "learning_rate": 1.898173768677366e-05,
+      "loss": 0.0667,
+      "step": 880
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.9116869451062006,
+      "learning_rate": 1.894484412470024e-05,
+      "loss": 0.0533,
+      "step": 890
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 0.8152499611885836,
+      "learning_rate": 1.890795056262682e-05,
+      "loss": 0.0534,
+      "step": 900
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 2.9440345318883696,
+      "learning_rate": 1.8871057000553406e-05,
+      "loss": 0.0419,
+      "step": 910
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 2.7588896552457074,
+      "learning_rate": 1.8834163438479987e-05,
+      "loss": 0.0418,
+      "step": 920
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.8886319208504987,
+      "learning_rate": 1.8797269876406568e-05,
+      "loss": 0.0397,
+      "step": 930
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.8820615671771377,
+      "learning_rate": 1.876037631433315e-05,
+      "loss": 0.0454,
+      "step": 940
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.2868214865586614,
+      "learning_rate": 1.872348275225973e-05,
+      "loss": 0.0373,
+      "step": 950
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 2.016876778284122,
+      "learning_rate": 1.8686589190186315e-05,
+      "loss": 0.0619,
+      "step": 960
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 0.8992764743601639,
+      "learning_rate": 1.8649695628112896e-05,
+      "loss": 0.0403,
+      "step": 970
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 0.42416507318131186,
+      "learning_rate": 1.8612802066039477e-05,
+      "loss": 0.067,
+      "step": 980
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.8601423331120188,
+      "learning_rate": 1.8575908503966062e-05,
+      "loss": 0.066,
+      "step": 990
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.3893250379275068,
+      "learning_rate": 1.8539014941892643e-05,
+      "loss": 0.0437,
+      "step": 1000
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 2.3240219930504376,
+      "learning_rate": 1.850212137981922e-05,
+      "loss": 0.0488,
+      "step": 1010
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.8849896953543195,
+      "learning_rate": 1.8465227817745805e-05,
+      "loss": 0.0472,
+      "step": 1020
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.8655212477880712,
+      "learning_rate": 1.8428334255672386e-05,
+      "loss": 0.0546,
+      "step": 1030
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 1.2016665548799828,
+      "learning_rate": 1.8391440693598967e-05,
+      "loss": 0.0956,
+      "step": 1040
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 1.492659569010543,
+      "learning_rate": 1.8354547131525552e-05,
+      "loss": 0.0449,
+      "step": 1050
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 0.4224837643985645,
+      "learning_rate": 1.8317653569452133e-05,
+      "loss": 0.0545,
+      "step": 1060
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 1.8405573730480247,
+      "learning_rate": 1.8280760007378714e-05,
+      "loss": 0.0328,
+      "step": 1070
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.3988661237020132,
+      "learning_rate": 1.8243866445305295e-05,
+      "loss": 0.0451,
+      "step": 1080
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.5573396543240564,
+      "learning_rate": 1.8206972883231876e-05,
+      "loss": 0.0502,
+      "step": 1090
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.7855566721285819,
+      "learning_rate": 1.817007932115846e-05,
+      "loss": 0.0341,
+      "step": 1100
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 1.497005175064917,
+      "learning_rate": 1.8133185759085042e-05,
+      "loss": 0.0513,
+      "step": 1110
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 1.650975953086994,
+      "learning_rate": 1.8096292197011623e-05,
+      "loss": 0.0473,
+      "step": 1120
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.8418054071334755,
+      "learning_rate": 1.8059398634938204e-05,
+      "loss": 0.0444,
+      "step": 1130
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 1.33651471077315,
+      "learning_rate": 1.8022505072864785e-05,
+      "loss": 0.0476,
+      "step": 1140
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 1.1529220423121023,
+      "learning_rate": 1.7985611510791367e-05,
+      "loss": 0.0379,
+      "step": 1150
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 0.6912013035706749,
+      "learning_rate": 1.794871794871795e-05,
+      "loss": 0.0526,
+      "step": 1160
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 2.806793867054328,
+      "learning_rate": 1.7911824386644532e-05,
+      "loss": 0.047,
+      "step": 1170
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 3.323409487413899,
+      "learning_rate": 1.7874930824571113e-05,
+      "loss": 0.0455,
+      "step": 1180
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 1.0877783561592873,
+      "learning_rate": 1.7838037262497695e-05,
+      "loss": 0.0555,
+      "step": 1190
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.9158165193439103,
+      "learning_rate": 1.7801143700424276e-05,
+      "loss": 0.0458,
+      "step": 1200
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.0863638939536142,
+      "learning_rate": 1.776425013835086e-05,
+      "loss": 0.0462,
+      "step": 1210
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 2.3398031954044978,
+      "learning_rate": 1.772735657627744e-05,
+      "loss": 0.032,
+      "step": 1220
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.9272917966743286,
+      "learning_rate": 1.7690463014204022e-05,
+      "loss": 0.0464,
+      "step": 1230
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 1.4608044079217914,
+      "learning_rate": 1.7653569452130607e-05,
+      "loss": 0.0544,
+      "step": 1240
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 2.0799767984416713,
+      "learning_rate": 1.7616675890057188e-05,
+      "loss": 0.0629,
+      "step": 1250
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 2.438242481318271,
+      "learning_rate": 1.757978232798377e-05,
+      "loss": 0.0586,
+      "step": 1260
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.22666963149041416,
+      "learning_rate": 1.754288876591035e-05,
+      "loss": 0.0419,
+      "step": 1270
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 1.6906855856145777,
+      "learning_rate": 1.750599520383693e-05,
+      "loss": 0.0417,
+      "step": 1280
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 0.8417294291584847,
+      "learning_rate": 1.7469101641763513e-05,
+      "loss": 0.0496,
+      "step": 1290
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 2.384648750559231,
+      "learning_rate": 1.7432208079690097e-05,
+      "loss": 0.037,
+      "step": 1300
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 1.0969098360012755,
+      "learning_rate": 1.7395314517616678e-05,
+      "loss": 0.0392,
+      "step": 1310
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 4.265947825753472,
+      "learning_rate": 1.735842095554326e-05,
+      "loss": 0.0434,
+      "step": 1320
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 1.0043877066833349,
+      "learning_rate": 1.732152739346984e-05,
+      "loss": 0.0423,
+      "step": 1330
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 0.3834891081373628,
+      "learning_rate": 1.728463383139642e-05,
+      "loss": 0.0423,
+      "step": 1340
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.8917610552257311,
+      "learning_rate": 1.7247740269323006e-05,
+      "loss": 0.0491,
+      "step": 1350
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 1.7777935423085822,
+      "learning_rate": 1.7210846707249587e-05,
+      "loss": 0.0465,
+      "step": 1360
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 2.882686457165806,
+      "learning_rate": 1.717395314517617e-05,
+      "loss": 0.0366,
+      "step": 1370
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 1.4070944899475468,
+      "learning_rate": 1.713705958310275e-05,
+      "loss": 0.0364,
+      "step": 1380
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 1.368387839772382,
+      "learning_rate": 1.710016602102933e-05,
+      "loss": 0.0302,
+      "step": 1390
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 1.2202288141019488,
+      "learning_rate": 1.7063272458955912e-05,
+      "loss": 0.0455,
+      "step": 1400
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 1.7419694841079256,
+      "learning_rate": 1.7026378896882496e-05,
+      "loss": 0.0402,
+      "step": 1410
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 1.5792094366611027,
+      "learning_rate": 1.6989485334809077e-05,
+      "loss": 0.0364,
+      "step": 1420
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 1.0092569099059323,
+      "learning_rate": 1.695259177273566e-05,
+      "loss": 0.0335,
+      "step": 1430
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.5374731387141227,
+      "learning_rate": 1.691569821066224e-05,
+      "loss": 0.0428,
+      "step": 1440
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.6268956058587998,
+      "learning_rate": 1.687880464858882e-05,
+      "loss": 0.0483,
+      "step": 1450
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 1.6628348966888566,
+      "learning_rate": 1.6841911086515402e-05,
+      "loss": 0.0487,
+      "step": 1460
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 8.037103442054311,
+      "learning_rate": 1.6805017524441987e-05,
+      "loss": 0.0679,
+      "step": 1470
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 1.620270154467724,
+      "learning_rate": 1.6768123962368568e-05,
+      "loss": 0.0437,
+      "step": 1480
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 1.0436875513091504,
+      "learning_rate": 1.6731230400295152e-05,
+      "loss": 0.0346,
+      "step": 1490
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.8374706362800531,
+      "learning_rate": 1.6694336838221733e-05,
+      "loss": 0.0517,
+      "step": 1500
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 6024,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 4,
+  "save_steps": 500,
+  "total_flos": 0.0,
+  "train_batch_size": 1,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-1506/training_args.bin

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

checkpoint-1506/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)

checkpoint-3012/1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+  "word_embedding_dimension": 1024,
+  "pooling_mode_cls_token": true,
+  "pooling_mode_mean_tokens": false,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false,
+  "pooling_mode_weightedmean_tokens": false,
+  "pooling_mode_lasttoken": false,
+  "include_prompt": true
+}

checkpoint-3012/README.md

@@ -0,0 +1,57 @@
+---
+library_name: sentence-transformers
+pipeline_tag: sentence-similarity
+tags:
+- sentence-transformers
+- feature-extraction
+- sentence-similarity
+
+---
+
+# {MODEL_NAME}
+
+This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 1024 dimensional dense vector space and can be used for tasks like clustering or semantic search.
+
+<!--- Describe your model here -->
+
+## Usage (Sentence-Transformers)
+
+Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
+
+```
+pip install -U sentence-transformers
+```
+
+Then you can use the model like this:
+
+```python
+from sentence_transformers import SentenceTransformer
+sentences = ["This is an example sentence", "Each sentence is converted"]
+
+model = SentenceTransformer('{MODEL_NAME}')
+embeddings = model.encode(sentences)
+print(embeddings)
+```
+
+
+
+## Evaluation Results
+
+<!--- Describe how your model was evaluated -->
+
+For an automated evaluation of this model, see the *Sentence Embeddings Benchmark*: [https://seb.sbert.net](https://seb.sbert.net?model_name={MODEL_NAME})
+
+
+
+## Full Model Architecture
+```
+SentenceTransformer(
+  (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False}) with Transformer model: XLMRobertaModel 
+  (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
+  (2): Normalize()
+)
+```
+
+## Citing & Authors
+
+<!--- Describe where people can find more information -->

checkpoint-3012/colbert_linear.pt

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

checkpoint-3012/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "saved_models/bgem3_unified_finetune_20240330/checkpoint-3012",
+  "architectures": [
+    "XLMRobertaModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 8194,
+  "model_type": "xlm-roberta",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.39.1",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 250002
+}

checkpoint-3012/config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.6.0",
+    "transformers": "4.39.1",
+    "pytorch": "2.0.1+cu117"
+  },
+  "prompts": {},
+  "default_prompt_name": null
+}

checkpoint-3012/global_step3012/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt

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

checkpoint-3012/global_step3012/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt

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

checkpoint-3012/global_step3012/mp_rank_00_model_states.pt

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

checkpoint-3012/latest

@@ -0,0 +1 @@
+global_step3012

checkpoint-3012/model.safetensors

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

checkpoint-3012/modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.models.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.models.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.models.Normalize"
+  }
+]

checkpoint-3012/rng_state_0.pth

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

checkpoint-3012/rng_state_1.pth

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

checkpoint-3012/sentence_bert_config.json

@@ -0,0 +1,4 @@
+{
+  "max_seq_length": 8192,
+  "do_lower_case": false
+}

checkpoint-3012/sentencepiece.bpe.model

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

checkpoint-3012/sparse_linear.pt

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

checkpoint-3012/special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint-3012/tokenizer.json

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

checkpoint-3012/tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "250001": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_max_length": 8192,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}

checkpoint-3012/trainer_state.json

@@ -0,0 +1,2128 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.0,
+  "eval_steps": 500,
+  "global_step": 3012,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.01,
+      "grad_norm": 41.35739895402257,
+      "learning_rate": 7.193423539345941e-06,
+      "loss": 0.5141,
+      "step": 10
+    },
+    {
+      "epoch": 0.01,
+      "grad_norm": 10.443694874625423,
+      "learning_rate": 9.358859796204429e-06,
+      "loss": 0.4195,
+      "step": 20
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 5.720021591145773,
+      "learning_rate": 1.0625558804168632e-05,
+      "loss": 0.2851,
+      "step": 30
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 7.59250257002781,
+      "learning_rate": 1.1524296053062918e-05,
+      "loss": 0.2174,
+      "step": 40
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 11.521550428114184,
+      "learning_rate": 1.2221410821833392e-05,
+      "loss": 0.1817,
+      "step": 50
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 6.027998448017936,
+      "learning_rate": 1.2790995061027121e-05,
+      "loss": 0.1886,
+      "step": 60
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 4.929816146115093,
+      "learning_rate": 1.3272571673439616e-05,
+      "loss": 0.1553,
+      "step": 70
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 2.3881515971871345,
+      "learning_rate": 1.3689732309921406e-05,
+      "loss": 0.129,
+      "step": 80
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 5.441118636680569,
+      "learning_rate": 1.4057694068991321e-05,
+      "loss": 0.1433,
+      "step": 90
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 6.085213496930297,
+      "learning_rate": 1.4386847078691883e-05,
+      "loss": 0.1092,
+      "step": 100
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 4.713361075579101,
+      "learning_rate": 1.4684602194465794e-05,
+      "loss": 0.1231,
+      "step": 110
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 3.8259500358417924,
+      "learning_rate": 1.495643131788561e-05,
+      "loss": 0.0697,
+      "step": 120
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 2.67920682727699,
+      "learning_rate": 1.5206489871327869e-05,
+      "loss": 0.084,
+      "step": 130
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 8.189491553913532,
+      "learning_rate": 1.54380079302981e-05,
+      "loss": 0.1023,
+      "step": 140
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 12.000369384166694,
+      "learning_rate": 1.5653546086656083e-05,
+      "loss": 0.0972,
+      "step": 150
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 2.3726942012239953,
+      "learning_rate": 1.5855168566779895e-05,
+      "loss": 0.1036,
+      "step": 160
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 2.0242117206977044,
+      "learning_rate": 1.604456377435124e-05,
+      "loss": 0.1081,
+      "step": 170
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 2.2261934744219967,
+      "learning_rate": 1.6223130325849813e-05,
+      "loss": 0.082,
+      "step": 180
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 4.70421073268978,
+      "learning_rate": 1.6392039793463407e-05,
+      "loss": 0.0744,
+      "step": 190
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 12.064244277626111,
+      "learning_rate": 1.6552283335550368e-05,
+      "loss": 0.0934,
+      "step": 200
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 2.150990470911233,
+      "learning_rate": 1.67047069382623e-05,
+      "loss": 0.0737,
+      "step": 210
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 4.074258308547766,
+      "learning_rate": 1.6850038451324284e-05,
+      "loss": 0.0841,
+      "step": 220
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 2.6764310841274237,
+      "learning_rate": 1.6988908609137504e-05,
+      "loss": 0.0821,
+      "step": 230
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 3.1391283710987063,
+      "learning_rate": 1.71218675747441e-05,
+      "loss": 0.0747,
+      "step": 240
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 3.1103728720325754,
+      "learning_rate": 1.7249398104320845e-05,
+      "loss": 0.0907,
+      "step": 250
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 5.354173276825422,
+      "learning_rate": 1.7371926128186358e-05,
+      "loss": 0.0765,
+      "step": 260
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 2.8677808267225005,
+      "learning_rate": 1.7489829333814013e-05,
+      "loss": 0.0777,
+      "step": 270
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 2.6124301773304874,
+      "learning_rate": 1.760344418715659e-05,
+      "loss": 0.0897,
+      "step": 280
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 2.915238130503912,
+      "learning_rate": 1.7713071721324668e-05,
+      "loss": 0.0866,
+      "step": 290
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 4.061775729833908,
+      "learning_rate": 1.781898234351457e-05,
+      "loss": 0.0804,
+      "step": 300
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 1.5648704857906242,
+      "learning_rate": 1.7921419853452233e-05,
+      "loss": 0.0834,
+      "step": 310
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 2.225837991191409,
+      "learning_rate": 1.8020604823638384e-05,
+      "loss": 0.0646,
+      "step": 320
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 4.041622096074851,
+      "learning_rate": 1.811673745928848e-05,
+      "loss": 0.0636,
+      "step": 330
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.21100008000919218,
+      "learning_rate": 1.821000003120973e-05,
+      "loss": 0.0656,
+      "step": 340
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 2.2923304434569443,
+      "learning_rate": 1.8300558955927067e-05,
+      "loss": 0.0756,
+      "step": 350
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 3.6524729986322013,
+      "learning_rate": 1.83885665827083e-05,
+      "loss": 0.0466,
+      "step": 360
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 2.1793660025907284,
+      "learning_rate": 1.847416273569235e-05,
+      "loss": 0.0704,
+      "step": 370
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 4.363503091034068,
+      "learning_rate": 1.8557476050321896e-05,
+      "loss": 0.0712,
+      "step": 380
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 3.0942807753734476,
+      "learning_rate": 1.863862513615056e-05,
+      "loss": 0.0831,
+      "step": 390
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 1.789230871814937,
+      "learning_rate": 1.8717719592408857e-05,
+      "loss": 0.078,
+      "step": 400
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 3.4361028432853677,
+      "learning_rate": 1.879486089815082e-05,
+      "loss": 0.0663,
+      "step": 410
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 3.690439378170989,
+      "learning_rate": 1.8870143195120794e-05,
+      "loss": 0.0738,
+      "step": 420
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 2.159404934400987,
+      "learning_rate": 1.8943653978491198e-05,
+      "loss": 0.0768,
+      "step": 430
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 0.8042070009727055,
+      "learning_rate": 1.901547470818277e-05,
+      "loss": 0.0777,
+      "step": 440
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 1.70758347995601,
+      "learning_rate": 1.9085681351478775e-05,
+      "loss": 0.05,
+      "step": 450
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 1.2113560599237037,
+      "learning_rate": 1.9154344865995993e-05,
+      "loss": 0.0715,
+      "step": 460
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 2.5816071085486527,
+      "learning_rate": 1.9221531630710657e-05,
+      "loss": 0.0688,
+      "step": 470
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 1.4447841193953872,
+      "learning_rate": 1.9287303831602588e-05,
+      "loss": 0.0659,
+      "step": 480
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 1.2668557011358101,
+      "learning_rate": 1.9351719807533285e-05,
+      "loss": 0.0515,
+      "step": 490
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 2.0353893651728865,
+      "learning_rate": 1.9414834361179333e-05,
+      "loss": 0.0687,
+      "step": 500
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 2.005863891715338,
+      "learning_rate": 1.947669903917393e-05,
+      "loss": 0.0459,
+      "step": 510
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.4330303529877615,
+      "learning_rate": 1.9537362385044847e-05,
+      "loss": 0.0557,
+      "step": 520
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 1.8683959370360874,
+      "learning_rate": 1.959687016805845e-05,
+      "loss": 0.0656,
+      "step": 530
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 1.5963373312962135,
+      "learning_rate": 1.9655265590672502e-05,
+      "loss": 0.0519,
+      "step": 540
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 1.824110595633705,
+      "learning_rate": 1.9712589476953243e-05,
+      "loss": 0.0557,
+      "step": 550
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 2.024299349665376,
+      "learning_rate": 1.976888044401508e-05,
+      "loss": 0.0617,
+      "step": 560
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 1.0490238284158413,
+      "learning_rate": 1.98241750582861e-05,
+      "loss": 0.0529,
+      "step": 570
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.2602675746455845,
+      "learning_rate": 1.9878507978183157e-05,
+      "loss": 0.0617,
+      "step": 580
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 1.7353304044327778,
+      "learning_rate": 1.9931912084590654e-05,
+      "loss": 0.0464,
+      "step": 590
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 2.043608410331547,
+      "learning_rate": 1.998441860037306e-05,
+      "loss": 0.0716,
+      "step": 600
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 0.9522735357022469,
+      "learning_rate": 1.997786386275595e-05,
+      "loss": 0.0496,
+      "step": 610
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 1.3740017908881452,
+      "learning_rate": 1.9940970300682533e-05,
+      "loss": 0.0686,
+      "step": 620
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.1533525092583308,
+      "learning_rate": 1.9904076738609114e-05,
+      "loss": 0.0562,
+      "step": 630
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 1.564671647601224,
+      "learning_rate": 1.9867183176535695e-05,
+      "loss": 0.0553,
+      "step": 640
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 0.8996693997352794,
+      "learning_rate": 1.9830289614462276e-05,
+      "loss": 0.0524,
+      "step": 650
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.9340857910489522,
+      "learning_rate": 1.979339605238886e-05,
+      "loss": 0.0534,
+      "step": 660
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 1.892682744806536,
+      "learning_rate": 1.9756502490315442e-05,
+      "loss": 0.0625,
+      "step": 670
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 1.2063853192894494,
+      "learning_rate": 1.9719608928242023e-05,
+      "loss": 0.0519,
+      "step": 680
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.984312822464147,
+      "learning_rate": 1.9682715366168604e-05,
+      "loss": 0.0494,
+      "step": 690
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 1.6705446030280595,
+      "learning_rate": 1.9645821804095185e-05,
+      "loss": 0.059,
+      "step": 700
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 1.7422620817223426,
+      "learning_rate": 1.960892824202177e-05,
+      "loss": 0.0454,
+      "step": 710
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.3182961229762868,
+      "learning_rate": 1.957203467994835e-05,
+      "loss": 0.0581,
+      "step": 720
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 1.330360819481426,
+      "learning_rate": 1.9535141117874932e-05,
+      "loss": 0.0575,
+      "step": 730
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 1.9840741868184866,
+      "learning_rate": 1.9498247555801517e-05,
+      "loss": 0.063,
+      "step": 740
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.6018064731760029,
+      "learning_rate": 1.9461353993728094e-05,
+      "loss": 0.0519,
+      "step": 750
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 1.391561342963203,
+      "learning_rate": 1.9424460431654675e-05,
+      "loss": 0.0572,
+      "step": 760
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 1.0319629863193043,
+      "learning_rate": 1.938756686958126e-05,
+      "loss": 0.0619,
+      "step": 770
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.7398111190442345,
+      "learning_rate": 1.935067330750784e-05,
+      "loss": 0.05,
+      "step": 780
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 1.675180749962208,
+      "learning_rate": 1.9313779745434422e-05,
+      "loss": 0.0654,
+      "step": 790
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 1.486943722740635,
+      "learning_rate": 1.9276886183361007e-05,
+      "loss": 0.0463,
+      "step": 800
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1444156835752686,
+      "learning_rate": 1.9239992621287588e-05,
+      "loss": 0.0598,
+      "step": 810
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 1.1260777714975718,
+      "learning_rate": 1.920309905921417e-05,
+      "loss": 0.039,
+      "step": 820
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 1.328313971146618,
+      "learning_rate": 1.916620549714075e-05,
+      "loss": 0.0591,
+      "step": 830
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 0.8805517771854091,
+      "learning_rate": 1.912931193506733e-05,
+      "loss": 0.0392,
+      "step": 840
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 0.9704226543415952,
+      "learning_rate": 1.9092418372993916e-05,
+      "loss": 0.056,
+      "step": 850
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 1.4328262810995938,
+      "learning_rate": 1.9055524810920497e-05,
+      "loss": 0.0602,
+      "step": 860
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 2.6133255890647167,
+      "learning_rate": 1.9018631248847078e-05,
+      "loss": 0.0507,
+      "step": 870
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 1.4574352859865667,
+      "learning_rate": 1.898173768677366e-05,
+      "loss": 0.0667,
+      "step": 880
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 1.9116869451062006,
+      "learning_rate": 1.894484412470024e-05,
+      "loss": 0.0533,
+      "step": 890
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 0.8152499611885836,
+      "learning_rate": 1.890795056262682e-05,
+      "loss": 0.0534,
+      "step": 900
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 2.9440345318883696,
+      "learning_rate": 1.8871057000553406e-05,
+      "loss": 0.0419,
+      "step": 910
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 2.7588896552457074,
+      "learning_rate": 1.8834163438479987e-05,
+      "loss": 0.0418,
+      "step": 920
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.8886319208504987,
+      "learning_rate": 1.8797269876406568e-05,
+      "loss": 0.0397,
+      "step": 930
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 1.8820615671771377,
+      "learning_rate": 1.876037631433315e-05,
+      "loss": 0.0454,
+      "step": 940
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 1.2868214865586614,
+      "learning_rate": 1.872348275225973e-05,
+      "loss": 0.0373,
+      "step": 950
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 2.016876778284122,
+      "learning_rate": 1.8686589190186315e-05,
+      "loss": 0.0619,
+      "step": 960
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 0.8992764743601639,
+      "learning_rate": 1.8649695628112896e-05,
+      "loss": 0.0403,
+      "step": 970
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 0.42416507318131186,
+      "learning_rate": 1.8612802066039477e-05,
+      "loss": 0.067,
+      "step": 980
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.8601423331120188,
+      "learning_rate": 1.8575908503966062e-05,
+      "loss": 0.066,
+      "step": 990
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 1.3893250379275068,
+      "learning_rate": 1.8539014941892643e-05,
+      "loss": 0.0437,
+      "step": 1000
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 2.3240219930504376,
+      "learning_rate": 1.850212137981922e-05,
+      "loss": 0.0488,
+      "step": 1010
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.8849896953543195,
+      "learning_rate": 1.8465227817745805e-05,
+      "loss": 0.0472,
+      "step": 1020
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.8655212477880712,
+      "learning_rate": 1.8428334255672386e-05,
+      "loss": 0.0546,
+      "step": 1030
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 1.2016665548799828,
+      "learning_rate": 1.8391440693598967e-05,
+      "loss": 0.0956,
+      "step": 1040
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 1.492659569010543,
+      "learning_rate": 1.8354547131525552e-05,
+      "loss": 0.0449,
+      "step": 1050
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 0.4224837643985645,
+      "learning_rate": 1.8317653569452133e-05,
+      "loss": 0.0545,
+      "step": 1060
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 1.8405573730480247,
+      "learning_rate": 1.8280760007378714e-05,
+      "loss": 0.0328,
+      "step": 1070
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.3988661237020132,
+      "learning_rate": 1.8243866445305295e-05,
+      "loss": 0.0451,
+      "step": 1080
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.5573396543240564,
+      "learning_rate": 1.8206972883231876e-05,
+      "loss": 0.0502,
+      "step": 1090
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.7855566721285819,
+      "learning_rate": 1.817007932115846e-05,
+      "loss": 0.0341,
+      "step": 1100
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 1.497005175064917,
+      "learning_rate": 1.8133185759085042e-05,
+      "loss": 0.0513,
+      "step": 1110
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 1.650975953086994,
+      "learning_rate": 1.8096292197011623e-05,
+      "loss": 0.0473,
+      "step": 1120
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.8418054071334755,
+      "learning_rate": 1.8059398634938204e-05,
+      "loss": 0.0444,
+      "step": 1130
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 1.33651471077315,
+      "learning_rate": 1.8022505072864785e-05,
+      "loss": 0.0476,
+      "step": 1140
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 1.1529220423121023,
+      "learning_rate": 1.7985611510791367e-05,
+      "loss": 0.0379,
+      "step": 1150
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 0.6912013035706749,
+      "learning_rate": 1.794871794871795e-05,
+      "loss": 0.0526,
+      "step": 1160
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 2.806793867054328,
+      "learning_rate": 1.7911824386644532e-05,
+      "loss": 0.047,
+      "step": 1170
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 3.323409487413899,
+      "learning_rate": 1.7874930824571113e-05,
+      "loss": 0.0455,
+      "step": 1180
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 1.0877783561592873,
+      "learning_rate": 1.7838037262497695e-05,
+      "loss": 0.0555,
+      "step": 1190
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.9158165193439103,
+      "learning_rate": 1.7801143700424276e-05,
+      "loss": 0.0458,
+      "step": 1200
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 1.0863638939536142,
+      "learning_rate": 1.776425013835086e-05,
+      "loss": 0.0462,
+      "step": 1210
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 2.3398031954044978,
+      "learning_rate": 1.772735657627744e-05,
+      "loss": 0.032,
+      "step": 1220
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.9272917966743286,
+      "learning_rate": 1.7690463014204022e-05,
+      "loss": 0.0464,
+      "step": 1230
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 1.4608044079217914,
+      "learning_rate": 1.7653569452130607e-05,
+      "loss": 0.0544,
+      "step": 1240
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 2.0799767984416713,
+      "learning_rate": 1.7616675890057188e-05,
+      "loss": 0.0629,
+      "step": 1250
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 2.438242481318271,
+      "learning_rate": 1.757978232798377e-05,
+      "loss": 0.0586,
+      "step": 1260
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.22666963149041416,
+      "learning_rate": 1.754288876591035e-05,
+      "loss": 0.0419,
+      "step": 1270
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 1.6906855856145777,
+      "learning_rate": 1.750599520383693e-05,
+      "loss": 0.0417,
+      "step": 1280
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 0.8417294291584847,
+      "learning_rate": 1.7469101641763513e-05,
+      "loss": 0.0496,
+      "step": 1290
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 2.384648750559231,
+      "learning_rate": 1.7432208079690097e-05,
+      "loss": 0.037,
+      "step": 1300
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 1.0969098360012755,
+      "learning_rate": 1.7395314517616678e-05,
+      "loss": 0.0392,
+      "step": 1310
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 4.265947825753472,
+      "learning_rate": 1.735842095554326e-05,
+      "loss": 0.0434,
+      "step": 1320
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 1.0043877066833349,
+      "learning_rate": 1.732152739346984e-05,
+      "loss": 0.0423,
+      "step": 1330
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 0.3834891081373628,
+      "learning_rate": 1.728463383139642e-05,
+      "loss": 0.0423,
+      "step": 1340
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.8917610552257311,
+      "learning_rate": 1.7247740269323006e-05,
+      "loss": 0.0491,
+      "step": 1350
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 1.7777935423085822,
+      "learning_rate": 1.7210846707249587e-05,
+      "loss": 0.0465,
+      "step": 1360
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 2.882686457165806,
+      "learning_rate": 1.717395314517617e-05,
+      "loss": 0.0366,
+      "step": 1370
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 1.4070944899475468,
+      "learning_rate": 1.713705958310275e-05,
+      "loss": 0.0364,
+      "step": 1380
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 1.368387839772382,
+      "learning_rate": 1.710016602102933e-05,
+      "loss": 0.0302,
+      "step": 1390
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 1.2202288141019488,
+      "learning_rate": 1.7063272458955912e-05,
+      "loss": 0.0455,
+      "step": 1400
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 1.7419694841079256,
+      "learning_rate": 1.7026378896882496e-05,
+      "loss": 0.0402,
+      "step": 1410
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 1.5792094366611027,
+      "learning_rate": 1.6989485334809077e-05,
+      "loss": 0.0364,
+      "step": 1420
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 1.0092569099059323,
+      "learning_rate": 1.695259177273566e-05,
+      "loss": 0.0335,
+      "step": 1430
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.5374731387141227,
+      "learning_rate": 1.691569821066224e-05,
+      "loss": 0.0428,
+      "step": 1440
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 1.6268956058587998,
+      "learning_rate": 1.687880464858882e-05,
+      "loss": 0.0483,
+      "step": 1450
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 1.6628348966888566,
+      "learning_rate": 1.6841911086515402e-05,
+      "loss": 0.0487,
+      "step": 1460
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 8.037103442054311,
+      "learning_rate": 1.6805017524441987e-05,
+      "loss": 0.0679,
+      "step": 1470
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 1.620270154467724,
+      "learning_rate": 1.6768123962368568e-05,
+      "loss": 0.0437,
+      "step": 1480
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 1.0436875513091504,
+      "learning_rate": 1.6731230400295152e-05,
+      "loss": 0.0346,
+      "step": 1490
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.8374706362800531,
+      "learning_rate": 1.6694336838221733e-05,
+      "loss": 0.0517,
+      "step": 1500
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 1.1939983249441695,
+      "learning_rate": 1.6657443276148314e-05,
+      "loss": 0.0393,
+      "step": 1510
+    },
+    {
+      "epoch": 1.01,
+      "grad_norm": 1.0650702970882078,
+      "learning_rate": 1.6620549714074896e-05,
+      "loss": 0.0433,
+      "step": 1520
+    },
+    {
+      "epoch": 1.02,
+      "grad_norm": 0.9588706576483126,
+      "learning_rate": 1.6583656152001477e-05,
+      "loss": 0.0381,
+      "step": 1530
+    },
+    {
+      "epoch": 1.02,
+      "grad_norm": 1.4412233810414918,
+      "learning_rate": 1.6546762589928058e-05,
+      "loss": 0.0346,
+      "step": 1540
+    },
+    {
+      "epoch": 1.03,
+      "grad_norm": 0.9533871774762268,
+      "learning_rate": 1.6509869027854642e-05,
+      "loss": 0.0239,
+      "step": 1550
+    },
+    {
+      "epoch": 1.04,
+      "grad_norm": 0.44748765523495393,
+      "learning_rate": 1.6472975465781223e-05,
+      "loss": 0.0424,
+      "step": 1560
+    },
+    {
+      "epoch": 1.04,
+      "grad_norm": 0.6542278196418172,
+      "learning_rate": 1.6436081903707805e-05,
+      "loss": 0.0322,
+      "step": 1570
+    },
+    {
+      "epoch": 1.05,
+      "grad_norm": 0.774821420829752,
+      "learning_rate": 1.6399188341634386e-05,
+      "loss": 0.0337,
+      "step": 1580
+    },
+    {
+      "epoch": 1.06,
+      "grad_norm": 3.4366984039627053,
+      "learning_rate": 1.6362294779560967e-05,
+      "loss": 0.0402,
+      "step": 1590
+    },
+    {
+      "epoch": 1.06,
+      "grad_norm": 0.9409785268524657,
+      "learning_rate": 1.6325401217487548e-05,
+      "loss": 0.0351,
+      "step": 1600
+    },
+    {
+      "epoch": 1.07,
+      "grad_norm": 1.3755255908395096,
+      "learning_rate": 1.6288507655414133e-05,
+      "loss": 0.0287,
+      "step": 1610
+    },
+    {
+      "epoch": 1.08,
+      "grad_norm": 0.9105883051718862,
+      "learning_rate": 1.6251614093340714e-05,
+      "loss": 0.0377,
+      "step": 1620
+    },
+    {
+      "epoch": 1.08,
+      "grad_norm": 0.669451044858242,
+      "learning_rate": 1.6214720531267295e-05,
+      "loss": 0.0397,
+      "step": 1630
+    },
+    {
+      "epoch": 1.09,
+      "grad_norm": 1.0173498911812129,
+      "learning_rate": 1.6177826969193876e-05,
+      "loss": 0.0362,
+      "step": 1640
+    },
+    {
+      "epoch": 1.1,
+      "grad_norm": 1.443726695479527,
+      "learning_rate": 1.6140933407120457e-05,
+      "loss": 0.0388,
+      "step": 1650
+    },
+    {
+      "epoch": 1.1,
+      "grad_norm": 0.48614710000952543,
+      "learning_rate": 1.610403984504704e-05,
+      "loss": 0.044,
+      "step": 1660
+    },
+    {
+      "epoch": 1.11,
+      "grad_norm": 0.8955875806318787,
+      "learning_rate": 1.6067146282973623e-05,
+      "loss": 0.0361,
+      "step": 1670
+    },
+    {
+      "epoch": 1.12,
+      "grad_norm": 0.7136011214906065,
+      "learning_rate": 1.6030252720900204e-05,
+      "loss": 0.0399,
+      "step": 1680
+    },
+    {
+      "epoch": 1.12,
+      "grad_norm": 1.342383821315822,
+      "learning_rate": 1.599335915882679e-05,
+      "loss": 0.0398,
+      "step": 1690
+    },
+    {
+      "epoch": 1.13,
+      "grad_norm": 1.9154366056765233,
+      "learning_rate": 1.5956465596753366e-05,
+      "loss": 0.0351,
+      "step": 1700
+    },
+    {
+      "epoch": 1.14,
+      "grad_norm": 3.380571358597728,
+      "learning_rate": 1.5919572034679947e-05,
+      "loss": 0.0355,
+      "step": 1710
+    },
+    {
+      "epoch": 1.14,
+      "grad_norm": 1.4160947580259224,
+      "learning_rate": 1.5882678472606532e-05,
+      "loss": 0.0333,
+      "step": 1720
+    },
+    {
+      "epoch": 1.15,
+      "grad_norm": 0.6737638581076882,
+      "learning_rate": 1.5845784910533113e-05,
+      "loss": 0.0296,
+      "step": 1730
+    },
+    {
+      "epoch": 1.16,
+      "grad_norm": 1.206337590570484,
+      "learning_rate": 1.5808891348459694e-05,
+      "loss": 0.0308,
+      "step": 1740
+    },
+    {
+      "epoch": 1.16,
+      "grad_norm": 0.7350933133142535,
+      "learning_rate": 1.577199778638628e-05,
+      "loss": 0.0335,
+      "step": 1750
+    },
+    {
+      "epoch": 1.17,
+      "grad_norm": 1.3603619905771789,
+      "learning_rate": 1.573510422431286e-05,
+      "loss": 0.0254,
+      "step": 1760
+    },
+    {
+      "epoch": 1.18,
+      "grad_norm": 0.6337146921122443,
+      "learning_rate": 1.569821066223944e-05,
+      "loss": 0.0381,
+      "step": 1770
+    },
+    {
+      "epoch": 1.18,
+      "grad_norm": 1.3833442570741445,
+      "learning_rate": 1.5661317100166022e-05,
+      "loss": 0.0295,
+      "step": 1780
+    },
+    {
+      "epoch": 1.19,
+      "grad_norm": 1.1660396503165016,
+      "learning_rate": 1.5624423538092603e-05,
+      "loss": 0.0313,
+      "step": 1790
+    },
+    {
+      "epoch": 1.2,
+      "grad_norm": 1.085393536126589,
+      "learning_rate": 1.5587529976019188e-05,
+      "loss": 0.0294,
+      "step": 1800
+    },
+    {
+      "epoch": 1.2,
+      "grad_norm": 0.9428955580680065,
+      "learning_rate": 1.555063641394577e-05,
+      "loss": 0.0403,
+      "step": 1810
+    },
+    {
+      "epoch": 1.21,
+      "grad_norm": 1.2824389825772662,
+      "learning_rate": 1.551374285187235e-05,
+      "loss": 0.0326,
+      "step": 1820
+    },
+    {
+      "epoch": 1.22,
+      "grad_norm": 1.071530839035684,
+      "learning_rate": 1.547684928979893e-05,
+      "loss": 0.0224,
+      "step": 1830
+    },
+    {
+      "epoch": 1.22,
+      "grad_norm": 0.7250437295813942,
+      "learning_rate": 1.5439955727725512e-05,
+      "loss": 0.0253,
+      "step": 1840
+    },
+    {
+      "epoch": 1.23,
+      "grad_norm": 0.4676245052117833,
+      "learning_rate": 1.5403062165652093e-05,
+      "loss": 0.0245,
+      "step": 1850
+    },
+    {
+      "epoch": 1.24,
+      "grad_norm": 1.3654237978697314,
+      "learning_rate": 1.5366168603578678e-05,
+      "loss": 0.0274,
+      "step": 1860
+    },
+    {
+      "epoch": 1.24,
+      "grad_norm": 0.8371654653693097,
+      "learning_rate": 1.532927504150526e-05,
+      "loss": 0.0408,
+      "step": 1870
+    },
+    {
+      "epoch": 1.25,
+      "grad_norm": 0.918371136718616,
+      "learning_rate": 1.529238147943184e-05,
+      "loss": 0.0252,
+      "step": 1880
+    },
+    {
+      "epoch": 1.25,
+      "grad_norm": 1.0100840319244568,
+      "learning_rate": 1.5255487917358423e-05,
+      "loss": 0.0282,
+      "step": 1890
+    },
+    {
+      "epoch": 1.26,
+      "grad_norm": 1.2752226504227586,
+      "learning_rate": 1.5218594355285004e-05,
+      "loss": 0.0481,
+      "step": 1900
+    },
+    {
+      "epoch": 1.27,
+      "grad_norm": 1.016628356639124,
+      "learning_rate": 1.5181700793211587e-05,
+      "loss": 0.0534,
+      "step": 1910
+    },
+    {
+      "epoch": 1.27,
+      "grad_norm": 1.2018849311057271,
+      "learning_rate": 1.5144807231138168e-05,
+      "loss": 0.0357,
+      "step": 1920
+    },
+    {
+      "epoch": 1.28,
+      "grad_norm": 1.3543149671289594,
+      "learning_rate": 1.5107913669064749e-05,
+      "loss": 0.0354,
+      "step": 1930
+    },
+    {
+      "epoch": 1.29,
+      "grad_norm": 1.0460241475499934,
+      "learning_rate": 1.5071020106991332e-05,
+      "loss": 0.0328,
+      "step": 1940
+    },
+    {
+      "epoch": 1.29,
+      "grad_norm": 0.7712993931217346,
+      "learning_rate": 1.5034126544917913e-05,
+      "loss": 0.0352,
+      "step": 1950
+    },
+    {
+      "epoch": 1.3,
+      "grad_norm": 0.6403552627247137,
+      "learning_rate": 1.4997232982844494e-05,
+      "loss": 0.0374,
+      "step": 1960
+    },
+    {
+      "epoch": 1.31,
+      "grad_norm": 0.6985534187744373,
+      "learning_rate": 1.4960339420771077e-05,
+      "loss": 0.0236,
+      "step": 1970
+    },
+    {
+      "epoch": 1.31,
+      "grad_norm": 2.7973859948867315,
+      "learning_rate": 1.4923445858697658e-05,
+      "loss": 0.0333,
+      "step": 1980
+    },
+    {
+      "epoch": 1.32,
+      "grad_norm": 0.7856751051336177,
+      "learning_rate": 1.488655229662424e-05,
+      "loss": 0.0283,
+      "step": 1990
+    },
+    {
+      "epoch": 1.33,
+      "grad_norm": 1.1090831103953211,
+      "learning_rate": 1.4849658734550822e-05,
+      "loss": 0.0424,
+      "step": 2000
+    },
+    {
+      "epoch": 1.33,
+      "grad_norm": 0.37705371680669364,
+      "learning_rate": 1.4812765172477403e-05,
+      "loss": 0.0264,
+      "step": 2010
+    },
+    {
+      "epoch": 1.34,
+      "grad_norm": 2.011193423521568,
+      "learning_rate": 1.4775871610403986e-05,
+      "loss": 0.0325,
+      "step": 2020
+    },
+    {
+      "epoch": 1.35,
+      "grad_norm": 1.00662492673962,
+      "learning_rate": 1.4738978048330567e-05,
+      "loss": 0.0325,
+      "step": 2030
+    },
+    {
+      "epoch": 1.35,
+      "grad_norm": 1.0249189181939053,
+      "learning_rate": 1.4702084486257148e-05,
+      "loss": 0.0296,
+      "step": 2040
+    },
+    {
+      "epoch": 1.36,
+      "grad_norm": 1.294162625472094,
+      "learning_rate": 1.4665190924183733e-05,
+      "loss": 0.0362,
+      "step": 2050
+    },
+    {
+      "epoch": 1.37,
+      "grad_norm": 0.8753798886295676,
+      "learning_rate": 1.4628297362110312e-05,
+      "loss": 0.0352,
+      "step": 2060
+    },
+    {
+      "epoch": 1.37,
+      "grad_norm": 0.32404034083576144,
+      "learning_rate": 1.4591403800036893e-05,
+      "loss": 0.0319,
+      "step": 2070
+    },
+    {
+      "epoch": 1.38,
+      "grad_norm": 1.3881694244187606,
+      "learning_rate": 1.4554510237963478e-05,
+      "loss": 0.0383,
+      "step": 2080
+    },
+    {
+      "epoch": 1.39,
+      "grad_norm": 0.9106958124280001,
+      "learning_rate": 1.4517616675890059e-05,
+      "loss": 0.0366,
+      "step": 2090
+    },
+    {
+      "epoch": 1.39,
+      "grad_norm": 1.3920757998129725,
+      "learning_rate": 1.448072311381664e-05,
+      "loss": 0.0359,
+      "step": 2100
+    },
+    {
+      "epoch": 1.4,
+      "grad_norm": 2.5378193320802556,
+      "learning_rate": 1.4443829551743223e-05,
+      "loss": 0.0375,
+      "step": 2110
+    },
+    {
+      "epoch": 1.41,
+      "grad_norm": 1.6327275888367767,
+      "learning_rate": 1.4406935989669804e-05,
+      "loss": 0.0491,
+      "step": 2120
+    },
+    {
+      "epoch": 1.41,
+      "grad_norm": 1.457863332305797,
+      "learning_rate": 1.4370042427596385e-05,
+      "loss": 0.0379,
+      "step": 2130
+    },
+    {
+      "epoch": 1.42,
+      "grad_norm": 0.19454376571103232,
+      "learning_rate": 1.4333148865522968e-05,
+      "loss": 0.0357,
+      "step": 2140
+    },
+    {
+      "epoch": 1.43,
+      "grad_norm": 0.7421971370673671,
+      "learning_rate": 1.429625530344955e-05,
+      "loss": 0.0418,
+      "step": 2150
+    },
+    {
+      "epoch": 1.43,
+      "grad_norm": 0.0937417741981473,
+      "learning_rate": 1.4259361741376132e-05,
+      "loss": 0.0248,
+      "step": 2160
+    },
+    {
+      "epoch": 1.44,
+      "grad_norm": 0.7544900086855553,
+      "learning_rate": 1.4222468179302713e-05,
+      "loss": 0.0242,
+      "step": 2170
+    },
+    {
+      "epoch": 1.45,
+      "grad_norm": 0.10357082442667653,
+      "learning_rate": 1.4185574617229294e-05,
+      "loss": 0.0277,
+      "step": 2180
+    },
+    {
+      "epoch": 1.45,
+      "grad_norm": 0.429204704240934,
+      "learning_rate": 1.4148681055155877e-05,
+      "loss": 0.0253,
+      "step": 2190
+    },
+    {
+      "epoch": 1.46,
+      "grad_norm": 1.7367754983146269,
+      "learning_rate": 1.4111787493082458e-05,
+      "loss": 0.0337,
+      "step": 2200
+    },
+    {
+      "epoch": 1.47,
+      "grad_norm": 0.6428519364176438,
+      "learning_rate": 1.407489393100904e-05,
+      "loss": 0.0331,
+      "step": 2210
+    },
+    {
+      "epoch": 1.47,
+      "grad_norm": 0.7848514705508334,
+      "learning_rate": 1.4038000368935622e-05,
+      "loss": 0.031,
+      "step": 2220
+    },
+    {
+      "epoch": 1.48,
+      "grad_norm": 1.2604412780544694,
+      "learning_rate": 1.4001106806862203e-05,
+      "loss": 0.0381,
+      "step": 2230
+    },
+    {
+      "epoch": 1.49,
+      "grad_norm": 1.0154534415085918,
+      "learning_rate": 1.3964213244788784e-05,
+      "loss": 0.0294,
+      "step": 2240
+    },
+    {
+      "epoch": 1.49,
+      "grad_norm": 1.2166332545112624,
+      "learning_rate": 1.3927319682715367e-05,
+      "loss": 0.0275,
+      "step": 2250
+    },
+    {
+      "epoch": 1.5,
+      "grad_norm": 0.8758369938501386,
+      "learning_rate": 1.3890426120641948e-05,
+      "loss": 0.0359,
+      "step": 2260
+    },
+    {
+      "epoch": 1.51,
+      "grad_norm": 0.9645035178583596,
+      "learning_rate": 1.385353255856853e-05,
+      "loss": 0.0257,
+      "step": 2270
+    },
+    {
+      "epoch": 1.51,
+      "grad_norm": 0.24217577106008117,
+      "learning_rate": 1.3816638996495112e-05,
+      "loss": 0.0188,
+      "step": 2280
+    },
+    {
+      "epoch": 1.52,
+      "grad_norm": 0.8656287403232453,
+      "learning_rate": 1.3779745434421693e-05,
+      "loss": 0.035,
+      "step": 2290
+    },
+    {
+      "epoch": 1.53,
+      "grad_norm": 1.5535880950760363,
+      "learning_rate": 1.3742851872348278e-05,
+      "loss": 0.0325,
+      "step": 2300
+    },
+    {
+      "epoch": 1.53,
+      "grad_norm": 0.6356266461691455,
+      "learning_rate": 1.3705958310274859e-05,
+      "loss": 0.0292,
+      "step": 2310
+    },
+    {
+      "epoch": 1.54,
+      "grad_norm": 1.3067413428110854,
+      "learning_rate": 1.3669064748201439e-05,
+      "loss": 0.0206,
+      "step": 2320
+    },
+    {
+      "epoch": 1.55,
+      "grad_norm": 1.3327319855425117,
+      "learning_rate": 1.3632171186128023e-05,
+      "loss": 0.023,
+      "step": 2330
+    },
+    {
+      "epoch": 1.55,
+      "grad_norm": 2.308037550483557,
+      "learning_rate": 1.3595277624054604e-05,
+      "loss": 0.0283,
+      "step": 2340
+    },
+    {
+      "epoch": 1.56,
+      "grad_norm": 1.447528110026965,
+      "learning_rate": 1.3558384061981185e-05,
+      "loss": 0.0344,
+      "step": 2350
+    },
+    {
+      "epoch": 1.57,
+      "grad_norm": 1.4781377553893515,
+      "learning_rate": 1.3521490499907768e-05,
+      "loss": 0.0338,
+      "step": 2360
+    },
+    {
+      "epoch": 1.57,
+      "grad_norm": 1.2213361639825704,
+      "learning_rate": 1.348459693783435e-05,
+      "loss": 0.0213,
+      "step": 2370
+    },
+    {
+      "epoch": 1.58,
+      "grad_norm": 0.806579624483452,
+      "learning_rate": 1.344770337576093e-05,
+      "loss": 0.0308,
+      "step": 2380
+    },
+    {
+      "epoch": 1.59,
+      "grad_norm": 0.36305855519021074,
+      "learning_rate": 1.3410809813687513e-05,
+      "loss": 0.0307,
+      "step": 2390
+    },
+    {
+      "epoch": 1.59,
+      "grad_norm": 1.7963102033556697,
+      "learning_rate": 1.3373916251614094e-05,
+      "loss": 0.0438,
+      "step": 2400
+    },
+    {
+      "epoch": 1.6,
+      "grad_norm": 0.8741626138535608,
+      "learning_rate": 1.3337022689540676e-05,
+      "loss": 0.0282,
+      "step": 2410
+    },
+    {
+      "epoch": 1.61,
+      "grad_norm": 0.8183941890622122,
+      "learning_rate": 1.3300129127467258e-05,
+      "loss": 0.03,
+      "step": 2420
+    },
+    {
+      "epoch": 1.61,
+      "grad_norm": 1.315171338874376,
+      "learning_rate": 1.326323556539384e-05,
+      "loss": 0.0347,
+      "step": 2430
+    },
+    {
+      "epoch": 1.62,
+      "grad_norm": 1.268446896432008,
+      "learning_rate": 1.3226342003320422e-05,
+      "loss": 0.0376,
+      "step": 2440
+    },
+    {
+      "epoch": 1.63,
+      "grad_norm": 0.913219454382149,
+      "learning_rate": 1.3189448441247003e-05,
+      "loss": 0.0358,
+      "step": 2450
+    },
+    {
+      "epoch": 1.63,
+      "grad_norm": 1.1043754536185766,
+      "learning_rate": 1.3152554879173585e-05,
+      "loss": 0.0284,
+      "step": 2460
+    },
+    {
+      "epoch": 1.64,
+      "grad_norm": 1.203565789266508,
+      "learning_rate": 1.3115661317100167e-05,
+      "loss": 0.0292,
+      "step": 2470
+    },
+    {
+      "epoch": 1.65,
+      "grad_norm": 1.9568074416295906,
+      "learning_rate": 1.3078767755026749e-05,
+      "loss": 0.0262,
+      "step": 2480
+    },
+    {
+      "epoch": 1.65,
+      "grad_norm": 0.2859744392048369,
+      "learning_rate": 1.304187419295333e-05,
+      "loss": 0.0295,
+      "step": 2490
+    },
+    {
+      "epoch": 1.66,
+      "grad_norm": 1.077880610534017,
+      "learning_rate": 1.3004980630879912e-05,
+      "loss": 0.0288,
+      "step": 2500
+    },
+    {
+      "epoch": 1.67,
+      "grad_norm": 0.9333475338900054,
+      "learning_rate": 1.2968087068806494e-05,
+      "loss": 0.0355,
+      "step": 2510
+    },
+    {
+      "epoch": 1.67,
+      "grad_norm": 1.904061045593678,
+      "learning_rate": 1.2931193506733075e-05,
+      "loss": 0.0251,
+      "step": 2520
+    },
+    {
+      "epoch": 1.68,
+      "grad_norm": 0.58158119774247,
+      "learning_rate": 1.2894299944659658e-05,
+      "loss": 0.0233,
+      "step": 2530
+    },
+    {
+      "epoch": 1.69,
+      "grad_norm": 0.7277481905024581,
+      "learning_rate": 1.2857406382586239e-05,
+      "loss": 0.0298,
+      "step": 2540
+    },
+    {
+      "epoch": 1.69,
+      "grad_norm": 1.8262592151991606,
+      "learning_rate": 1.2820512820512823e-05,
+      "loss": 0.0276,
+      "step": 2550
+    },
+    {
+      "epoch": 1.7,
+      "grad_norm": 0.6465268769367779,
+      "learning_rate": 1.2783619258439404e-05,
+      "loss": 0.027,
+      "step": 2560
+    },
+    {
+      "epoch": 1.71,
+      "grad_norm": 0.6081326854882737,
+      "learning_rate": 1.2746725696365985e-05,
+      "loss": 0.0337,
+      "step": 2570
+    },
+    {
+      "epoch": 1.71,
+      "grad_norm": 0.7068771700231296,
+      "learning_rate": 1.2709832134292568e-05,
+      "loss": 0.0236,
+      "step": 2580
+    },
+    {
+      "epoch": 1.72,
+      "grad_norm": 1.8049539476685674,
+      "learning_rate": 1.267293857221915e-05,
+      "loss": 0.0282,
+      "step": 2590
+    },
+    {
+      "epoch": 1.73,
+      "grad_norm": 0.553245442957324,
+      "learning_rate": 1.263604501014573e-05,
+      "loss": 0.0226,
+      "step": 2600
+    },
+    {
+      "epoch": 1.73,
+      "grad_norm": 0.8295872212818136,
+      "learning_rate": 1.2599151448072313e-05,
+      "loss": 0.0232,
+      "step": 2610
+    },
+    {
+      "epoch": 1.74,
+      "grad_norm": 0.9798348798123743,
+      "learning_rate": 1.2562257885998895e-05,
+      "loss": 0.0456,
+      "step": 2620
+    },
+    {
+      "epoch": 1.75,
+      "grad_norm": 1.19245304335107,
+      "learning_rate": 1.2525364323925476e-05,
+      "loss": 0.0284,
+      "step": 2630
+    },
+    {
+      "epoch": 1.75,
+      "grad_norm": 1.5998767447680797,
+      "learning_rate": 1.2488470761852058e-05,
+      "loss": 0.0307,
+      "step": 2640
+    },
+    {
+      "epoch": 1.76,
+      "grad_norm": 1.7431018978542536,
+      "learning_rate": 1.245157719977864e-05,
+      "loss": 0.0216,
+      "step": 2650
+    },
+    {
+      "epoch": 1.77,
+      "grad_norm": 1.0864224894613053,
+      "learning_rate": 1.241468363770522e-05,
+      "loss": 0.0388,
+      "step": 2660
+    },
+    {
+      "epoch": 1.77,
+      "grad_norm": 2.164074288346446,
+      "learning_rate": 1.2377790075631804e-05,
+      "loss": 0.0302,
+      "step": 2670
+    },
+    {
+      "epoch": 1.78,
+      "grad_norm": 0.6890867086297026,
+      "learning_rate": 1.2340896513558385e-05,
+      "loss": 0.0258,
+      "step": 2680
+    },
+    {
+      "epoch": 1.79,
+      "grad_norm": 1.1150440915546074,
+      "learning_rate": 1.2304002951484968e-05,
+      "loss": 0.0279,
+      "step": 2690
+    },
+    {
+      "epoch": 1.79,
+      "grad_norm": 0.7515709240098276,
+      "learning_rate": 1.2267109389411549e-05,
+      "loss": 0.0419,
+      "step": 2700
+    },
+    {
+      "epoch": 1.8,
+      "grad_norm": 0.7795946895425347,
+      "learning_rate": 1.223021582733813e-05,
+      "loss": 0.0152,
+      "step": 2710
+    },
+    {
+      "epoch": 1.81,
+      "grad_norm": 0.7206757082802774,
+      "learning_rate": 1.2193322265264713e-05,
+      "loss": 0.039,
+      "step": 2720
+    },
+    {
+      "epoch": 1.81,
+      "grad_norm": 0.9261417250470075,
+      "learning_rate": 1.2156428703191294e-05,
+      "loss": 0.031,
+      "step": 2730
+    },
+    {
+      "epoch": 1.82,
+      "grad_norm": 0.9738836548814439,
+      "learning_rate": 1.2119535141117875e-05,
+      "loss": 0.0318,
+      "step": 2740
+    },
+    {
+      "epoch": 1.83,
+      "grad_norm": 1.6517158198941853,
+      "learning_rate": 1.2082641579044458e-05,
+      "loss": 0.0346,
+      "step": 2750
+    },
+    {
+      "epoch": 1.83,
+      "grad_norm": 0.8931816948527265,
+      "learning_rate": 1.2045748016971039e-05,
+      "loss": 0.0278,
+      "step": 2760
+    },
+    {
+      "epoch": 1.84,
+      "grad_norm": 1.5388621088460506,
+      "learning_rate": 1.200885445489762e-05,
+      "loss": 0.0296,
+      "step": 2770
+    },
+    {
+      "epoch": 1.85,
+      "grad_norm": 0.6080136375748658,
+      "learning_rate": 1.1971960892824204e-05,
+      "loss": 0.0314,
+      "step": 2780
+    },
+    {
+      "epoch": 1.85,
+      "grad_norm": 0.7855564824393652,
+      "learning_rate": 1.1935067330750784e-05,
+      "loss": 0.0371,
+      "step": 2790
+    },
+    {
+      "epoch": 1.86,
+      "grad_norm": 1.391898825501054,
+      "learning_rate": 1.1898173768677365e-05,
+      "loss": 0.0215,
+      "step": 2800
+    },
+    {
+      "epoch": 1.87,
+      "grad_norm": 0.023622261403316365,
+      "learning_rate": 1.186128020660395e-05,
+      "loss": 0.0289,
+      "step": 2810
+    },
+    {
+      "epoch": 1.87,
+      "grad_norm": 1.2844725300311655,
+      "learning_rate": 1.182438664453053e-05,
+      "loss": 0.0272,
+      "step": 2820
+    },
+    {
+      "epoch": 1.88,
+      "grad_norm": 0.008098058731733223,
+      "learning_rate": 1.1787493082457114e-05,
+      "loss": 0.021,
+      "step": 2830
+    },
+    {
+      "epoch": 1.89,
+      "grad_norm": 2.4455783687272565,
+      "learning_rate": 1.1750599520383695e-05,
+      "loss": 0.0387,
+      "step": 2840
+    },
+    {
+      "epoch": 1.89,
+      "grad_norm": 0.5832035368747368,
+      "learning_rate": 1.1713705958310276e-05,
+      "loss": 0.0392,
+      "step": 2850
+    },
+    {
+      "epoch": 1.9,
+      "grad_norm": 0.4288935348488189,
+      "learning_rate": 1.1676812396236859e-05,
+      "loss": 0.0243,
+      "step": 2860
+    },
+    {
+      "epoch": 1.91,
+      "grad_norm": 0.85140142754813,
+      "learning_rate": 1.163991883416344e-05,
+      "loss": 0.0288,
+      "step": 2870
+    },
+    {
+      "epoch": 1.91,
+      "grad_norm": 0.6208450731238502,
+      "learning_rate": 1.1603025272090021e-05,
+      "loss": 0.0257,
+      "step": 2880
+    },
+    {
+      "epoch": 1.92,
+      "grad_norm": 0.5506545117162834,
+      "learning_rate": 1.1566131710016604e-05,
+      "loss": 0.0257,
+      "step": 2890
+    },
+    {
+      "epoch": 1.93,
+      "grad_norm": 0.6983830876671564,
+      "learning_rate": 1.1529238147943185e-05,
+      "loss": 0.0153,
+      "step": 2900
+    },
+    {
+      "epoch": 1.93,
+      "grad_norm": 0.5365569446878222,
+      "learning_rate": 1.1492344585869766e-05,
+      "loss": 0.0287,
+      "step": 2910
+    },
+    {
+      "epoch": 1.94,
+      "grad_norm": 0.48404097576321364,
+      "learning_rate": 1.1455451023796349e-05,
+      "loss": 0.0189,
+      "step": 2920
+    },
+    {
+      "epoch": 1.95,
+      "grad_norm": 0.2656124897423078,
+      "learning_rate": 1.141855746172293e-05,
+      "loss": 0.0321,
+      "step": 2930
+    },
+    {
+      "epoch": 1.95,
+      "grad_norm": 0.9816015622085352,
+      "learning_rate": 1.1381663899649511e-05,
+      "loss": 0.0392,
+      "step": 2940
+    },
+    {
+      "epoch": 1.96,
+      "grad_norm": 1.6810737744401563,
+      "learning_rate": 1.1344770337576094e-05,
+      "loss": 0.022,
+      "step": 2950
+    },
+    {
+      "epoch": 1.97,
+      "grad_norm": 1.4543583055955254,
+      "learning_rate": 1.1307876775502675e-05,
+      "loss": 0.0331,
+      "step": 2960
+    },
+    {
+      "epoch": 1.97,
+      "grad_norm": 1.4710753680279536,
+      "learning_rate": 1.1270983213429258e-05,
+      "loss": 0.0383,
+      "step": 2970
+    },
+    {
+      "epoch": 1.98,
+      "grad_norm": 1.2478883072438642,
+      "learning_rate": 1.1234089651355839e-05,
+      "loss": 0.0316,
+      "step": 2980
+    },
+    {
+      "epoch": 1.99,
+      "grad_norm": 2.2702061081470153,
+      "learning_rate": 1.119719608928242e-05,
+      "loss": 0.0282,
+      "step": 2990
+    },
+    {
+      "epoch": 1.99,
+      "grad_norm": 0.6940800534656072,
+      "learning_rate": 1.1160302527209005e-05,
+      "loss": 0.0243,
+      "step": 3000
+    },
+    {
+      "epoch": 2.0,
+      "grad_norm": 0.576849535693127,
+      "learning_rate": 1.1123408965135584e-05,
+      "loss": 0.0209,
+      "step": 3010
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 6024,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 4,
+  "save_steps": 500,
+  "total_flos": 0.0,
+  "train_batch_size": 1,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-3012/training_args.bin

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

checkpoint-3012/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)

checkpoint-4518/1_Pooling/config.json

@@ -0,0 +1,10 @@
+{
+  "word_embedding_dimension": 1024,
+  "pooling_mode_cls_token": true,
+  "pooling_mode_mean_tokens": false,
+  "pooling_mode_max_tokens": false,
+  "pooling_mode_mean_sqrt_len_tokens": false,
+  "pooling_mode_weightedmean_tokens": false,
+  "pooling_mode_lasttoken": false,
+  "include_prompt": true
+}

checkpoint-4518/README.md

@@ -0,0 +1,57 @@
+---
+library_name: sentence-transformers
+pipeline_tag: sentence-similarity
+tags:
+- sentence-transformers
+- feature-extraction
+- sentence-similarity
+
+---
+
+# {MODEL_NAME}
+
+This is a [sentence-transformers](https://www.SBERT.net) model: It maps sentences & paragraphs to a 1024 dimensional dense vector space and can be used for tasks like clustering or semantic search.
+
+<!--- Describe your model here -->
+
+## Usage (Sentence-Transformers)
+
+Using this model becomes easy when you have [sentence-transformers](https://www.SBERT.net) installed:
+
+```
+pip install -U sentence-transformers
+```
+
+Then you can use the model like this:
+
+```python
+from sentence_transformers import SentenceTransformer
+sentences = ["This is an example sentence", "Each sentence is converted"]
+
+model = SentenceTransformer('{MODEL_NAME}')
+embeddings = model.encode(sentences)
+print(embeddings)
+```
+
+
+
+## Evaluation Results
+
+<!--- Describe how your model was evaluated -->
+
+For an automated evaluation of this model, see the *Sentence Embeddings Benchmark*: [https://seb.sbert.net](https://seb.sbert.net?model_name={MODEL_NAME})
+
+
+
+## Full Model Architecture
+```
+SentenceTransformer(
+  (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False}) with Transformer model: XLMRobertaModel 
+  (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
+  (2): Normalize()
+)
+```
+
+## Citing & Authors
+
+<!--- Describe where people can find more information -->

checkpoint-4518/colbert_linear.pt

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

checkpoint-4518/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "saved_models/bgem3_unified_finetune_20240330/checkpoint-4518",
+  "architectures": [
+    "XLMRobertaModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 8194,
+  "model_type": "xlm-roberta",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.39.1",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 250002
+}

checkpoint-4518/config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+  "__version__": {
+    "sentence_transformers": "2.6.0",
+    "transformers": "4.39.1",
+    "pytorch": "2.0.1+cu117"
+  },
+  "prompts": {},
+  "default_prompt_name": null
+}