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-127/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+checkpoint-254/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+checkpoint-381/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,165 @@
+---
+library_name: transformers
+license: apache-2.0
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+tags:
+- generated_from_trainer
+datasets:
+- train.jsonl
+model-index:
+- name: outputs/out
+  results: []
+---
+
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+[<img src="https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/axolotl-ai-cloud/axolotl)
+<details><summary>See axolotl config</summary>
+
+axolotl version: `0.6.0`
+```yaml
+base_model: Qwen/Qwen2.5-1.5B-Instruct
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+trust_remote_code: false
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+output_dir: ./outputs/out
+remove_unused_columns: false
+
+chat_template: qwen_25
+# chat_template: qwen_25
+datasets:
+  - path: train.jsonl
+    type: chat_template
+    field_messages: messages
+    message_field_role: role
+    message_field_content: content
+    roles:
+      user:
+        - user
+      assistant:
+        - assistant
+
+dataset_prepared_path: mr1-sft-1
+# dataset_prepared_path: ko_r1
+val_set_size: 0.005
+eval_sample_packing: False
+
+sequence_len: 512
+sample_packing: False
+pad_to_sequence_len: False
+
+wandb_project: mergedbench
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+plugins:
+  - axolotl.integrations.liger.LigerPlugin
+liger_rope: true
+liger_rms_norm: true
+liger_swiglu: true
+liger_fused_linear_cross_entropy: true
+
+gradient_accumulation_steps: 1
+micro_batch_size: 128
+eval_batch_size: 4
+num_epochs: 3
+optimizer: paged_adamw_8bit
+lr_scheduler: cosine
+learning_rate: 2e-5
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16: 
+tf32: false
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+early_stopping_patience:
+resume_from_checkpoint:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+
+warmup_steps: 10
+evals_per_epoch: 3
+eval_max_new_tokens: 128
+eval_table_size:
+saves_per_epoch: 1
+debug:
+deepspeed: deepspeed_configs/zero1.json
+weight_decay: 0.01
+fsdp:
+fsdp_config:
+special_tokens:
+  eos_token: 
+```
+
+</details><br>
+
+# outputs/out
+
+This model is a fine-tuned version of [Qwen/Qwen2.5-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct) on the train.jsonl dataset.
+It achieves the following results on the evaluation set:
+- Loss: 0.3103
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 2e-05
+- train_batch_size: 128
+- eval_batch_size: 4
+- seed: 42
+- distributed_type: multi-GPU
+- num_devices: 2
+- total_train_batch_size: 256
+- total_eval_batch_size: 8
+- optimizer: Use OptimizerNames.PAGED_ADAMW_8BIT with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_steps: 10
+- num_epochs: 3.0
+
+### Training results
+
+| Training Loss | Epoch  | Step | Validation Loss |
+|:-------------:|:------:|:----:|:---------------:|
+| 4.6099        | 0.0079 | 1    | 3.1001          |
+| 0.0071        | 0.3386 | 43   | 0.3896          |
+| 0.0098        | 0.6772 | 86   | 0.3527          |
+| 0.0026        | 1.0157 | 129  | 0.3306          |
+| 0.0128        | 1.3543 | 172  | 0.3166          |
+| 0.0042        | 1.6929 | 215  | 0.3484          |
+| 0.0019        | 2.0315 | 258  | 0.2931          |
+| 0.0039        | 2.3701 | 301  | 0.3032          |
+| 0.0           | 2.7087 | 344  | 0.3103          |
+
+
+### Framework versions
+
+- Transformers 4.48.1
+- Pytorch 2.5.1+cu121
+- Datasets 3.2.0
+- Tokenizers 0.21.0

added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

checkpoint-127/added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

checkpoint-127/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "Qwen/Qwen2.5-1.5B-Instruct",
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.1",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151665
+}

checkpoint-127/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "4.48.1"
+}

checkpoint-127/latest

@@ -0,0 +1 @@
+global_step127

checkpoint-127/model.safetensors

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

checkpoint-127/rng_state_0.pth

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

checkpoint-127/rng_state_1.pth

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

checkpoint-127/scheduler.pt

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

checkpoint-127/special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint-127/tokenizer.json

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

checkpoint-127/tokenizer_config.json

@@ -0,0 +1,208 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

checkpoint-127/trainer_state.json

@@ -0,0 +1,946 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 43,
+  "global_step": 127,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.007874015748031496,
+      "grad_norm": 118.11203002929688,
+      "learning_rate": 2.0000000000000003e-06,
+      "loss": 4.6099,
+      "step": 1
+    },
+    {
+      "epoch": 0.007874015748031496,
+      "eval_loss": 3.1001100540161133,
+      "eval_runtime": 5.3966,
+      "eval_samples_per_second": 30.204,
+      "eval_steps_per_second": 3.891,
+      "step": 1
+    },
+    {
+      "epoch": 0.015748031496062992,
+      "grad_norm": 118.4310302734375,
+      "learning_rate": 4.000000000000001e-06,
+      "loss": 4.5857,
+      "step": 2
+    },
+    {
+      "epoch": 0.023622047244094488,
+      "grad_norm": 103.37439727783203,
+      "learning_rate": 6e-06,
+      "loss": 4.3069,
+      "step": 3
+    },
+    {
+      "epoch": 0.031496062992125984,
+      "grad_norm": 75.05075073242188,
+      "learning_rate": 8.000000000000001e-06,
+      "loss": 3.8754,
+      "step": 4
+    },
+    {
+      "epoch": 0.03937007874015748,
+      "grad_norm": 50.459983825683594,
+      "learning_rate": 1e-05,
+      "loss": 3.2841,
+      "step": 5
+    },
+    {
+      "epoch": 0.047244094488188976,
+      "grad_norm": 47.4603385925293,
+      "learning_rate": 1.2e-05,
+      "loss": 2.4285,
+      "step": 6
+    },
+    {
+      "epoch": 0.05511811023622047,
+      "grad_norm": 32.362667083740234,
+      "learning_rate": 1.4e-05,
+      "loss": 1.8177,
+      "step": 7
+    },
+    {
+      "epoch": 0.06299212598425197,
+      "grad_norm": 22.846933364868164,
+      "learning_rate": 1.6000000000000003e-05,
+      "loss": 1.1567,
+      "step": 8
+    },
+    {
+      "epoch": 0.07086614173228346,
+      "grad_norm": 17.060213088989258,
+      "learning_rate": 1.8e-05,
+      "loss": 0.8257,
+      "step": 9
+    },
+    {
+      "epoch": 0.07874015748031496,
+      "grad_norm": 14.415579795837402,
+      "learning_rate": 2e-05,
+      "loss": 0.4257,
+      "step": 10
+    },
+    {
+      "epoch": 0.08661417322834646,
+      "grad_norm": 7.753712177276611,
+      "learning_rate": 1.999964147509006e-05,
+      "loss": 0.2976,
+      "step": 11
+    },
+    {
+      "epoch": 0.09448818897637795,
+      "grad_norm": 26.883708953857422,
+      "learning_rate": 1.9998565926068253e-05,
+      "loss": 0.3365,
+      "step": 12
+    },
+    {
+      "epoch": 0.10236220472440945,
+      "grad_norm": 10.675631523132324,
+      "learning_rate": 1.9996773430056806e-05,
+      "loss": 0.2161,
+      "step": 13
+    },
+    {
+      "epoch": 0.11023622047244094,
+      "grad_norm": 6.670111179351807,
+      "learning_rate": 1.999426411558661e-05,
+      "loss": 0.1816,
+      "step": 14
+    },
+    {
+      "epoch": 0.11811023622047244,
+      "grad_norm": 8.878239631652832,
+      "learning_rate": 1.9991038162588018e-05,
+      "loss": 0.1567,
+      "step": 15
+    },
+    {
+      "epoch": 0.12598425196850394,
+      "grad_norm": 2.9917383193969727,
+      "learning_rate": 1.9987095802377933e-05,
+      "loss": 0.0813,
+      "step": 16
+    },
+    {
+      "epoch": 0.13385826771653545,
+      "grad_norm": 1.0548763275146484,
+      "learning_rate": 1.9982437317643218e-05,
+      "loss": 0.0217,
+      "step": 17
+    },
+    {
+      "epoch": 0.14173228346456693,
+      "grad_norm": 2.8778488636016846,
+      "learning_rate": 1.9977063042420438e-05,
+      "loss": 0.0618,
+      "step": 18
+    },
+    {
+      "epoch": 0.14960629921259844,
+      "grad_norm": 0.9811734557151794,
+      "learning_rate": 1.99709733620719e-05,
+      "loss": 0.0175,
+      "step": 19
+    },
+    {
+      "epoch": 0.15748031496062992,
+      "grad_norm": 0.7218202948570251,
+      "learning_rate": 1.996416871325803e-05,
+      "loss": 0.0302,
+      "step": 20
+    },
+    {
+      "epoch": 0.16535433070866143,
+      "grad_norm": 1.2746995687484741,
+      "learning_rate": 1.995664958390604e-05,
+      "loss": 0.0453,
+      "step": 21
+    },
+    {
+      "epoch": 0.1732283464566929,
+      "grad_norm": 0.9413469433784485,
+      "learning_rate": 1.9948416513174976e-05,
+      "loss": 0.0175,
+      "step": 22
+    },
+    {
+      "epoch": 0.18110236220472442,
+      "grad_norm": 1.4161137342453003,
+      "learning_rate": 1.9939470091417012e-05,
+      "loss": 0.0277,
+      "step": 23
+    },
+    {
+      "epoch": 0.1889763779527559,
+      "grad_norm": 2.2721235752105713,
+      "learning_rate": 1.992981096013517e-05,
+      "loss": 0.0589,
+      "step": 24
+    },
+    {
+      "epoch": 0.1968503937007874,
+      "grad_norm": 1.143970251083374,
+      "learning_rate": 1.9919439811937283e-05,
+      "loss": 0.0182,
+      "step": 25
+    },
+    {
+      "epoch": 0.2047244094488189,
+      "grad_norm": 0.8054028749465942,
+      "learning_rate": 1.9908357390486342e-05,
+      "loss": 0.0211,
+      "step": 26
+    },
+    {
+      "epoch": 0.2125984251968504,
+      "grad_norm": 1.4449081420898438,
+      "learning_rate": 1.989656449044718e-05,
+      "loss": 0.0244,
+      "step": 27
+    },
+    {
+      "epoch": 0.2204724409448819,
+      "grad_norm": 0.49216631054878235,
+      "learning_rate": 1.988406195742948e-05,
+      "loss": 0.005,
+      "step": 28
+    },
+    {
+      "epoch": 0.2283464566929134,
+      "grad_norm": 0.9945647716522217,
+      "learning_rate": 1.987085068792715e-05,
+      "loss": 0.0373,
+      "step": 29
+    },
+    {
+      "epoch": 0.23622047244094488,
+      "grad_norm": 1.1753748655319214,
+      "learning_rate": 1.9856931629254032e-05,
+      "loss": 0.0217,
+      "step": 30
+    },
+    {
+      "epoch": 0.2440944881889764,
+      "grad_norm": 0.5960403680801392,
+      "learning_rate": 1.984230577947597e-05,
+      "loss": 0.0157,
+      "step": 31
+    },
+    {
+      "epoch": 0.25196850393700787,
+      "grad_norm": 0.3657272160053253,
+      "learning_rate": 1.9826974187339267e-05,
+      "loss": 0.0082,
+      "step": 32
+    },
+    {
+      "epoch": 0.25984251968503935,
+      "grad_norm": 1.1290266513824463,
+      "learning_rate": 1.981093795219546e-05,
+      "loss": 0.0236,
+      "step": 33
+    },
+    {
+      "epoch": 0.2677165354330709,
+      "grad_norm": 1.673962116241455,
+      "learning_rate": 1.9794198223922496e-05,
+      "loss": 0.0182,
+      "step": 34
+    },
+    {
+      "epoch": 0.2755905511811024,
+      "grad_norm": 0.540355384349823,
+      "learning_rate": 1.9776756202842297e-05,
+      "loss": 0.011,
+      "step": 35
+    },
+    {
+      "epoch": 0.28346456692913385,
+      "grad_norm": 0.3380790054798126,
+      "learning_rate": 1.9758613139634662e-05,
+      "loss": 0.0048,
+      "step": 36
+    },
+    {
+      "epoch": 0.29133858267716534,
+      "grad_norm": 1.886232852935791,
+      "learning_rate": 1.9739770335247616e-05,
+      "loss": 0.0157,
+      "step": 37
+    },
+    {
+      "epoch": 0.2992125984251969,
+      "grad_norm": 2.140639305114746,
+      "learning_rate": 1.972022914080411e-05,
+      "loss": 0.0393,
+      "step": 38
+    },
+    {
+      "epoch": 0.30708661417322836,
+      "grad_norm": 0.35308870673179626,
+      "learning_rate": 1.9699990957505136e-05,
+      "loss": 0.0074,
+      "step": 39
+    },
+    {
+      "epoch": 0.31496062992125984,
+      "grad_norm": 0.3918301463127136,
+      "learning_rate": 1.9679057236529266e-05,
+      "loss": 0.0083,
+      "step": 40
+    },
+    {
+      "epoch": 0.3228346456692913,
+      "grad_norm": 0.4406338632106781,
+      "learning_rate": 1.965742947892858e-05,
+      "loss": 0.0152,
+      "step": 41
+    },
+    {
+      "epoch": 0.33070866141732286,
+      "grad_norm": 0.6819682121276855,
+      "learning_rate": 1.9635109235521057e-05,
+      "loss": 0.0091,
+      "step": 42
+    },
+    {
+      "epoch": 0.33858267716535434,
+      "grad_norm": 0.6794927716255188,
+      "learning_rate": 1.961209810677934e-05,
+      "loss": 0.0071,
+      "step": 43
+    },
+    {
+      "epoch": 0.33858267716535434,
+      "eval_loss": 0.3895845115184784,
+      "eval_runtime": 6.5602,
+      "eval_samples_per_second": 24.847,
+      "eval_steps_per_second": 3.201,
+      "step": 43
+    },
+    {
+      "epoch": 0.3464566929133858,
+      "grad_norm": 0.3874967694282532,
+      "learning_rate": 1.9588397742716004e-05,
+      "loss": 0.0089,
+      "step": 44
+    },
+    {
+      "epoch": 0.3543307086614173,
+      "grad_norm": 0.5577577352523804,
+      "learning_rate": 1.9564009842765225e-05,
+      "loss": 0.0098,
+      "step": 45
+    },
+    {
+      "epoch": 0.36220472440944884,
+      "grad_norm": 0.8152347207069397,
+      "learning_rate": 1.9538936155660934e-05,
+      "loss": 0.0118,
+      "step": 46
+    },
+    {
+      "epoch": 0.3700787401574803,
+      "grad_norm": 0.2971118688583374,
+      "learning_rate": 1.951317847931141e-05,
+      "loss": 0.0084,
+      "step": 47
+    },
+    {
+      "epoch": 0.3779527559055118,
+      "grad_norm": 1.0286651849746704,
+      "learning_rate": 1.9486738660670373e-05,
+      "loss": 0.0123,
+      "step": 48
+    },
+    {
+      "epoch": 0.3858267716535433,
+      "grad_norm": 0.5227222442626953,
+      "learning_rate": 1.945961859560454e-05,
+      "loss": 0.0144,
+      "step": 49
+    },
+    {
+      "epoch": 0.3937007874015748,
+      "grad_norm": 0.461935818195343,
+      "learning_rate": 1.943182022875769e-05,
+      "loss": 0.0119,
+      "step": 50
+    },
+    {
+      "epoch": 0.4015748031496063,
+      "grad_norm": 1.2550626993179321,
+      "learning_rate": 1.940334555341122e-05,
+      "loss": 0.013,
+      "step": 51
+    },
+    {
+      "epoch": 0.4094488188976378,
+      "grad_norm": 0.37549659609794617,
+      "learning_rate": 1.9374196611341212e-05,
+      "loss": 0.0181,
+      "step": 52
+    },
+    {
+      "epoch": 0.41732283464566927,
+      "grad_norm": 0.3444191515445709,
+      "learning_rate": 1.9344375492672024e-05,
+      "loss": 0.0111,
+      "step": 53
+    },
+    {
+      "epoch": 0.4251968503937008,
+      "grad_norm": 0.3489387333393097,
+      "learning_rate": 1.9313884335726443e-05,
+      "loss": 0.0111,
+      "step": 54
+    },
+    {
+      "epoch": 0.4330708661417323,
+      "grad_norm": 0.26080814003944397,
+      "learning_rate": 1.9282725326872324e-05,
+      "loss": 0.0091,
+      "step": 55
+    },
+    {
+      "epoch": 0.4409448818897638,
+      "grad_norm": 0.1390451341867447,
+      "learning_rate": 1.9250900700365837e-05,
+      "loss": 0.0033,
+      "step": 56
+    },
+    {
+      "epoch": 0.44881889763779526,
+      "grad_norm": 0.20499111711978912,
+      "learning_rate": 1.921841273819125e-05,
+      "loss": 0.0066,
+      "step": 57
+    },
+    {
+      "epoch": 0.4566929133858268,
+      "grad_norm": 2.185487747192383,
+      "learning_rate": 1.918526376989731e-05,
+      "loss": 0.0095,
+      "step": 58
+    },
+    {
+      "epoch": 0.4645669291338583,
+      "grad_norm": 0.23939816653728485,
+      "learning_rate": 1.9151456172430186e-05,
+      "loss": 0.0048,
+      "step": 59
+    },
+    {
+      "epoch": 0.47244094488188976,
+      "grad_norm": 0.41510018706321716,
+      "learning_rate": 1.911699236996305e-05,
+      "loss": 0.0077,
+      "step": 60
+    },
+    {
+      "epoch": 0.48031496062992124,
+      "grad_norm": 0.264318585395813,
+      "learning_rate": 1.9081874833722234e-05,
+      "loss": 0.0129,
+      "step": 61
+    },
+    {
+      "epoch": 0.4881889763779528,
+      "grad_norm": 1.0443968772888184,
+      "learning_rate": 1.9046106081810047e-05,
+      "loss": 0.0035,
+      "step": 62
+    },
+    {
+      "epoch": 0.49606299212598426,
+      "grad_norm": 0.2800132632255554,
+      "learning_rate": 1.900968867902419e-05,
+      "loss": 0.0057,
+      "step": 63
+    },
+    {
+      "epoch": 0.5039370078740157,
+      "grad_norm": 1.114960789680481,
+      "learning_rate": 1.8972625236673887e-05,
+      "loss": 0.0123,
+      "step": 64
+    },
+    {
+      "epoch": 0.5118110236220472,
+      "grad_norm": 0.5027065873146057,
+      "learning_rate": 1.8934918412392596e-05,
+      "loss": 0.0052,
+      "step": 65
+    },
+    {
+      "epoch": 0.5196850393700787,
+      "grad_norm": 0.5564169883728027,
+      "learning_rate": 1.8896570909947477e-05,
+      "loss": 0.0085,
+      "step": 66
+    },
+    {
+      "epoch": 0.5275590551181102,
+      "grad_norm": 0.7567198872566223,
+      "learning_rate": 1.8857585479045493e-05,
+      "loss": 0.0054,
+      "step": 67
+    },
+    {
+      "epoch": 0.5354330708661418,
+      "grad_norm": 0.13573969900608063,
+      "learning_rate": 1.8817964915136277e-05,
+      "loss": 0.0008,
+      "step": 68
+    },
+    {
+      "epoch": 0.5433070866141733,
+      "grad_norm": 0.2704390287399292,
+      "learning_rate": 1.8777712059211643e-05,
+      "loss": 0.0078,
+      "step": 69
+    },
+    {
+      "epoch": 0.5511811023622047,
+      "grad_norm": 0.6014392971992493,
+      "learning_rate": 1.8736829797601903e-05,
+      "loss": 0.0059,
+      "step": 70
+    },
+    {
+      "epoch": 0.5590551181102362,
+      "grad_norm": 0.5487034916877747,
+      "learning_rate": 1.8695321061768886e-05,
+      "loss": 0.0097,
+      "step": 71
+    },
+    {
+      "epoch": 0.5669291338582677,
+      "grad_norm": 0.6670834422111511,
+      "learning_rate": 1.8653188828095754e-05,
+      "loss": 0.011,
+      "step": 72
+    },
+    {
+      "epoch": 0.5748031496062992,
+      "grad_norm": 0.1795203685760498,
+      "learning_rate": 1.8610436117673557e-05,
+      "loss": 0.0067,
+      "step": 73
+    },
+    {
+      "epoch": 0.5826771653543307,
+      "grad_norm": 1.768436074256897,
+      "learning_rate": 1.8567065996084628e-05,
+      "loss": 0.0096,
+      "step": 74
+    },
+    {
+      "epoch": 0.5905511811023622,
+      "grad_norm": 0.26233312487602234,
+      "learning_rate": 1.8523081573182754e-05,
+      "loss": 0.0124,
+      "step": 75
+    },
+    {
+      "epoch": 0.5984251968503937,
+      "grad_norm": 0.3775719404220581,
+      "learning_rate": 1.847848600287019e-05,
+      "loss": 0.0052,
+      "step": 76
+    },
+    {
+      "epoch": 0.6062992125984252,
+      "grad_norm": 1.0016565322875977,
+      "learning_rate": 1.8433282482871497e-05,
+      "loss": 0.0058,
+      "step": 77
+    },
+    {
+      "epoch": 0.6141732283464567,
+      "grad_norm": 0.20153792202472687,
+      "learning_rate": 1.8387474254504265e-05,
+      "loss": 0.0056,
+      "step": 78
+    },
+    {
+      "epoch": 0.6220472440944882,
+      "grad_norm": 0.5119822025299072,
+      "learning_rate": 1.8341064602446686e-05,
+      "loss": 0.0079,
+      "step": 79
+    },
+    {
+      "epoch": 0.6299212598425197,
+      "grad_norm": 1.5781004428863525,
+      "learning_rate": 1.829405685450202e-05,
+      "loss": 0.008,
+      "step": 80
+    },
+    {
+      "epoch": 0.6377952755905512,
+      "grad_norm": 0.23826757073402405,
+      "learning_rate": 1.824645438135999e-05,
+      "loss": 0.0041,
+      "step": 81
+    },
+    {
+      "epoch": 0.6456692913385826,
+      "grad_norm": 0.6386727690696716,
+      "learning_rate": 1.8198260596355077e-05,
+      "loss": 0.0188,
+      "step": 82
+    },
+    {
+      "epoch": 0.6535433070866141,
+      "grad_norm": 0.9503199458122253,
+      "learning_rate": 1.814947895522176e-05,
+      "loss": 0.008,
+      "step": 83
+    },
+    {
+      "epoch": 0.6614173228346457,
+      "grad_norm": 0.2040701061487198,
+      "learning_rate": 1.8100112955846746e-05,
+      "loss": 0.0038,
+      "step": 84
+    },
+    {
+      "epoch": 0.6692913385826772,
+      "grad_norm": 0.3660199046134949,
+      "learning_rate": 1.805016613801813e-05,
+      "loss": 0.0148,
+      "step": 85
+    },
+    {
+      "epoch": 0.6771653543307087,
+      "grad_norm": 1.0502821207046509,
+      "learning_rate": 1.7999642083171576e-05,
+      "loss": 0.0098,
+      "step": 86
+    },
+    {
+      "epoch": 0.6771653543307087,
+      "eval_loss": 0.3526817262172699,
+      "eval_runtime": 6.6167,
+      "eval_samples_per_second": 24.635,
+      "eval_steps_per_second": 3.174,
+      "step": 86
+    },
+    {
+      "epoch": 0.6850393700787402,
+      "grad_norm": 0.13735969364643097,
+      "learning_rate": 1.7948544414133534e-05,
+      "loss": 0.0022,
+      "step": 87
+    },
+    {
+      "epoch": 0.6929133858267716,
+      "grad_norm": 0.6425012946128845,
+      "learning_rate": 1.7896876794861443e-05,
+      "loss": 0.0086,
+      "step": 88
+    },
+    {
+      "epoch": 0.7007874015748031,
+      "grad_norm": 0.7540380954742432,
+      "learning_rate": 1.7844642930181008e-05,
+      "loss": 0.0062,
+      "step": 89
+    },
+    {
+      "epoch": 0.7086614173228346,
+      "grad_norm": 0.6727365255355835,
+      "learning_rate": 1.779184656552056e-05,
+      "loss": 0.0027,
+      "step": 90
+    },
+    {
+      "epoch": 0.7165354330708661,
+      "grad_norm": 0.14059337973594666,
+      "learning_rate": 1.773849148664247e-05,
+      "loss": 0.0056,
+      "step": 91
+    },
+    {
+      "epoch": 0.7244094488188977,
+      "grad_norm": 0.33292093873023987,
+      "learning_rate": 1.7684581519371714e-05,
+      "loss": 0.0047,
+      "step": 92
+    },
+    {
+      "epoch": 0.7322834645669292,
+      "grad_norm": 0.3809877932071686,
+      "learning_rate": 1.7630120529321518e-05,
+      "loss": 0.0139,
+      "step": 93
+    },
+    {
+      "epoch": 0.7401574803149606,
+      "grad_norm": 1.729589819908142,
+      "learning_rate": 1.7575112421616203e-05,
+      "loss": 0.0128,
+      "step": 94
+    },
+    {
+      "epoch": 0.7480314960629921,
+      "grad_norm": 0.18192608654499054,
+      "learning_rate": 1.751956114061113e-05,
+      "loss": 0.0025,
+      "step": 95
+    },
+    {
+      "epoch": 0.7559055118110236,
+      "grad_norm": 1.0333118438720703,
+      "learning_rate": 1.7463470669609907e-05,
+      "loss": 0.006,
+      "step": 96
+    },
+    {
+      "epoch": 0.7637795275590551,
+      "grad_norm": 0.7247685194015503,
+      "learning_rate": 1.7406845030578747e-05,
+      "loss": 0.0073,
+      "step": 97
+    },
+    {
+      "epoch": 0.7716535433070866,
+      "grad_norm": 0.06979379802942276,
+      "learning_rate": 1.734968828385808e-05,
+      "loss": 0.0005,
+      "step": 98
+    },
+    {
+      "epoch": 0.7795275590551181,
+      "grad_norm": 0.5137119293212891,
+      "learning_rate": 1.729200452787139e-05,
+      "loss": 0.0082,
+      "step": 99
+    },
+    {
+      "epoch": 0.7874015748031497,
+      "grad_norm": 0.4704137146472931,
+      "learning_rate": 1.7233797898831376e-05,
+      "loss": 0.005,
+      "step": 100
+    },
+    {
+      "epoch": 0.7952755905511811,
+      "grad_norm": 0.28564465045928955,
+      "learning_rate": 1.717507257044331e-05,
+      "loss": 0.0052,
+      "step": 101
+    },
+    {
+      "epoch": 0.8031496062992126,
+      "grad_norm": 0.17685537040233612,
+      "learning_rate": 1.711583275360582e-05,
+      "loss": 0.0024,
+      "step": 102
+    },
+    {
+      "epoch": 0.8110236220472441,
+      "grad_norm": 0.45714935660362244,
+      "learning_rate": 1.7056082696108896e-05,
+      "loss": 0.0072,
+      "step": 103
+    },
+    {
+      "epoch": 0.8188976377952756,
+      "grad_norm": 0.4373086988925934,
+      "learning_rate": 1.699582668232934e-05,
+      "loss": 0.0051,
+      "step": 104
+    },
+    {
+      "epoch": 0.8267716535433071,
+      "grad_norm": 0.8478983640670776,
+      "learning_rate": 1.6935069032923525e-05,
+      "loss": 0.022,
+      "step": 105
+    },
+    {
+      "epoch": 0.8346456692913385,
+      "grad_norm": 0.16181086003780365,
+      "learning_rate": 1.6873814104517617e-05,
+      "loss": 0.0058,
+      "step": 106
+    },
+    {
+      "epoch": 0.84251968503937,
+      "grad_norm": 0.09503592550754547,
+      "learning_rate": 1.6812066289395157e-05,
+      "loss": 0.0009,
+      "step": 107
+    },
+    {
+      "epoch": 0.8503937007874016,
+      "grad_norm": 0.7462632060050964,
+      "learning_rate": 1.6749830015182106e-05,
+      "loss": 0.0044,
+      "step": 108
+    },
+    {
+      "epoch": 0.8582677165354331,
+      "grad_norm": 0.07221701741218567,
+      "learning_rate": 1.6687109744529394e-05,
+      "loss": 0.0015,
+      "step": 109
+    },
+    {
+      "epoch": 0.8661417322834646,
+      "grad_norm": 0.08999036252498627,
+      "learning_rate": 1.6623909974792888e-05,
+      "loss": 0.0023,
+      "step": 110
+    },
+    {
+      "epoch": 0.8740157480314961,
+      "grad_norm": 0.42536938190460205,
+      "learning_rate": 1.656023523771095e-05,
+      "loss": 0.005,
+      "step": 111
+    },
+    {
+      "epoch": 0.8818897637795275,
+      "grad_norm": 0.7885191440582275,
+      "learning_rate": 1.6496090099079452e-05,
+      "loss": 0.0103,
+      "step": 112
+    },
+    {
+      "epoch": 0.889763779527559,
+      "grad_norm": 0.16610018908977509,
+      "learning_rate": 1.64314791584244e-05,
+      "loss": 0.006,
+      "step": 113
+    },
+    {
+      "epoch": 0.8976377952755905,
+      "grad_norm": 0.32151034474372864,
+      "learning_rate": 1.6366407048672135e-05,
+      "loss": 0.0086,
+      "step": 114
+    },
+    {
+      "epoch": 0.905511811023622,
+      "grad_norm": 0.557732343673706,
+      "learning_rate": 1.6300878435817115e-05,
+      "loss": 0.0064,
+      "step": 115
+    },
+    {
+      "epoch": 0.9133858267716536,
+      "grad_norm": 0.2238176167011261,
+      "learning_rate": 1.6234898018587336e-05,
+      "loss": 0.0065,
+      "step": 116
+    },
+    {
+      "epoch": 0.9212598425196851,
+      "grad_norm": 0.2980042099952698,
+      "learning_rate": 1.616847052810744e-05,
+      "loss": 0.0095,
+      "step": 117
+    },
+    {
+      "epoch": 0.9291338582677166,
+      "grad_norm": 0.1529705822467804,
+      "learning_rate": 1.6101600727559423e-05,
+      "loss": 0.0062,
+      "step": 118
+    },
+    {
+      "epoch": 0.937007874015748,
+      "grad_norm": 0.017149658873677254,
+      "learning_rate": 1.603429341184114e-05,
+      "loss": 0.0002,
+      "step": 119
+    },
+    {
+      "epoch": 0.9448818897637795,
+      "grad_norm": 0.4514746367931366,
+      "learning_rate": 1.596655340722244e-05,
+      "loss": 0.0067,
+      "step": 120
+    },
+    {
+      "epoch": 0.952755905511811,
+      "grad_norm": 0.11766134947538376,
+      "learning_rate": 1.5898385570999146e-05,
+      "loss": 0.0053,
+      "step": 121
+    },
+    {
+      "epoch": 0.9606299212598425,
+      "grad_norm": 0.4089784026145935,
+      "learning_rate": 1.5829794791144723e-05,
+      "loss": 0.0085,
+      "step": 122
+    },
+    {
+      "epoch": 0.968503937007874,
+      "grad_norm": 0.1353057473897934,
+      "learning_rate": 1.57607859859598e-05,
+      "loss": 0.0013,
+      "step": 123
+    },
+    {
+      "epoch": 0.9763779527559056,
+      "grad_norm": 0.6548481583595276,
+      "learning_rate": 1.5691364103719515e-05,
+      "loss": 0.0117,
+      "step": 124
+    },
+    {
+      "epoch": 0.984251968503937,
+      "grad_norm": 0.1571267992258072,
+      "learning_rate": 1.5621534122318682e-05,
+      "loss": 0.0049,
+      "step": 125
+    },
+    {
+      "epoch": 0.9921259842519685,
+      "grad_norm": 1.2177189588546753,
+      "learning_rate": 1.5551301048914863e-05,
+      "loss": 0.0161,
+      "step": 126
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.414489209651947,
+      "learning_rate": 1.5480669919569313e-05,
+      "loss": 0.0181,
+      "step": 127
+    }
+  ],
+  "logging_steps": 1,
+  "max_steps": 381,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 3,
+  "save_steps": 127,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 1.3087271069889331e+17,
+  "train_batch_size": 128,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-127/training_args.bin

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

checkpoint-127/zero_to_fp32.py

@@ -0,0 +1,760 @@
+#!/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 . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+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, weights_only=False)
+
+        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 tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # 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}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] 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, exclude_frozen_parameters):
+    """
+    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,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+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 _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+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() if _has_callable(shape, 'numel') else math.prod(shape)
+            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,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _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")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+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 = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * 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
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        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}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        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,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _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 to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    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``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    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.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    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")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    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_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``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``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        # an memory-efficient approach for sharding
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+        shard_state_dict = to_torch_tensor(shard_state_dict)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+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_dir",
+                        type=str,
+                        help="directory to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    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("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    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_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)

checkpoint-254/added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

checkpoint-254/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "Qwen/Qwen2.5-1.5B-Instruct",
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.1",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151665
+}

checkpoint-254/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "4.48.1"
+}

checkpoint-254/latest

@@ -0,0 +1 @@
+global_step254

checkpoint-254/model.safetensors

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

checkpoint-254/rng_state_0.pth

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

checkpoint-254/rng_state_1.pth

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

checkpoint-254/scheduler.pt

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

checkpoint-254/special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint-254/tokenizer.json

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

checkpoint-254/tokenizer_config.json

@@ -0,0 +1,208 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

checkpoint-254/trainer_state.json

@@ -0,0 +1,1859 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.0,
+  "eval_steps": 43,
+  "global_step": 254,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.007874015748031496,
+      "grad_norm": 118.11203002929688,
+      "learning_rate": 2.0000000000000003e-06,
+      "loss": 4.6099,
+      "step": 1
+    },
+    {
+      "epoch": 0.007874015748031496,
+      "eval_loss": 3.1001100540161133,
+      "eval_runtime": 5.3966,
+      "eval_samples_per_second": 30.204,
+      "eval_steps_per_second": 3.891,
+      "step": 1
+    },
+    {
+      "epoch": 0.015748031496062992,
+      "grad_norm": 118.4310302734375,
+      "learning_rate": 4.000000000000001e-06,
+      "loss": 4.5857,
+      "step": 2
+    },
+    {
+      "epoch": 0.023622047244094488,
+      "grad_norm": 103.37439727783203,
+      "learning_rate": 6e-06,
+      "loss": 4.3069,
+      "step": 3
+    },
+    {
+      "epoch": 0.031496062992125984,
+      "grad_norm": 75.05075073242188,
+      "learning_rate": 8.000000000000001e-06,
+      "loss": 3.8754,
+      "step": 4
+    },
+    {
+      "epoch": 0.03937007874015748,
+      "grad_norm": 50.459983825683594,
+      "learning_rate": 1e-05,
+      "loss": 3.2841,
+      "step": 5
+    },
+    {
+      "epoch": 0.047244094488188976,
+      "grad_norm": 47.4603385925293,
+      "learning_rate": 1.2e-05,
+      "loss": 2.4285,
+      "step": 6
+    },
+    {
+      "epoch": 0.05511811023622047,
+      "grad_norm": 32.362667083740234,
+      "learning_rate": 1.4e-05,
+      "loss": 1.8177,
+      "step": 7
+    },
+    {
+      "epoch": 0.06299212598425197,
+      "grad_norm": 22.846933364868164,
+      "learning_rate": 1.6000000000000003e-05,
+      "loss": 1.1567,
+      "step": 8
+    },
+    {
+      "epoch": 0.07086614173228346,
+      "grad_norm": 17.060213088989258,
+      "learning_rate": 1.8e-05,
+      "loss": 0.8257,
+      "step": 9
+    },
+    {
+      "epoch": 0.07874015748031496,
+      "grad_norm": 14.415579795837402,
+      "learning_rate": 2e-05,
+      "loss": 0.4257,
+      "step": 10
+    },
+    {
+      "epoch": 0.08661417322834646,
+      "grad_norm": 7.753712177276611,
+      "learning_rate": 1.999964147509006e-05,
+      "loss": 0.2976,
+      "step": 11
+    },
+    {
+      "epoch": 0.09448818897637795,
+      "grad_norm": 26.883708953857422,
+      "learning_rate": 1.9998565926068253e-05,
+      "loss": 0.3365,
+      "step": 12
+    },
+    {
+      "epoch": 0.10236220472440945,
+      "grad_norm": 10.675631523132324,
+      "learning_rate": 1.9996773430056806e-05,
+      "loss": 0.2161,
+      "step": 13
+    },
+    {
+      "epoch": 0.11023622047244094,
+      "grad_norm": 6.670111179351807,
+      "learning_rate": 1.999426411558661e-05,
+      "loss": 0.1816,
+      "step": 14
+    },
+    {
+      "epoch": 0.11811023622047244,
+      "grad_norm": 8.878239631652832,
+      "learning_rate": 1.9991038162588018e-05,
+      "loss": 0.1567,
+      "step": 15
+    },
+    {
+      "epoch": 0.12598425196850394,
+      "grad_norm": 2.9917383193969727,
+      "learning_rate": 1.9987095802377933e-05,
+      "loss": 0.0813,
+      "step": 16
+    },
+    {
+      "epoch": 0.13385826771653545,
+      "grad_norm": 1.0548763275146484,
+      "learning_rate": 1.9982437317643218e-05,
+      "loss": 0.0217,
+      "step": 17
+    },
+    {
+      "epoch": 0.14173228346456693,
+      "grad_norm": 2.8778488636016846,
+      "learning_rate": 1.9977063042420438e-05,
+      "loss": 0.0618,
+      "step": 18
+    },
+    {
+      "epoch": 0.14960629921259844,
+      "grad_norm": 0.9811734557151794,
+      "learning_rate": 1.99709733620719e-05,
+      "loss": 0.0175,
+      "step": 19
+    },
+    {
+      "epoch": 0.15748031496062992,
+      "grad_norm": 0.7218202948570251,
+      "learning_rate": 1.996416871325803e-05,
+      "loss": 0.0302,
+      "step": 20
+    },
+    {
+      "epoch": 0.16535433070866143,
+      "grad_norm": 1.2746995687484741,
+      "learning_rate": 1.995664958390604e-05,
+      "loss": 0.0453,
+      "step": 21
+    },
+    {
+      "epoch": 0.1732283464566929,
+      "grad_norm": 0.9413469433784485,
+      "learning_rate": 1.9948416513174976e-05,
+      "loss": 0.0175,
+      "step": 22
+    },
+    {
+      "epoch": 0.18110236220472442,
+      "grad_norm": 1.4161137342453003,
+      "learning_rate": 1.9939470091417012e-05,
+      "loss": 0.0277,
+      "step": 23
+    },
+    {
+      "epoch": 0.1889763779527559,
+      "grad_norm": 2.2721235752105713,
+      "learning_rate": 1.992981096013517e-05,
+      "loss": 0.0589,
+      "step": 24
+    },
+    {
+      "epoch": 0.1968503937007874,
+      "grad_norm": 1.143970251083374,
+      "learning_rate": 1.9919439811937283e-05,
+      "loss": 0.0182,
+      "step": 25
+    },
+    {
+      "epoch": 0.2047244094488189,
+      "grad_norm": 0.8054028749465942,
+      "learning_rate": 1.9908357390486342e-05,
+      "loss": 0.0211,
+      "step": 26
+    },
+    {
+      "epoch": 0.2125984251968504,
+      "grad_norm": 1.4449081420898438,
+      "learning_rate": 1.989656449044718e-05,
+      "loss": 0.0244,
+      "step": 27
+    },
+    {
+      "epoch": 0.2204724409448819,
+      "grad_norm": 0.49216631054878235,
+      "learning_rate": 1.988406195742948e-05,
+      "loss": 0.005,
+      "step": 28
+    },
+    {
+      "epoch": 0.2283464566929134,
+      "grad_norm": 0.9945647716522217,
+      "learning_rate": 1.987085068792715e-05,
+      "loss": 0.0373,
+      "step": 29
+    },
+    {
+      "epoch": 0.23622047244094488,
+      "grad_norm": 1.1753748655319214,
+      "learning_rate": 1.9856931629254032e-05,
+      "loss": 0.0217,
+      "step": 30
+    },
+    {
+      "epoch": 0.2440944881889764,
+      "grad_norm": 0.5960403680801392,
+      "learning_rate": 1.984230577947597e-05,
+      "loss": 0.0157,
+      "step": 31
+    },
+    {
+      "epoch": 0.25196850393700787,
+      "grad_norm": 0.3657272160053253,
+      "learning_rate": 1.9826974187339267e-05,
+      "loss": 0.0082,
+      "step": 32
+    },
+    {
+      "epoch": 0.25984251968503935,
+      "grad_norm": 1.1290266513824463,
+      "learning_rate": 1.981093795219546e-05,
+      "loss": 0.0236,
+      "step": 33
+    },
+    {
+      "epoch": 0.2677165354330709,
+      "grad_norm": 1.673962116241455,
+      "learning_rate": 1.9794198223922496e-05,
+      "loss": 0.0182,
+      "step": 34
+    },
+    {
+      "epoch": 0.2755905511811024,
+      "grad_norm": 0.540355384349823,
+      "learning_rate": 1.9776756202842297e-05,
+      "loss": 0.011,
+      "step": 35
+    },
+    {
+      "epoch": 0.28346456692913385,
+      "grad_norm": 0.3380790054798126,
+      "learning_rate": 1.9758613139634662e-05,
+      "loss": 0.0048,
+      "step": 36
+    },
+    {
+      "epoch": 0.29133858267716534,
+      "grad_norm": 1.886232852935791,
+      "learning_rate": 1.9739770335247616e-05,
+      "loss": 0.0157,
+      "step": 37
+    },
+    {
+      "epoch": 0.2992125984251969,
+      "grad_norm": 2.140639305114746,
+      "learning_rate": 1.972022914080411e-05,
+      "loss": 0.0393,
+      "step": 38
+    },
+    {
+      "epoch": 0.30708661417322836,
+      "grad_norm": 0.35308870673179626,
+      "learning_rate": 1.9699990957505136e-05,
+      "loss": 0.0074,
+      "step": 39
+    },
+    {
+      "epoch": 0.31496062992125984,
+      "grad_norm": 0.3918301463127136,
+      "learning_rate": 1.9679057236529266e-05,
+      "loss": 0.0083,
+      "step": 40
+    },
+    {
+      "epoch": 0.3228346456692913,
+      "grad_norm": 0.4406338632106781,
+      "learning_rate": 1.965742947892858e-05,
+      "loss": 0.0152,
+      "step": 41
+    },
+    {
+      "epoch": 0.33070866141732286,
+      "grad_norm": 0.6819682121276855,
+      "learning_rate": 1.9635109235521057e-05,
+      "loss": 0.0091,
+      "step": 42
+    },
+    {
+      "epoch": 0.33858267716535434,
+      "grad_norm": 0.6794927716255188,
+      "learning_rate": 1.961209810677934e-05,
+      "loss": 0.0071,
+      "step": 43
+    },
+    {
+      "epoch": 0.33858267716535434,
+      "eval_loss": 0.3895845115184784,
+      "eval_runtime": 6.5602,
+      "eval_samples_per_second": 24.847,
+      "eval_steps_per_second": 3.201,
+      "step": 43
+    },
+    {
+      "epoch": 0.3464566929133858,
+      "grad_norm": 0.3874967694282532,
+      "learning_rate": 1.9588397742716004e-05,
+      "loss": 0.0089,
+      "step": 44
+    },
+    {
+      "epoch": 0.3543307086614173,
+      "grad_norm": 0.5577577352523804,
+      "learning_rate": 1.9564009842765225e-05,
+      "loss": 0.0098,
+      "step": 45
+    },
+    {
+      "epoch": 0.36220472440944884,
+      "grad_norm": 0.8152347207069397,
+      "learning_rate": 1.9538936155660934e-05,
+      "loss": 0.0118,
+      "step": 46
+    },
+    {
+      "epoch": 0.3700787401574803,
+      "grad_norm": 0.2971118688583374,
+      "learning_rate": 1.951317847931141e-05,
+      "loss": 0.0084,
+      "step": 47
+    },
+    {
+      "epoch": 0.3779527559055118,
+      "grad_norm": 1.0286651849746704,
+      "learning_rate": 1.9486738660670373e-05,
+      "loss": 0.0123,
+      "step": 48
+    },
+    {
+      "epoch": 0.3858267716535433,
+      "grad_norm": 0.5227222442626953,
+      "learning_rate": 1.945961859560454e-05,
+      "loss": 0.0144,
+      "step": 49
+    },
+    {
+      "epoch": 0.3937007874015748,
+      "grad_norm": 0.461935818195343,
+      "learning_rate": 1.943182022875769e-05,
+      "loss": 0.0119,
+      "step": 50
+    },
+    {
+      "epoch": 0.4015748031496063,
+      "grad_norm": 1.2550626993179321,
+      "learning_rate": 1.940334555341122e-05,
+      "loss": 0.013,
+      "step": 51
+    },
+    {
+      "epoch": 0.4094488188976378,
+      "grad_norm": 0.37549659609794617,
+      "learning_rate": 1.9374196611341212e-05,
+      "loss": 0.0181,
+      "step": 52
+    },
+    {
+      "epoch": 0.41732283464566927,
+      "grad_norm": 0.3444191515445709,
+      "learning_rate": 1.9344375492672024e-05,
+      "loss": 0.0111,
+      "step": 53
+    },
+    {
+      "epoch": 0.4251968503937008,
+      "grad_norm": 0.3489387333393097,
+      "learning_rate": 1.9313884335726443e-05,
+      "loss": 0.0111,
+      "step": 54
+    },
+    {
+      "epoch": 0.4330708661417323,
+      "grad_norm": 0.26080814003944397,
+      "learning_rate": 1.9282725326872324e-05,
+      "loss": 0.0091,
+      "step": 55
+    },
+    {
+      "epoch": 0.4409448818897638,
+      "grad_norm": 0.1390451341867447,
+      "learning_rate": 1.9250900700365837e-05,
+      "loss": 0.0033,
+      "step": 56
+    },
+    {
+      "epoch": 0.44881889763779526,
+      "grad_norm": 0.20499111711978912,
+      "learning_rate": 1.921841273819125e-05,
+      "loss": 0.0066,
+      "step": 57
+    },
+    {
+      "epoch": 0.4566929133858268,
+      "grad_norm": 2.185487747192383,
+      "learning_rate": 1.918526376989731e-05,
+      "loss": 0.0095,
+      "step": 58
+    },
+    {
+      "epoch": 0.4645669291338583,
+      "grad_norm": 0.23939816653728485,
+      "learning_rate": 1.9151456172430186e-05,
+      "loss": 0.0048,
+      "step": 59
+    },
+    {
+      "epoch": 0.47244094488188976,
+      "grad_norm": 0.41510018706321716,
+      "learning_rate": 1.911699236996305e-05,
+      "loss": 0.0077,
+      "step": 60
+    },
+    {
+      "epoch": 0.48031496062992124,
+      "grad_norm": 0.264318585395813,
+      "learning_rate": 1.9081874833722234e-05,
+      "loss": 0.0129,
+      "step": 61
+    },
+    {
+      "epoch": 0.4881889763779528,
+      "grad_norm": 1.0443968772888184,
+      "learning_rate": 1.9046106081810047e-05,
+      "loss": 0.0035,
+      "step": 62
+    },
+    {
+      "epoch": 0.49606299212598426,
+      "grad_norm": 0.2800132632255554,
+      "learning_rate": 1.900968867902419e-05,
+      "loss": 0.0057,
+      "step": 63
+    },
+    {
+      "epoch": 0.5039370078740157,
+      "grad_norm": 1.114960789680481,
+      "learning_rate": 1.8972625236673887e-05,
+      "loss": 0.0123,
+      "step": 64
+    },
+    {
+      "epoch": 0.5118110236220472,
+      "grad_norm": 0.5027065873146057,
+      "learning_rate": 1.8934918412392596e-05,
+      "loss": 0.0052,
+      "step": 65
+    },
+    {
+      "epoch": 0.5196850393700787,
+      "grad_norm": 0.5564169883728027,
+      "learning_rate": 1.8896570909947477e-05,
+      "loss": 0.0085,
+      "step": 66
+    },
+    {
+      "epoch": 0.5275590551181102,
+      "grad_norm": 0.7567198872566223,
+      "learning_rate": 1.8857585479045493e-05,
+      "loss": 0.0054,
+      "step": 67
+    },
+    {
+      "epoch": 0.5354330708661418,
+      "grad_norm": 0.13573969900608063,
+      "learning_rate": 1.8817964915136277e-05,
+      "loss": 0.0008,
+      "step": 68
+    },
+    {
+      "epoch": 0.5433070866141733,
+      "grad_norm": 0.2704390287399292,
+      "learning_rate": 1.8777712059211643e-05,
+      "loss": 0.0078,
+      "step": 69
+    },
+    {
+      "epoch": 0.5511811023622047,
+      "grad_norm": 0.6014392971992493,
+      "learning_rate": 1.8736829797601903e-05,
+      "loss": 0.0059,
+      "step": 70
+    },
+    {
+      "epoch": 0.5590551181102362,
+      "grad_norm": 0.5487034916877747,
+      "learning_rate": 1.8695321061768886e-05,
+      "loss": 0.0097,
+      "step": 71
+    },
+    {
+      "epoch": 0.5669291338582677,
+      "grad_norm": 0.6670834422111511,
+      "learning_rate": 1.8653188828095754e-05,
+      "loss": 0.011,
+      "step": 72
+    },
+    {
+      "epoch": 0.5748031496062992,
+      "grad_norm": 0.1795203685760498,
+      "learning_rate": 1.8610436117673557e-05,
+      "loss": 0.0067,
+      "step": 73
+    },
+    {
+      "epoch": 0.5826771653543307,
+      "grad_norm": 1.768436074256897,
+      "learning_rate": 1.8567065996084628e-05,
+      "loss": 0.0096,
+      "step": 74
+    },
+    {
+      "epoch": 0.5905511811023622,
+      "grad_norm": 0.26233312487602234,
+      "learning_rate": 1.8523081573182754e-05,
+      "loss": 0.0124,
+      "step": 75
+    },
+    {
+      "epoch": 0.5984251968503937,
+      "grad_norm": 0.3775719404220581,
+      "learning_rate": 1.847848600287019e-05,
+      "loss": 0.0052,
+      "step": 76
+    },
+    {
+      "epoch": 0.6062992125984252,
+      "grad_norm": 1.0016565322875977,
+      "learning_rate": 1.8433282482871497e-05,
+      "loss": 0.0058,
+      "step": 77
+    },
+    {
+      "epoch": 0.6141732283464567,
+      "grad_norm": 0.20153792202472687,
+      "learning_rate": 1.8387474254504265e-05,
+      "loss": 0.0056,
+      "step": 78
+    },
+    {
+      "epoch": 0.6220472440944882,
+      "grad_norm": 0.5119822025299072,
+      "learning_rate": 1.8341064602446686e-05,
+      "loss": 0.0079,
+      "step": 79
+    },
+    {
+      "epoch": 0.6299212598425197,
+      "grad_norm": 1.5781004428863525,
+      "learning_rate": 1.829405685450202e-05,
+      "loss": 0.008,
+      "step": 80
+    },
+    {
+      "epoch": 0.6377952755905512,
+      "grad_norm": 0.23826757073402405,
+      "learning_rate": 1.824645438135999e-05,
+      "loss": 0.0041,
+      "step": 81
+    },
+    {
+      "epoch": 0.6456692913385826,
+      "grad_norm": 0.6386727690696716,
+      "learning_rate": 1.8198260596355077e-05,
+      "loss": 0.0188,
+      "step": 82
+    },
+    {
+      "epoch": 0.6535433070866141,
+      "grad_norm": 0.9503199458122253,
+      "learning_rate": 1.814947895522176e-05,
+      "loss": 0.008,
+      "step": 83
+    },
+    {
+      "epoch": 0.6614173228346457,
+      "grad_norm": 0.2040701061487198,
+      "learning_rate": 1.8100112955846746e-05,
+      "loss": 0.0038,
+      "step": 84
+    },
+    {
+      "epoch": 0.6692913385826772,
+      "grad_norm": 0.3660199046134949,
+      "learning_rate": 1.805016613801813e-05,
+      "loss": 0.0148,
+      "step": 85
+    },
+    {
+      "epoch": 0.6771653543307087,
+      "grad_norm": 1.0502821207046509,
+      "learning_rate": 1.7999642083171576e-05,
+      "loss": 0.0098,
+      "step": 86
+    },
+    {
+      "epoch": 0.6771653543307087,
+      "eval_loss": 0.3526817262172699,
+      "eval_runtime": 6.6167,
+      "eval_samples_per_second": 24.635,
+      "eval_steps_per_second": 3.174,
+      "step": 86
+    },
+    {
+      "epoch": 0.6850393700787402,
+      "grad_norm": 0.13735969364643097,
+      "learning_rate": 1.7948544414133534e-05,
+      "loss": 0.0022,
+      "step": 87
+    },
+    {
+      "epoch": 0.6929133858267716,
+      "grad_norm": 0.6425012946128845,
+      "learning_rate": 1.7896876794861443e-05,
+      "loss": 0.0086,
+      "step": 88
+    },
+    {
+      "epoch": 0.7007874015748031,
+      "grad_norm": 0.7540380954742432,
+      "learning_rate": 1.7844642930181008e-05,
+      "loss": 0.0062,
+      "step": 89
+    },
+    {
+      "epoch": 0.7086614173228346,
+      "grad_norm": 0.6727365255355835,
+      "learning_rate": 1.779184656552056e-05,
+      "loss": 0.0027,
+      "step": 90
+    },
+    {
+      "epoch": 0.7165354330708661,
+      "grad_norm": 0.14059337973594666,
+      "learning_rate": 1.773849148664247e-05,
+      "loss": 0.0056,
+      "step": 91
+    },
+    {
+      "epoch": 0.7244094488188977,
+      "grad_norm": 0.33292093873023987,
+      "learning_rate": 1.7684581519371714e-05,
+      "loss": 0.0047,
+      "step": 92
+    },
+    {
+      "epoch": 0.7322834645669292,
+      "grad_norm": 0.3809877932071686,
+      "learning_rate": 1.7630120529321518e-05,
+      "loss": 0.0139,
+      "step": 93
+    },
+    {
+      "epoch": 0.7401574803149606,
+      "grad_norm": 1.729589819908142,
+      "learning_rate": 1.7575112421616203e-05,
+      "loss": 0.0128,
+      "step": 94
+    },
+    {
+      "epoch": 0.7480314960629921,
+      "grad_norm": 0.18192608654499054,
+      "learning_rate": 1.751956114061113e-05,
+      "loss": 0.0025,
+      "step": 95
+    },
+    {
+      "epoch": 0.7559055118110236,
+      "grad_norm": 1.0333118438720703,
+      "learning_rate": 1.7463470669609907e-05,
+      "loss": 0.006,
+      "step": 96
+    },
+    {
+      "epoch": 0.7637795275590551,
+      "grad_norm": 0.7247685194015503,
+      "learning_rate": 1.7406845030578747e-05,
+      "loss": 0.0073,
+      "step": 97
+    },
+    {
+      "epoch": 0.7716535433070866,
+      "grad_norm": 0.06979379802942276,
+      "learning_rate": 1.734968828385808e-05,
+      "loss": 0.0005,
+      "step": 98
+    },
+    {
+      "epoch": 0.7795275590551181,
+      "grad_norm": 0.5137119293212891,
+      "learning_rate": 1.729200452787139e-05,
+      "loss": 0.0082,
+      "step": 99
+    },
+    {
+      "epoch": 0.7874015748031497,
+      "grad_norm": 0.4704137146472931,
+      "learning_rate": 1.7233797898831376e-05,
+      "loss": 0.005,
+      "step": 100
+    },
+    {
+      "epoch": 0.7952755905511811,
+      "grad_norm": 0.28564465045928955,
+      "learning_rate": 1.717507257044331e-05,
+      "loss": 0.0052,
+      "step": 101
+    },
+    {
+      "epoch": 0.8031496062992126,
+      "grad_norm": 0.17685537040233612,
+      "learning_rate": 1.711583275360582e-05,
+      "loss": 0.0024,
+      "step": 102
+    },
+    {
+      "epoch": 0.8110236220472441,
+      "grad_norm": 0.45714935660362244,
+      "learning_rate": 1.7056082696108896e-05,
+      "loss": 0.0072,
+      "step": 103
+    },
+    {
+      "epoch": 0.8188976377952756,
+      "grad_norm": 0.4373086988925934,
+      "learning_rate": 1.699582668232934e-05,
+      "loss": 0.0051,
+      "step": 104
+    },
+    {
+      "epoch": 0.8267716535433071,
+      "grad_norm": 0.8478983640670776,
+      "learning_rate": 1.6935069032923525e-05,
+      "loss": 0.022,
+      "step": 105
+    },
+    {
+      "epoch": 0.8346456692913385,
+      "grad_norm": 0.16181086003780365,
+      "learning_rate": 1.6873814104517617e-05,
+      "loss": 0.0058,
+      "step": 106
+    },
+    {
+      "epoch": 0.84251968503937,
+      "grad_norm": 0.09503592550754547,
+      "learning_rate": 1.6812066289395157e-05,
+      "loss": 0.0009,
+      "step": 107
+    },
+    {
+      "epoch": 0.8503937007874016,
+      "grad_norm": 0.7462632060050964,
+      "learning_rate": 1.6749830015182106e-05,
+      "loss": 0.0044,
+      "step": 108
+    },
+    {
+      "epoch": 0.8582677165354331,
+      "grad_norm": 0.07221701741218567,
+      "learning_rate": 1.6687109744529394e-05,
+      "loss": 0.0015,
+      "step": 109
+    },
+    {
+      "epoch": 0.8661417322834646,
+      "grad_norm": 0.08999036252498627,
+      "learning_rate": 1.6623909974792888e-05,
+      "loss": 0.0023,
+      "step": 110
+    },
+    {
+      "epoch": 0.8740157480314961,
+      "grad_norm": 0.42536938190460205,
+      "learning_rate": 1.656023523771095e-05,
+      "loss": 0.005,
+      "step": 111
+    },
+    {
+      "epoch": 0.8818897637795275,
+      "grad_norm": 0.7885191440582275,
+      "learning_rate": 1.6496090099079452e-05,
+      "loss": 0.0103,
+      "step": 112
+    },
+    {
+      "epoch": 0.889763779527559,
+      "grad_norm": 0.16610018908977509,
+      "learning_rate": 1.64314791584244e-05,
+      "loss": 0.006,
+      "step": 113
+    },
+    {
+      "epoch": 0.8976377952755905,
+      "grad_norm": 0.32151034474372864,
+      "learning_rate": 1.6366407048672135e-05,
+      "loss": 0.0086,
+      "step": 114
+    },
+    {
+      "epoch": 0.905511811023622,
+      "grad_norm": 0.557732343673706,
+      "learning_rate": 1.6300878435817115e-05,
+      "loss": 0.0064,
+      "step": 115
+    },
+    {
+      "epoch": 0.9133858267716536,
+      "grad_norm": 0.2238176167011261,
+      "learning_rate": 1.6234898018587336e-05,
+      "loss": 0.0065,
+      "step": 116
+    },
+    {
+      "epoch": 0.9212598425196851,
+      "grad_norm": 0.2980042099952698,
+      "learning_rate": 1.616847052810744e-05,
+      "loss": 0.0095,
+      "step": 117
+    },
+    {
+      "epoch": 0.9291338582677166,
+      "grad_norm": 0.1529705822467804,
+      "learning_rate": 1.6101600727559423e-05,
+      "loss": 0.0062,
+      "step": 118
+    },
+    {
+      "epoch": 0.937007874015748,
+      "grad_norm": 0.017149658873677254,
+      "learning_rate": 1.603429341184114e-05,
+      "loss": 0.0002,
+      "step": 119
+    },
+    {
+      "epoch": 0.9448818897637795,
+      "grad_norm": 0.4514746367931366,
+      "learning_rate": 1.596655340722244e-05,
+      "loss": 0.0067,
+      "step": 120
+    },
+    {
+      "epoch": 0.952755905511811,
+      "grad_norm": 0.11766134947538376,
+      "learning_rate": 1.5898385570999146e-05,
+      "loss": 0.0053,
+      "step": 121
+    },
+    {
+      "epoch": 0.9606299212598425,
+      "grad_norm": 0.4089784026145935,
+      "learning_rate": 1.5829794791144723e-05,
+      "loss": 0.0085,
+      "step": 122
+    },
+    {
+      "epoch": 0.968503937007874,
+      "grad_norm": 0.1353057473897934,
+      "learning_rate": 1.57607859859598e-05,
+      "loss": 0.0013,
+      "step": 123
+    },
+    {
+      "epoch": 0.9763779527559056,
+      "grad_norm": 0.6548481583595276,
+      "learning_rate": 1.5691364103719515e-05,
+      "loss": 0.0117,
+      "step": 124
+    },
+    {
+      "epoch": 0.984251968503937,
+      "grad_norm": 0.1571267992258072,
+      "learning_rate": 1.5621534122318682e-05,
+      "loss": 0.0049,
+      "step": 125
+    },
+    {
+      "epoch": 0.9921259842519685,
+      "grad_norm": 1.2177189588546753,
+      "learning_rate": 1.5551301048914863e-05,
+      "loss": 0.0161,
+      "step": 126
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.414489209651947,
+      "learning_rate": 1.5480669919569313e-05,
+      "loss": 0.0181,
+      "step": 127
+    },
+    {
+      "epoch": 1.0078740157480315,
+      "grad_norm": 0.10985995829105377,
+      "learning_rate": 1.54096457988859e-05,
+      "loss": 0.0049,
+      "step": 128
+    },
+    {
+      "epoch": 1.015748031496063,
+      "grad_norm": 0.12780147790908813,
+      "learning_rate": 1.533823377964791e-05,
+      "loss": 0.0026,
+      "step": 129
+    },
+    {
+      "epoch": 1.015748031496063,
+      "eval_loss": 0.33064374327659607,
+      "eval_runtime": 6.9286,
+      "eval_samples_per_second": 23.526,
+      "eval_steps_per_second": 3.031,
+      "step": 129
+    },
+    {
+      "epoch": 1.0236220472440944,
+      "grad_norm": 0.5142458081245422,
+      "learning_rate": 1.52664389824529e-05,
+      "loss": 0.0082,
+      "step": 130
+    },
+    {
+      "epoch": 1.031496062992126,
+      "grad_norm": 0.15617145597934723,
+      "learning_rate": 1.5194266555345505e-05,
+      "loss": 0.0016,
+      "step": 131
+    },
+    {
+      "epoch": 1.0393700787401574,
+      "grad_norm": 0.5782387852668762,
+      "learning_rate": 1.5121721673448319e-05,
+      "loss": 0.0117,
+      "step": 132
+    },
+    {
+      "epoch": 1.047244094488189,
+      "grad_norm": 0.08414836972951889,
+      "learning_rate": 1.5048809538590789e-05,
+      "loss": 0.0021,
+      "step": 133
+    },
+    {
+      "epoch": 1.0551181102362204,
+      "grad_norm": 0.28253939747810364,
+      "learning_rate": 1.4975535378936228e-05,
+      "loss": 0.0055,
+      "step": 134
+    },
+    {
+      "epoch": 1.0629921259842519,
+      "grad_norm": 0.47917842864990234,
+      "learning_rate": 1.490190444860694e-05,
+      "loss": 0.0046,
+      "step": 135
+    },
+    {
+      "epoch": 1.0708661417322836,
+      "grad_norm": 0.1895662248134613,
+      "learning_rate": 1.482792202730745e-05,
+      "loss": 0.006,
+      "step": 136
+    },
+    {
+      "epoch": 1.078740157480315,
+      "grad_norm": 0.13722768425941467,
+      "learning_rate": 1.475359341994595e-05,
+      "loss": 0.0031,
+      "step": 137
+    },
+    {
+      "epoch": 1.0866141732283465,
+      "grad_norm": 0.10731153190135956,
+      "learning_rate": 1.4678923956253894e-05,
+      "loss": 0.0005,
+      "step": 138
+    },
+    {
+      "epoch": 1.094488188976378,
+      "grad_norm": 0.12261265516281128,
+      "learning_rate": 1.460391899040383e-05,
+      "loss": 0.0031,
+      "step": 139
+    },
+    {
+      "epoch": 1.1023622047244095,
+      "grad_norm": 0.0038245893083512783,
+      "learning_rate": 1.4528583900625481e-05,
+      "loss": 0.0,
+      "step": 140
+    },
+    {
+      "epoch": 1.110236220472441,
+      "grad_norm": 0.28762558102607727,
+      "learning_rate": 1.4452924088820101e-05,
+      "loss": 0.004,
+      "step": 141
+    },
+    {
+      "epoch": 1.1181102362204725,
+      "grad_norm": 0.17267552018165588,
+      "learning_rate": 1.4376944980173138e-05,
+      "loss": 0.0002,
+      "step": 142
+    },
+    {
+      "epoch": 1.125984251968504,
+      "grad_norm": 0.12727122008800507,
+      "learning_rate": 1.4300652022765207e-05,
+      "loss": 0.0029,
+      "step": 143
+    },
+    {
+      "epoch": 1.1338582677165354,
+      "grad_norm": 0.25049135088920593,
+      "learning_rate": 1.4224050687181442e-05,
+      "loss": 0.0108,
+      "step": 144
+    },
+    {
+      "epoch": 1.141732283464567,
+      "grad_norm": 0.16092728078365326,
+      "learning_rate": 1.4147146466119235e-05,
+      "loss": 0.0024,
+      "step": 145
+    },
+    {
+      "epoch": 1.1496062992125984,
+      "grad_norm": 0.13642658293247223,
+      "learning_rate": 1.406994487399437e-05,
+      "loss": 0.0037,
+      "step": 146
+    },
+    {
+      "epoch": 1.1574803149606299,
+      "grad_norm": 0.9029403328895569,
+      "learning_rate": 1.3992451446545624e-05,
+      "loss": 0.0034,
+      "step": 147
+    },
+    {
+      "epoch": 1.1653543307086613,
+      "grad_norm": 0.19518424570560455,
+      "learning_rate": 1.3914671740437811e-05,
+      "loss": 0.0057,
+      "step": 148
+    },
+    {
+      "epoch": 1.1732283464566928,
+      "grad_norm": 0.12140502035617828,
+      "learning_rate": 1.3836611332863356e-05,
+      "loss": 0.0041,
+      "step": 149
+    },
+    {
+      "epoch": 1.1811023622047245,
+      "grad_norm": 0.5148038864135742,
+      "learning_rate": 1.3758275821142382e-05,
+      "loss": 0.0026,
+      "step": 150
+    },
+    {
+      "epoch": 1.188976377952756,
+      "grad_norm": 1.828904390335083,
+      "learning_rate": 1.3679670822321347e-05,
+      "loss": 0.0024,
+      "step": 151
+    },
+    {
+      "epoch": 1.1968503937007875,
+      "grad_norm": 0.3571717143058777,
+      "learning_rate": 1.3600801972770272e-05,
+      "loss": 0.0106,
+      "step": 152
+    },
+    {
+      "epoch": 1.204724409448819,
+      "grad_norm": 0.051027003675699234,
+      "learning_rate": 1.3521674927778594e-05,
+      "loss": 0.0003,
+      "step": 153
+    },
+    {
+      "epoch": 1.2125984251968505,
+      "grad_norm": 0.6490982174873352,
+      "learning_rate": 1.3442295361149651e-05,
+      "loss": 0.0035,
+      "step": 154
+    },
+    {
+      "epoch": 1.220472440944882,
+      "grad_norm": 0.08408445864915848,
+      "learning_rate": 1.336266896479384e-05,
+      "loss": 0.0027,
+      "step": 155
+    },
+    {
+      "epoch": 1.2283464566929134,
+      "grad_norm": 0.09666562080383301,
+      "learning_rate": 1.328280144832047e-05,
+      "loss": 0.0019,
+      "step": 156
+    },
+    {
+      "epoch": 1.236220472440945,
+      "grad_norm": 0.03880690038204193,
+      "learning_rate": 1.3202698538628376e-05,
+      "loss": 0.0003,
+      "step": 157
+    },
+    {
+      "epoch": 1.2440944881889764,
+      "grad_norm": 0.11940775066614151,
+      "learning_rate": 1.3122365979495259e-05,
+      "loss": 0.0024,
+      "step": 158
+    },
+    {
+      "epoch": 1.2519685039370079,
+      "grad_norm": 0.1442880481481552,
+      "learning_rate": 1.3041809531165819e-05,
+      "loss": 0.0015,
+      "step": 159
+    },
+    {
+      "epoch": 1.2598425196850394,
+      "grad_norm": 0.1961939036846161,
+      "learning_rate": 1.2961034969938732e-05,
+      "loss": 0.0056,
+      "step": 160
+    },
+    {
+      "epoch": 1.2677165354330708,
+      "grad_norm": 0.26947638392448425,
+      "learning_rate": 1.288004808775246e-05,
+      "loss": 0.0028,
+      "step": 161
+    },
+    {
+      "epoch": 1.2755905511811023,
+      "grad_norm": 0.5154056549072266,
+      "learning_rate": 1.2798854691769927e-05,
+      "loss": 0.0037,
+      "step": 162
+    },
+    {
+      "epoch": 1.2834645669291338,
+      "grad_norm": 0.4292369782924652,
+      "learning_rate": 1.2717460603962132e-05,
+      "loss": 0.0029,
+      "step": 163
+    },
+    {
+      "epoch": 1.2913385826771653,
+      "grad_norm": 0.19139212369918823,
+      "learning_rate": 1.2635871660690677e-05,
+      "loss": 0.0061,
+      "step": 164
+    },
+    {
+      "epoch": 1.2992125984251968,
+      "grad_norm": 0.19960306584835052,
+      "learning_rate": 1.2554093712289267e-05,
+      "loss": 0.005,
+      "step": 165
+    },
+    {
+      "epoch": 1.3070866141732282,
+      "grad_norm": 0.4523830711841583,
+      "learning_rate": 1.2472132622644222e-05,
+      "loss": 0.0065,
+      "step": 166
+    },
+    {
+      "epoch": 1.3149606299212597,
+      "grad_norm": 0.49343299865722656,
+      "learning_rate": 1.2389994268773995e-05,
+      "loss": 0.0061,
+      "step": 167
+    },
+    {
+      "epoch": 1.3228346456692912,
+      "grad_norm": 0.01938088797032833,
+      "learning_rate": 1.2307684540407775e-05,
+      "loss": 0.0001,
+      "step": 168
+    },
+    {
+      "epoch": 1.330708661417323,
+      "grad_norm": 0.3082112669944763,
+      "learning_rate": 1.2225209339563144e-05,
+      "loss": 0.0053,
+      "step": 169
+    },
+    {
+      "epoch": 1.3385826771653544,
+      "grad_norm": 0.01982509344816208,
+      "learning_rate": 1.2142574580122903e-05,
+      "loss": 0.0001,
+      "step": 170
+    },
+    {
+      "epoch": 1.3464566929133859,
+      "grad_norm": 0.12388588488101959,
+      "learning_rate": 1.2059786187410984e-05,
+      "loss": 0.0049,
+      "step": 171
+    },
+    {
+      "epoch": 1.3543307086614174,
+      "grad_norm": 0.43759095668792725,
+      "learning_rate": 1.1976850097767598e-05,
+      "loss": 0.0128,
+      "step": 172
+    },
+    {
+      "epoch": 1.3543307086614174,
+      "eval_loss": 0.3166251480579376,
+      "eval_runtime": 6.9515,
+      "eval_samples_per_second": 23.448,
+      "eval_steps_per_second": 3.021,
+      "step": 172
+    },
+    {
+      "epoch": 1.3622047244094488,
+      "grad_norm": 0.46561670303344727,
+      "learning_rate": 1.1893772258123554e-05,
+      "loss": 0.008,
+      "step": 173
+    },
+    {
+      "epoch": 1.3700787401574803,
+      "grad_norm": 0.16612188518047333,
+      "learning_rate": 1.1810558625573856e-05,
+      "loss": 0.0024,
+      "step": 174
+    },
+    {
+      "epoch": 1.3779527559055118,
+      "grad_norm": 0.13628093898296356,
+      "learning_rate": 1.1727215166950519e-05,
+      "loss": 0.0045,
+      "step": 175
+    },
+    {
+      "epoch": 1.3858267716535433,
+      "grad_norm": 0.565229058265686,
+      "learning_rate": 1.1643747858394743e-05,
+      "loss": 0.0103,
+      "step": 176
+    },
+    {
+      "epoch": 1.3937007874015748,
+      "grad_norm": 0.14550763368606567,
+      "learning_rate": 1.156016268492839e-05,
+      "loss": 0.0028,
+      "step": 177
+    },
+    {
+      "epoch": 1.4015748031496063,
+      "grad_norm": 0.12460129708051682,
+      "learning_rate": 1.1476465640024814e-05,
+      "loss": 0.0031,
+      "step": 178
+    },
+    {
+      "epoch": 1.4094488188976377,
+      "grad_norm": 0.19089221954345703,
+      "learning_rate": 1.1392662725179114e-05,
+      "loss": 0.0035,
+      "step": 179
+    },
+    {
+      "epoch": 1.4173228346456692,
+      "grad_norm": 0.6106573343276978,
+      "learning_rate": 1.1308759949477786e-05,
+      "loss": 0.0088,
+      "step": 180
+    },
+    {
+      "epoch": 1.425196850393701,
+      "grad_norm": 0.20053207874298096,
+      "learning_rate": 1.1224763329167859e-05,
+      "loss": 0.0033,
+      "step": 181
+    },
+    {
+      "epoch": 1.4330708661417324,
+      "grad_norm": 0.1984691321849823,
+      "learning_rate": 1.1140678887225468e-05,
+      "loss": 0.0051,
+      "step": 182
+    },
+    {
+      "epoch": 1.4409448818897639,
+      "grad_norm": 0.19264858961105347,
+      "learning_rate": 1.1056512652924014e-05,
+      "loss": 0.0046,
+      "step": 183
+    },
+    {
+      "epoch": 1.4488188976377954,
+      "grad_norm": 0.10979076474905014,
+      "learning_rate": 1.0972270661401812e-05,
+      "loss": 0.0031,
+      "step": 184
+    },
+    {
+      "epoch": 1.4566929133858268,
+      "grad_norm": 0.1744084656238556,
+      "learning_rate": 1.0887958953229349e-05,
+      "loss": 0.0024,
+      "step": 185
+    },
+    {
+      "epoch": 1.4645669291338583,
+      "grad_norm": 0.20646224915981293,
+      "learning_rate": 1.0803583573976137e-05,
+      "loss": 0.008,
+      "step": 186
+    },
+    {
+      "epoch": 1.4724409448818898,
+      "grad_norm": 0.14391584694385529,
+      "learning_rate": 1.0719150573777226e-05,
+      "loss": 0.004,
+      "step": 187
+    },
+    {
+      "epoch": 1.4803149606299213,
+      "grad_norm": 0.36887863278388977,
+      "learning_rate": 1.0634666006899375e-05,
+      "loss": 0.0074,
+      "step": 188
+    },
+    {
+      "epoch": 1.4881889763779528,
+      "grad_norm": 0.21352627873420715,
+      "learning_rate": 1.055013593130693e-05,
+      "loss": 0.0082,
+      "step": 189
+    },
+    {
+      "epoch": 1.4960629921259843,
+      "grad_norm": 0.22443020343780518,
+      "learning_rate": 1.046556640822744e-05,
+      "loss": 0.0087,
+      "step": 190
+    },
+    {
+      "epoch": 1.5039370078740157,
+      "grad_norm": 0.4243764281272888,
+      "learning_rate": 1.0380963501717034e-05,
+      "loss": 0.0068,
+      "step": 191
+    },
+    {
+      "epoch": 1.5118110236220472,
+      "grad_norm": 0.17558562755584717,
+      "learning_rate": 1.0296333278225599e-05,
+      "loss": 0.0054,
+      "step": 192
+    },
+    {
+      "epoch": 1.5196850393700787,
+      "grad_norm": 0.14842620491981506,
+      "learning_rate": 1.0211681806161787e-05,
+      "loss": 0.0031,
+      "step": 193
+    },
+    {
+      "epoch": 1.5275590551181102,
+      "grad_norm": 0.09316081553697586,
+      "learning_rate": 1.0127015155457875e-05,
+      "loss": 0.0013,
+      "step": 194
+    },
+    {
+      "epoch": 1.5354330708661417,
+      "grad_norm": 0.19795025885105133,
+      "learning_rate": 1.0042339397134528e-05,
+      "loss": 0.0051,
+      "step": 195
+    },
+    {
+      "epoch": 1.5433070866141732,
+      "grad_norm": 0.21606990694999695,
+      "learning_rate": 9.957660602865477e-06,
+      "loss": 0.0041,
+      "step": 196
+    },
+    {
+      "epoch": 1.5511811023622046,
+      "grad_norm": 0.18036173284053802,
+      "learning_rate": 9.872984844542128e-06,
+      "loss": 0.0037,
+      "step": 197
+    },
+    {
+      "epoch": 1.5590551181102361,
+      "grad_norm": 0.18953870236873627,
+      "learning_rate": 9.788318193838218e-06,
+      "loss": 0.0041,
+      "step": 198
+    },
+    {
+      "epoch": 1.5669291338582676,
+      "grad_norm": 0.12346503138542175,
+      "learning_rate": 9.703666721774403e-06,
+      "loss": 0.0035,
+      "step": 199
+    },
+    {
+      "epoch": 1.574803149606299,
+      "grad_norm": 0.4576225280761719,
+      "learning_rate": 9.619036498282968e-06,
+      "loss": 0.0041,
+      "step": 200
+    },
+    {
+      "epoch": 1.5826771653543306,
+      "grad_norm": 0.10333681106567383,
+      "learning_rate": 9.534433591772562e-06,
+      "loss": 0.0011,
+      "step": 201
+    },
+    {
+      "epoch": 1.590551181102362,
+      "grad_norm": 0.19167865812778473,
+      "learning_rate": 9.449864068693072e-06,
+      "loss": 0.0062,
+      "step": 202
+    },
+    {
+      "epoch": 1.5984251968503937,
+      "grad_norm": 0.2258184254169464,
+      "learning_rate": 9.365333993100628e-06,
+      "loss": 0.003,
+      "step": 203
+    },
+    {
+      "epoch": 1.6062992125984252,
+      "grad_norm": 0.07945302873849869,
+      "learning_rate": 9.280849426222778e-06,
+      "loss": 0.0008,
+      "step": 204
+    },
+    {
+      "epoch": 1.6141732283464567,
+      "grad_norm": 0.17767398059368134,
+      "learning_rate": 9.196416426023868e-06,
+      "loss": 0.0053,
+      "step": 205
+    },
+    {
+      "epoch": 1.6220472440944882,
+      "grad_norm": 0.12704500555992126,
+      "learning_rate": 9.112041046770653e-06,
+      "loss": 0.0023,
+      "step": 206
+    },
+    {
+      "epoch": 1.6299212598425197,
+      "grad_norm": 0.4054742753505707,
+      "learning_rate": 9.027729338598188e-06,
+      "loss": 0.0045,
+      "step": 207
+    },
+    {
+      "epoch": 1.6377952755905512,
+      "grad_norm": 0.4463757574558258,
+      "learning_rate": 8.943487347075988e-06,
+      "loss": 0.007,
+      "step": 208
+    },
+    {
+      "epoch": 1.6456692913385826,
+      "grad_norm": 0.6517045497894287,
+      "learning_rate": 8.859321112774535e-06,
+      "loss": 0.0052,
+      "step": 209
+    },
+    {
+      "epoch": 1.6535433070866141,
+      "grad_norm": 0.1542089730501175,
+      "learning_rate": 8.775236670832146e-06,
+      "loss": 0.0047,
+      "step": 210
+    },
+    {
+      "epoch": 1.6614173228346458,
+      "grad_norm": 0.14716440439224243,
+      "learning_rate": 8.691240050522215e-06,
+      "loss": 0.0049,
+      "step": 211
+    },
+    {
+      "epoch": 1.6692913385826773,
+      "grad_norm": 0.2997347116470337,
+      "learning_rate": 8.607337274820888e-06,
+      "loss": 0.0076,
+      "step": 212
+    },
+    {
+      "epoch": 1.6771653543307088,
+      "grad_norm": 0.22548256814479828,
+      "learning_rate": 8.52353435997519e-06,
+      "loss": 0.0063,
+      "step": 213
+    },
+    {
+      "epoch": 1.6850393700787403,
+      "grad_norm": 0.7220733165740967,
+      "learning_rate": 8.439837315071612e-06,
+      "loss": 0.0089,
+      "step": 214
+    },
+    {
+      "epoch": 1.6929133858267718,
+      "grad_norm": 0.5101618766784668,
+      "learning_rate": 8.35625214160526e-06,
+      "loss": 0.0042,
+      "step": 215
+    },
+    {
+      "epoch": 1.6929133858267718,
+      "eval_loss": 0.3484288156032562,
+      "eval_runtime": 6.4482,
+      "eval_samples_per_second": 25.278,
+      "eval_steps_per_second": 3.257,
+      "step": 215
+    },
+    {
+      "epoch": 1.7007874015748032,
+      "grad_norm": 0.1698393076658249,
+      "learning_rate": 8.272784833049485e-06,
+      "loss": 0.0028,
+      "step": 216
+    },
+    {
+      "epoch": 1.7086614173228347,
+      "grad_norm": 0.5772718191146851,
+      "learning_rate": 8.18944137442615e-06,
+      "loss": 0.0082,
+      "step": 217
+    },
+    {
+      "epoch": 1.7165354330708662,
+      "grad_norm": 0.09606469422578812,
+      "learning_rate": 8.106227741876447e-06,
+      "loss": 0.0011,
+      "step": 218
+    },
+    {
+      "epoch": 1.7244094488188977,
+      "grad_norm": 0.14510361850261688,
+      "learning_rate": 8.023149902232404e-06,
+      "loss": 0.0015,
+      "step": 219
+    },
+    {
+      "epoch": 1.7322834645669292,
+      "grad_norm": 0.055804118514060974,
+      "learning_rate": 7.940213812589018e-06,
+      "loss": 0.0008,
+      "step": 220
+    },
+    {
+      "epoch": 1.7401574803149606,
+      "grad_norm": 0.13318321108818054,
+      "learning_rate": 7.857425419877097e-06,
+      "loss": 0.005,
+      "step": 221
+    },
+    {
+      "epoch": 1.7480314960629921,
+      "grad_norm": 0.23600782454013824,
+      "learning_rate": 7.774790660436857e-06,
+      "loss": 0.0063,
+      "step": 222
+    },
+    {
+      "epoch": 1.7559055118110236,
+      "grad_norm": 0.8483791351318359,
+      "learning_rate": 7.69231545959223e-06,
+      "loss": 0.0027,
+      "step": 223
+    },
+    {
+      "epoch": 1.763779527559055,
+      "grad_norm": 0.16536197066307068,
+      "learning_rate": 7.610005731226009e-06,
+      "loss": 0.0039,
+      "step": 224
+    },
+    {
+      "epoch": 1.7716535433070866,
+      "grad_norm": 0.14446765184402466,
+      "learning_rate": 7.52786737735578e-06,
+      "loss": 0.0036,
+      "step": 225
+    },
+    {
+      "epoch": 1.779527559055118,
+      "grad_norm": 0.8880365490913391,
+      "learning_rate": 7.445906287710733e-06,
+      "loss": 0.0061,
+      "step": 226
+    },
+    {
+      "epoch": 1.7874015748031495,
+      "grad_norm": 0.151743084192276,
+      "learning_rate": 7.364128339309326e-06,
+      "loss": 0.0028,
+      "step": 227
+    },
+    {
+      "epoch": 1.795275590551181,
+      "grad_norm": 0.1224551647901535,
+      "learning_rate": 7.282539396037868e-06,
+      "loss": 0.002,
+      "step": 228
+    },
+    {
+      "epoch": 1.8031496062992125,
+      "grad_norm": 0.4868486225605011,
+      "learning_rate": 7.201145308230075e-06,
+      "loss": 0.0031,
+      "step": 229
+    },
+    {
+      "epoch": 1.811023622047244,
+      "grad_norm": 0.2875569462776184,
+      "learning_rate": 7.119951912247545e-06,
+      "loss": 0.0082,
+      "step": 230
+    },
+    {
+      "epoch": 1.8188976377952755,
+      "grad_norm": 0.43524420261383057,
+      "learning_rate": 7.038965030061273e-06,
+      "loss": 0.0075,
+      "step": 231
+    },
+    {
+      "epoch": 1.826771653543307,
+      "grad_norm": 0.39634883403778076,
+      "learning_rate": 6.9581904688341854e-06,
+      "loss": 0.0032,
+      "step": 232
+    },
+    {
+      "epoch": 1.8346456692913384,
+      "grad_norm": 0.9809433817863464,
+      "learning_rate": 6.8776340205047446e-06,
+      "loss": 0.0085,
+      "step": 233
+    },
+    {
+      "epoch": 1.84251968503937,
+      "grad_norm": 0.20062875747680664,
+      "learning_rate": 6.797301461371626e-06,
+      "loss": 0.0043,
+      "step": 234
+    },
+    {
+      "epoch": 1.8503937007874016,
+      "grad_norm": 0.148948073387146,
+      "learning_rate": 6.7171985516795315e-06,
+      "loss": 0.0036,
+      "step": 235
+    },
+    {
+      "epoch": 1.858267716535433,
+      "grad_norm": 0.15658679604530334,
+      "learning_rate": 6.637331035206166e-06,
+      "loss": 0.0046,
+      "step": 236
+    },
+    {
+      "epoch": 1.8661417322834646,
+      "grad_norm": 0.22365815937519073,
+      "learning_rate": 6.557704638850352e-06,
+      "loss": 0.0081,
+      "step": 237
+    },
+    {
+      "epoch": 1.874015748031496,
+      "grad_norm": 0.10596666485071182,
+      "learning_rate": 6.4783250722214066e-06,
+      "loss": 0.0032,
+      "step": 238
+    },
+    {
+      "epoch": 1.8818897637795275,
+      "grad_norm": 0.2130754142999649,
+      "learning_rate": 6.399198027229732e-06,
+      "loss": 0.0056,
+      "step": 239
+    },
+    {
+      "epoch": 1.889763779527559,
+      "grad_norm": 0.05641167238354683,
+      "learning_rate": 6.320329177678656e-06,
+      "loss": 0.0008,
+      "step": 240
+    },
+    {
+      "epoch": 1.8976377952755905,
+      "grad_norm": 0.10349344462156296,
+      "learning_rate": 6.241724178857621e-06,
+      "loss": 0.0026,
+      "step": 241
+    },
+    {
+      "epoch": 1.905511811023622,
+      "grad_norm": 0.08451675623655319,
+      "learning_rate": 6.163388667136646e-06,
+      "loss": 0.0016,
+      "step": 242
+    },
+    {
+      "epoch": 1.9133858267716537,
+      "grad_norm": 0.13671623170375824,
+      "learning_rate": 6.085328259562195e-06,
+      "loss": 0.0034,
+      "step": 243
+    },
+    {
+      "epoch": 1.9212598425196852,
+      "grad_norm": 0.5500523447990417,
+      "learning_rate": 6.007548553454379e-06,
+      "loss": 0.0028,
+      "step": 244
+    },
+    {
+      "epoch": 1.9291338582677167,
+      "grad_norm": 0.06702329218387604,
+      "learning_rate": 5.93005512600563e-06,
+      "loss": 0.0009,
+      "step": 245
+    },
+    {
+      "epoch": 1.9370078740157481,
+      "grad_norm": 0.15156973898410797,
+      "learning_rate": 5.852853533880768e-06,
+      "loss": 0.0064,
+      "step": 246
+    },
+    {
+      "epoch": 1.9448818897637796,
+      "grad_norm": 0.2970314621925354,
+      "learning_rate": 5.7759493128185584e-06,
+      "loss": 0.0077,
+      "step": 247
+    },
+    {
+      "epoch": 1.952755905511811,
+      "grad_norm": 0.06406261771917343,
+      "learning_rate": 5.699347977234799e-06,
+      "loss": 0.0006,
+      "step": 248
+    },
+    {
+      "epoch": 1.9606299212598426,
+      "grad_norm": 0.2910393178462982,
+      "learning_rate": 5.623055019826862e-06,
+      "loss": 0.0036,
+      "step": 249
+    },
+    {
+      "epoch": 1.968503937007874,
+      "grad_norm": 0.6454993486404419,
+      "learning_rate": 5.547075911179902e-06,
+      "loss": 0.0084,
+      "step": 250
+    },
+    {
+      "epoch": 1.9763779527559056,
+      "grad_norm": 0.09460143744945526,
+      "learning_rate": 5.471416099374525e-06,
+      "loss": 0.0021,
+      "step": 251
+    },
+    {
+      "epoch": 1.984251968503937,
+      "grad_norm": 0.2024363875389099,
+      "learning_rate": 5.3960810095961705e-06,
+      "loss": 0.0052,
+      "step": 252
+    },
+    {
+      "epoch": 1.9921259842519685,
+      "grad_norm": 0.09423142671585083,
+      "learning_rate": 5.321076043746108e-06,
+      "loss": 0.0018,
+      "step": 253
+    },
+    {
+      "epoch": 2.0,
+      "grad_norm": 0.1085880920290947,
+      "learning_rate": 5.246406580054051e-06,
+      "loss": 0.0039,
+      "step": 254
+    }
+  ],
+  "logging_steps": 1,
+  "max_steps": 381,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 3,
+  "save_steps": 127,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 2.6174542139778662e+17,
+  "train_batch_size": 128,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-254/training_args.bin

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

checkpoint-254/zero_to_fp32.py

@@ -0,0 +1,760 @@
+#!/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 . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+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, weights_only=False)
+
+        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 tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # 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}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] 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, exclude_frozen_parameters):
+    """
+    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,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+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 _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+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() if _has_callable(shape, 'numel') else math.prod(shape)
+            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,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _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")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+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 = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * 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
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        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}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        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,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _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 to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    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``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    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.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    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")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    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_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``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``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        # an memory-efficient approach for sharding
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+        shard_state_dict = to_torch_tensor(shard_state_dict)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+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_dir",
+                        type=str,
+                        help="directory to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    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("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    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_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)

checkpoint-381/added_tokens.json

@@ -0,0 +1,24 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

checkpoint-381/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "Qwen/Qwen2.5-1.5B-Instruct",
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.1",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151665
+}

checkpoint-381/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "4.48.1"
+}

checkpoint-381/latest

@@ -0,0 +1 @@
+global_step381

checkpoint-381/model.safetensors

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

checkpoint-381/rng_state_0.pth

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

checkpoint-381/rng_state_1.pth

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

checkpoint-381/scheduler.pt

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

checkpoint-381/special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint-381/tokenizer.json

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

checkpoint-381/tokenizer_config.json

@@ -0,0 +1,208 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

checkpoint-381/trainer_state.json

@@ -0,0 +1,2772 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 3.0,
+  "eval_steps": 43,
+  "global_step": 381,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.007874015748031496,
+      "grad_norm": 118.11203002929688,
+      "learning_rate": 2.0000000000000003e-06,
+      "loss": 4.6099,
+      "step": 1
+    },
+    {
+      "epoch": 0.007874015748031496,
+      "eval_loss": 3.1001100540161133,
+      "eval_runtime": 5.3966,
+      "eval_samples_per_second": 30.204,
+      "eval_steps_per_second": 3.891,
+      "step": 1
+    },
+    {
+      "epoch": 0.015748031496062992,
+      "grad_norm": 118.4310302734375,
+      "learning_rate": 4.000000000000001e-06,
+      "loss": 4.5857,
+      "step": 2
+    },
+    {
+      "epoch": 0.023622047244094488,
+      "grad_norm": 103.37439727783203,
+      "learning_rate": 6e-06,
+      "loss": 4.3069,
+      "step": 3
+    },
+    {
+      "epoch": 0.031496062992125984,
+      "grad_norm": 75.05075073242188,
+      "learning_rate": 8.000000000000001e-06,
+      "loss": 3.8754,
+      "step": 4
+    },
+    {
+      "epoch": 0.03937007874015748,
+      "grad_norm": 50.459983825683594,
+      "learning_rate": 1e-05,
+      "loss": 3.2841,
+      "step": 5
+    },
+    {
+      "epoch": 0.047244094488188976,
+      "grad_norm": 47.4603385925293,
+      "learning_rate": 1.2e-05,
+      "loss": 2.4285,
+      "step": 6
+    },
+    {
+      "epoch": 0.05511811023622047,
+      "grad_norm": 32.362667083740234,
+      "learning_rate": 1.4e-05,
+      "loss": 1.8177,
+      "step": 7
+    },
+    {
+      "epoch": 0.06299212598425197,
+      "grad_norm": 22.846933364868164,
+      "learning_rate": 1.6000000000000003e-05,
+      "loss": 1.1567,
+      "step": 8
+    },
+    {
+      "epoch": 0.07086614173228346,
+      "grad_norm": 17.060213088989258,
+      "learning_rate": 1.8e-05,
+      "loss": 0.8257,
+      "step": 9
+    },
+    {
+      "epoch": 0.07874015748031496,
+      "grad_norm": 14.415579795837402,
+      "learning_rate": 2e-05,
+      "loss": 0.4257,
+      "step": 10
+    },
+    {
+      "epoch": 0.08661417322834646,
+      "grad_norm": 7.753712177276611,
+      "learning_rate": 1.999964147509006e-05,
+      "loss": 0.2976,
+      "step": 11
+    },
+    {
+      "epoch": 0.09448818897637795,
+      "grad_norm": 26.883708953857422,
+      "learning_rate": 1.9998565926068253e-05,
+      "loss": 0.3365,
+      "step": 12
+    },
+    {
+      "epoch": 0.10236220472440945,
+      "grad_norm": 10.675631523132324,
+      "learning_rate": 1.9996773430056806e-05,
+      "loss": 0.2161,
+      "step": 13
+    },
+    {
+      "epoch": 0.11023622047244094,
+      "grad_norm": 6.670111179351807,
+      "learning_rate": 1.999426411558661e-05,
+      "loss": 0.1816,
+      "step": 14
+    },
+    {
+      "epoch": 0.11811023622047244,
+      "grad_norm": 8.878239631652832,
+      "learning_rate": 1.9991038162588018e-05,
+      "loss": 0.1567,
+      "step": 15
+    },
+    {
+      "epoch": 0.12598425196850394,
+      "grad_norm": 2.9917383193969727,
+      "learning_rate": 1.9987095802377933e-05,
+      "loss": 0.0813,
+      "step": 16
+    },
+    {
+      "epoch": 0.13385826771653545,
+      "grad_norm": 1.0548763275146484,
+      "learning_rate": 1.9982437317643218e-05,
+      "loss": 0.0217,
+      "step": 17
+    },
+    {
+      "epoch": 0.14173228346456693,
+      "grad_norm": 2.8778488636016846,
+      "learning_rate": 1.9977063042420438e-05,
+      "loss": 0.0618,
+      "step": 18
+    },
+    {
+      "epoch": 0.14960629921259844,
+      "grad_norm": 0.9811734557151794,
+      "learning_rate": 1.99709733620719e-05,
+      "loss": 0.0175,
+      "step": 19
+    },
+    {
+      "epoch": 0.15748031496062992,
+      "grad_norm": 0.7218202948570251,
+      "learning_rate": 1.996416871325803e-05,
+      "loss": 0.0302,
+      "step": 20
+    },
+    {
+      "epoch": 0.16535433070866143,
+      "grad_norm": 1.2746995687484741,
+      "learning_rate": 1.995664958390604e-05,
+      "loss": 0.0453,
+      "step": 21
+    },
+    {
+      "epoch": 0.1732283464566929,
+      "grad_norm": 0.9413469433784485,
+      "learning_rate": 1.9948416513174976e-05,
+      "loss": 0.0175,
+      "step": 22
+    },
+    {
+      "epoch": 0.18110236220472442,
+      "grad_norm": 1.4161137342453003,
+      "learning_rate": 1.9939470091417012e-05,
+      "loss": 0.0277,
+      "step": 23
+    },
+    {
+      "epoch": 0.1889763779527559,
+      "grad_norm": 2.2721235752105713,
+      "learning_rate": 1.992981096013517e-05,
+      "loss": 0.0589,
+      "step": 24
+    },
+    {
+      "epoch": 0.1968503937007874,
+      "grad_norm": 1.143970251083374,
+      "learning_rate": 1.9919439811937283e-05,
+      "loss": 0.0182,
+      "step": 25
+    },
+    {
+      "epoch": 0.2047244094488189,
+      "grad_norm": 0.8054028749465942,
+      "learning_rate": 1.9908357390486342e-05,
+      "loss": 0.0211,
+      "step": 26
+    },
+    {
+      "epoch": 0.2125984251968504,
+      "grad_norm": 1.4449081420898438,
+      "learning_rate": 1.989656449044718e-05,
+      "loss": 0.0244,
+      "step": 27
+    },
+    {
+      "epoch": 0.2204724409448819,
+      "grad_norm": 0.49216631054878235,
+      "learning_rate": 1.988406195742948e-05,
+      "loss": 0.005,
+      "step": 28
+    },
+    {
+      "epoch": 0.2283464566929134,
+      "grad_norm": 0.9945647716522217,
+      "learning_rate": 1.987085068792715e-05,
+      "loss": 0.0373,
+      "step": 29
+    },
+    {
+      "epoch": 0.23622047244094488,
+      "grad_norm": 1.1753748655319214,
+      "learning_rate": 1.9856931629254032e-05,
+      "loss": 0.0217,
+      "step": 30
+    },
+    {
+      "epoch": 0.2440944881889764,
+      "grad_norm": 0.5960403680801392,
+      "learning_rate": 1.984230577947597e-05,
+      "loss": 0.0157,
+      "step": 31
+    },
+    {
+      "epoch": 0.25196850393700787,
+      "grad_norm": 0.3657272160053253,
+      "learning_rate": 1.9826974187339267e-05,
+      "loss": 0.0082,
+      "step": 32
+    },
+    {
+      "epoch": 0.25984251968503935,
+      "grad_norm": 1.1290266513824463,
+      "learning_rate": 1.981093795219546e-05,
+      "loss": 0.0236,
+      "step": 33
+    },
+    {
+      "epoch": 0.2677165354330709,
+      "grad_norm": 1.673962116241455,
+      "learning_rate": 1.9794198223922496e-05,
+      "loss": 0.0182,
+      "step": 34
+    },
+    {
+      "epoch": 0.2755905511811024,
+      "grad_norm": 0.540355384349823,
+      "learning_rate": 1.9776756202842297e-05,
+      "loss": 0.011,
+      "step": 35
+    },
+    {
+      "epoch": 0.28346456692913385,
+      "grad_norm": 0.3380790054798126,
+      "learning_rate": 1.9758613139634662e-05,
+      "loss": 0.0048,
+      "step": 36
+    },
+    {
+      "epoch": 0.29133858267716534,
+      "grad_norm": 1.886232852935791,
+      "learning_rate": 1.9739770335247616e-05,
+      "loss": 0.0157,
+      "step": 37
+    },
+    {
+      "epoch": 0.2992125984251969,
+      "grad_norm": 2.140639305114746,
+      "learning_rate": 1.972022914080411e-05,
+      "loss": 0.0393,
+      "step": 38
+    },
+    {
+      "epoch": 0.30708661417322836,
+      "grad_norm": 0.35308870673179626,
+      "learning_rate": 1.9699990957505136e-05,
+      "loss": 0.0074,
+      "step": 39
+    },
+    {
+      "epoch": 0.31496062992125984,
+      "grad_norm": 0.3918301463127136,
+      "learning_rate": 1.9679057236529266e-05,
+      "loss": 0.0083,
+      "step": 40
+    },
+    {
+      "epoch": 0.3228346456692913,
+      "grad_norm": 0.4406338632106781,
+      "learning_rate": 1.965742947892858e-05,
+      "loss": 0.0152,
+      "step": 41
+    },
+    {
+      "epoch": 0.33070866141732286,
+      "grad_norm": 0.6819682121276855,
+      "learning_rate": 1.9635109235521057e-05,
+      "loss": 0.0091,
+      "step": 42
+    },
+    {
+      "epoch": 0.33858267716535434,
+      "grad_norm": 0.6794927716255188,
+      "learning_rate": 1.961209810677934e-05,
+      "loss": 0.0071,
+      "step": 43
+    },
+    {
+      "epoch": 0.33858267716535434,
+      "eval_loss": 0.3895845115184784,
+      "eval_runtime": 6.5602,
+      "eval_samples_per_second": 24.847,
+      "eval_steps_per_second": 3.201,
+      "step": 43
+    },
+    {
+      "epoch": 0.3464566929133858,
+      "grad_norm": 0.3874967694282532,
+      "learning_rate": 1.9588397742716004e-05,
+      "loss": 0.0089,
+      "step": 44
+    },
+    {
+      "epoch": 0.3543307086614173,
+      "grad_norm": 0.5577577352523804,
+      "learning_rate": 1.9564009842765225e-05,
+      "loss": 0.0098,
+      "step": 45
+    },
+    {
+      "epoch": 0.36220472440944884,
+      "grad_norm": 0.8152347207069397,
+      "learning_rate": 1.9538936155660934e-05,
+      "loss": 0.0118,
+      "step": 46
+    },
+    {
+      "epoch": 0.3700787401574803,
+      "grad_norm": 0.2971118688583374,
+      "learning_rate": 1.951317847931141e-05,
+      "loss": 0.0084,
+      "step": 47
+    },
+    {
+      "epoch": 0.3779527559055118,
+      "grad_norm": 1.0286651849746704,
+      "learning_rate": 1.9486738660670373e-05,
+      "loss": 0.0123,
+      "step": 48
+    },
+    {
+      "epoch": 0.3858267716535433,
+      "grad_norm": 0.5227222442626953,
+      "learning_rate": 1.945961859560454e-05,
+      "loss": 0.0144,
+      "step": 49
+    },
+    {
+      "epoch": 0.3937007874015748,
+      "grad_norm": 0.461935818195343,
+      "learning_rate": 1.943182022875769e-05,
+      "loss": 0.0119,
+      "step": 50
+    },
+    {
+      "epoch": 0.4015748031496063,
+      "grad_norm": 1.2550626993179321,
+      "learning_rate": 1.940334555341122e-05,
+      "loss": 0.013,
+      "step": 51
+    },
+    {
+      "epoch": 0.4094488188976378,
+      "grad_norm": 0.37549659609794617,
+      "learning_rate": 1.9374196611341212e-05,
+      "loss": 0.0181,
+      "step": 52
+    },
+    {
+      "epoch": 0.41732283464566927,
+      "grad_norm": 0.3444191515445709,
+      "learning_rate": 1.9344375492672024e-05,
+      "loss": 0.0111,
+      "step": 53
+    },
+    {
+      "epoch": 0.4251968503937008,
+      "grad_norm": 0.3489387333393097,
+      "learning_rate": 1.9313884335726443e-05,
+      "loss": 0.0111,
+      "step": 54
+    },
+    {
+      "epoch": 0.4330708661417323,
+      "grad_norm": 0.26080814003944397,
+      "learning_rate": 1.9282725326872324e-05,
+      "loss": 0.0091,
+      "step": 55
+    },
+    {
+      "epoch": 0.4409448818897638,
+      "grad_norm": 0.1390451341867447,
+      "learning_rate": 1.9250900700365837e-05,
+      "loss": 0.0033,
+      "step": 56
+    },
+    {
+      "epoch": 0.44881889763779526,
+      "grad_norm": 0.20499111711978912,
+      "learning_rate": 1.921841273819125e-05,
+      "loss": 0.0066,
+      "step": 57
+    },
+    {
+      "epoch": 0.4566929133858268,
+      "grad_norm": 2.185487747192383,
+      "learning_rate": 1.918526376989731e-05,
+      "loss": 0.0095,
+      "step": 58
+    },
+    {
+      "epoch": 0.4645669291338583,
+      "grad_norm": 0.23939816653728485,
+      "learning_rate": 1.9151456172430186e-05,
+      "loss": 0.0048,
+      "step": 59
+    },
+    {
+      "epoch": 0.47244094488188976,
+      "grad_norm": 0.41510018706321716,
+      "learning_rate": 1.911699236996305e-05,
+      "loss": 0.0077,
+      "step": 60
+    },
+    {
+      "epoch": 0.48031496062992124,
+      "grad_norm": 0.264318585395813,
+      "learning_rate": 1.9081874833722234e-05,
+      "loss": 0.0129,
+      "step": 61
+    },
+    {
+      "epoch": 0.4881889763779528,
+      "grad_norm": 1.0443968772888184,
+      "learning_rate": 1.9046106081810047e-05,
+      "loss": 0.0035,
+      "step": 62
+    },
+    {
+      "epoch": 0.49606299212598426,
+      "grad_norm": 0.2800132632255554,
+      "learning_rate": 1.900968867902419e-05,
+      "loss": 0.0057,
+      "step": 63
+    },
+    {
+      "epoch": 0.5039370078740157,
+      "grad_norm": 1.114960789680481,
+      "learning_rate": 1.8972625236673887e-05,
+      "loss": 0.0123,
+      "step": 64
+    },
+    {
+      "epoch": 0.5118110236220472,
+      "grad_norm": 0.5027065873146057,
+      "learning_rate": 1.8934918412392596e-05,
+      "loss": 0.0052,
+      "step": 65
+    },
+    {
+      "epoch": 0.5196850393700787,
+      "grad_norm": 0.5564169883728027,
+      "learning_rate": 1.8896570909947477e-05,
+      "loss": 0.0085,
+      "step": 66
+    },
+    {
+      "epoch": 0.5275590551181102,
+      "grad_norm": 0.7567198872566223,
+      "learning_rate": 1.8857585479045493e-05,
+      "loss": 0.0054,
+      "step": 67
+    },
+    {
+      "epoch": 0.5354330708661418,
+      "grad_norm": 0.13573969900608063,
+      "learning_rate": 1.8817964915136277e-05,
+      "loss": 0.0008,
+      "step": 68
+    },
+    {
+      "epoch": 0.5433070866141733,
+      "grad_norm": 0.2704390287399292,
+      "learning_rate": 1.8777712059211643e-05,
+      "loss": 0.0078,
+      "step": 69
+    },
+    {
+      "epoch": 0.5511811023622047,
+      "grad_norm": 0.6014392971992493,
+      "learning_rate": 1.8736829797601903e-05,
+      "loss": 0.0059,
+      "step": 70
+    },
+    {
+      "epoch": 0.5590551181102362,
+      "grad_norm": 0.5487034916877747,
+      "learning_rate": 1.8695321061768886e-05,
+      "loss": 0.0097,
+      "step": 71
+    },
+    {
+      "epoch": 0.5669291338582677,
+      "grad_norm": 0.6670834422111511,
+      "learning_rate": 1.8653188828095754e-05,
+      "loss": 0.011,
+      "step": 72
+    },
+    {
+      "epoch": 0.5748031496062992,
+      "grad_norm": 0.1795203685760498,
+      "learning_rate": 1.8610436117673557e-05,
+      "loss": 0.0067,
+      "step": 73
+    },
+    {
+      "epoch": 0.5826771653543307,
+      "grad_norm": 1.768436074256897,
+      "learning_rate": 1.8567065996084628e-05,
+      "loss": 0.0096,
+      "step": 74
+    },
+    {
+      "epoch": 0.5905511811023622,
+      "grad_norm": 0.26233312487602234,
+      "learning_rate": 1.8523081573182754e-05,
+      "loss": 0.0124,
+      "step": 75
+    },
+    {
+      "epoch": 0.5984251968503937,
+      "grad_norm": 0.3775719404220581,
+      "learning_rate": 1.847848600287019e-05,
+      "loss": 0.0052,
+      "step": 76
+    },
+    {
+      "epoch": 0.6062992125984252,
+      "grad_norm": 1.0016565322875977,
+      "learning_rate": 1.8433282482871497e-05,
+      "loss": 0.0058,
+      "step": 77
+    },
+    {
+      "epoch": 0.6141732283464567,
+      "grad_norm": 0.20153792202472687,
+      "learning_rate": 1.8387474254504265e-05,
+      "loss": 0.0056,
+      "step": 78
+    },
+    {
+      "epoch": 0.6220472440944882,
+      "grad_norm": 0.5119822025299072,
+      "learning_rate": 1.8341064602446686e-05,
+      "loss": 0.0079,
+      "step": 79
+    },
+    {
+      "epoch": 0.6299212598425197,
+      "grad_norm": 1.5781004428863525,
+      "learning_rate": 1.829405685450202e-05,
+      "loss": 0.008,
+      "step": 80
+    },
+    {
+      "epoch": 0.6377952755905512,
+      "grad_norm": 0.23826757073402405,
+      "learning_rate": 1.824645438135999e-05,
+      "loss": 0.0041,
+      "step": 81
+    },
+    {
+      "epoch": 0.6456692913385826,
+      "grad_norm": 0.6386727690696716,
+      "learning_rate": 1.8198260596355077e-05,
+      "loss": 0.0188,
+      "step": 82
+    },
+    {
+      "epoch": 0.6535433070866141,
+      "grad_norm": 0.9503199458122253,
+      "learning_rate": 1.814947895522176e-05,
+      "loss": 0.008,
+      "step": 83
+    },
+    {
+      "epoch": 0.6614173228346457,
+      "grad_norm": 0.2040701061487198,
+      "learning_rate": 1.8100112955846746e-05,
+      "loss": 0.0038,
+      "step": 84
+    },
+    {
+      "epoch": 0.6692913385826772,
+      "grad_norm": 0.3660199046134949,
+      "learning_rate": 1.805016613801813e-05,
+      "loss": 0.0148,
+      "step": 85
+    },
+    {
+      "epoch": 0.6771653543307087,
+      "grad_norm": 1.0502821207046509,
+      "learning_rate": 1.7999642083171576e-05,
+      "loss": 0.0098,
+      "step": 86
+    },
+    {
+      "epoch": 0.6771653543307087,
+      "eval_loss": 0.3526817262172699,
+      "eval_runtime": 6.6167,
+      "eval_samples_per_second": 24.635,
+      "eval_steps_per_second": 3.174,
+      "step": 86
+    },
+    {
+      "epoch": 0.6850393700787402,
+      "grad_norm": 0.13735969364643097,
+      "learning_rate": 1.7948544414133534e-05,
+      "loss": 0.0022,
+      "step": 87
+    },
+    {
+      "epoch": 0.6929133858267716,
+      "grad_norm": 0.6425012946128845,
+      "learning_rate": 1.7896876794861443e-05,
+      "loss": 0.0086,
+      "step": 88
+    },
+    {
+      "epoch": 0.7007874015748031,
+      "grad_norm": 0.7540380954742432,
+      "learning_rate": 1.7844642930181008e-05,
+      "loss": 0.0062,
+      "step": 89
+    },
+    {
+      "epoch": 0.7086614173228346,
+      "grad_norm": 0.6727365255355835,
+      "learning_rate": 1.779184656552056e-05,
+      "loss": 0.0027,
+      "step": 90
+    },
+    {
+      "epoch": 0.7165354330708661,
+      "grad_norm": 0.14059337973594666,
+      "learning_rate": 1.773849148664247e-05,
+      "loss": 0.0056,
+      "step": 91
+    },
+    {
+      "epoch": 0.7244094488188977,
+      "grad_norm": 0.33292093873023987,
+      "learning_rate": 1.7684581519371714e-05,
+      "loss": 0.0047,
+      "step": 92
+    },
+    {
+      "epoch": 0.7322834645669292,
+      "grad_norm": 0.3809877932071686,
+      "learning_rate": 1.7630120529321518e-05,
+      "loss": 0.0139,
+      "step": 93
+    },
+    {
+      "epoch": 0.7401574803149606,
+      "grad_norm": 1.729589819908142,
+      "learning_rate": 1.7575112421616203e-05,
+      "loss": 0.0128,
+      "step": 94
+    },
+    {
+      "epoch": 0.7480314960629921,
+      "grad_norm": 0.18192608654499054,
+      "learning_rate": 1.751956114061113e-05,
+      "loss": 0.0025,
+      "step": 95
+    },
+    {
+      "epoch": 0.7559055118110236,
+      "grad_norm": 1.0333118438720703,
+      "learning_rate": 1.7463470669609907e-05,
+      "loss": 0.006,
+      "step": 96
+    },
+    {
+      "epoch": 0.7637795275590551,
+      "grad_norm": 0.7247685194015503,
+      "learning_rate": 1.7406845030578747e-05,
+      "loss": 0.0073,
+      "step": 97
+    },
+    {
+      "epoch": 0.7716535433070866,
+      "grad_norm": 0.06979379802942276,
+      "learning_rate": 1.734968828385808e-05,
+      "loss": 0.0005,
+      "step": 98
+    },
+    {
+      "epoch": 0.7795275590551181,
+      "grad_norm": 0.5137119293212891,
+      "learning_rate": 1.729200452787139e-05,
+      "loss": 0.0082,
+      "step": 99
+    },
+    {
+      "epoch": 0.7874015748031497,
+      "grad_norm": 0.4704137146472931,
+      "learning_rate": 1.7233797898831376e-05,
+      "loss": 0.005,
+      "step": 100
+    },
+    {
+      "epoch": 0.7952755905511811,
+      "grad_norm": 0.28564465045928955,
+      "learning_rate": 1.717507257044331e-05,
+      "loss": 0.0052,
+      "step": 101
+    },
+    {
+      "epoch": 0.8031496062992126,
+      "grad_norm": 0.17685537040233612,
+      "learning_rate": 1.711583275360582e-05,
+      "loss": 0.0024,
+      "step": 102
+    },
+    {
+      "epoch": 0.8110236220472441,
+      "grad_norm": 0.45714935660362244,
+      "learning_rate": 1.7056082696108896e-05,
+      "loss": 0.0072,
+      "step": 103
+    },
+    {
+      "epoch": 0.8188976377952756,
+      "grad_norm": 0.4373086988925934,
+      "learning_rate": 1.699582668232934e-05,
+      "loss": 0.0051,
+      "step": 104
+    },
+    {
+      "epoch": 0.8267716535433071,
+      "grad_norm": 0.8478983640670776,
+      "learning_rate": 1.6935069032923525e-05,
+      "loss": 0.022,
+      "step": 105
+    },
+    {
+      "epoch": 0.8346456692913385,
+      "grad_norm": 0.16181086003780365,
+      "learning_rate": 1.6873814104517617e-05,
+      "loss": 0.0058,
+      "step": 106
+    },
+    {
+      "epoch": 0.84251968503937,
+      "grad_norm": 0.09503592550754547,
+      "learning_rate": 1.6812066289395157e-05,
+      "loss": 0.0009,
+      "step": 107
+    },
+    {
+      "epoch": 0.8503937007874016,
+      "grad_norm": 0.7462632060050964,
+      "learning_rate": 1.6749830015182106e-05,
+      "loss": 0.0044,
+      "step": 108
+    },
+    {
+      "epoch": 0.8582677165354331,
+      "grad_norm": 0.07221701741218567,
+      "learning_rate": 1.6687109744529394e-05,
+      "loss": 0.0015,
+      "step": 109
+    },
+    {
+      "epoch": 0.8661417322834646,
+      "grad_norm": 0.08999036252498627,
+      "learning_rate": 1.6623909974792888e-05,
+      "loss": 0.0023,
+      "step": 110
+    },
+    {
+      "epoch": 0.8740157480314961,
+      "grad_norm": 0.42536938190460205,
+      "learning_rate": 1.656023523771095e-05,
+      "loss": 0.005,
+      "step": 111
+    },
+    {
+      "epoch": 0.8818897637795275,
+      "grad_norm": 0.7885191440582275,
+      "learning_rate": 1.6496090099079452e-05,
+      "loss": 0.0103,
+      "step": 112
+    },
+    {
+      "epoch": 0.889763779527559,
+      "grad_norm": 0.16610018908977509,
+      "learning_rate": 1.64314791584244e-05,
+      "loss": 0.006,
+      "step": 113
+    },
+    {
+      "epoch": 0.8976377952755905,
+      "grad_norm": 0.32151034474372864,
+      "learning_rate": 1.6366407048672135e-05,
+      "loss": 0.0086,
+      "step": 114
+    },
+    {
+      "epoch": 0.905511811023622,
+      "grad_norm": 0.557732343673706,
+      "learning_rate": 1.6300878435817115e-05,
+      "loss": 0.0064,
+      "step": 115
+    },
+    {
+      "epoch": 0.9133858267716536,
+      "grad_norm": 0.2238176167011261,
+      "learning_rate": 1.6234898018587336e-05,
+      "loss": 0.0065,
+      "step": 116
+    },
+    {
+      "epoch": 0.9212598425196851,
+      "grad_norm": 0.2980042099952698,
+      "learning_rate": 1.616847052810744e-05,
+      "loss": 0.0095,
+      "step": 117
+    },
+    {
+      "epoch": 0.9291338582677166,
+      "grad_norm": 0.1529705822467804,
+      "learning_rate": 1.6101600727559423e-05,
+      "loss": 0.0062,
+      "step": 118
+    },
+    {
+      "epoch": 0.937007874015748,
+      "grad_norm": 0.017149658873677254,
+      "learning_rate": 1.603429341184114e-05,
+      "loss": 0.0002,
+      "step": 119
+    },
+    {
+      "epoch": 0.9448818897637795,
+      "grad_norm": 0.4514746367931366,
+      "learning_rate": 1.596655340722244e-05,
+      "loss": 0.0067,
+      "step": 120
+    },
+    {
+      "epoch": 0.952755905511811,
+      "grad_norm": 0.11766134947538376,
+      "learning_rate": 1.5898385570999146e-05,
+      "loss": 0.0053,
+      "step": 121
+    },
+    {
+      "epoch": 0.9606299212598425,
+      "grad_norm": 0.4089784026145935,
+      "learning_rate": 1.5829794791144723e-05,
+      "loss": 0.0085,
+      "step": 122
+    },
+    {
+      "epoch": 0.968503937007874,
+      "grad_norm": 0.1353057473897934,
+      "learning_rate": 1.57607859859598e-05,
+      "loss": 0.0013,
+      "step": 123
+    },
+    {
+      "epoch": 0.9763779527559056,
+      "grad_norm": 0.6548481583595276,
+      "learning_rate": 1.5691364103719515e-05,
+      "loss": 0.0117,
+      "step": 124
+    },
+    {
+      "epoch": 0.984251968503937,
+      "grad_norm": 0.1571267992258072,
+      "learning_rate": 1.5621534122318682e-05,
+      "loss": 0.0049,
+      "step": 125
+    },
+    {
+      "epoch": 0.9921259842519685,
+      "grad_norm": 1.2177189588546753,
+      "learning_rate": 1.5551301048914863e-05,
+      "loss": 0.0161,
+      "step": 126
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.414489209651947,
+      "learning_rate": 1.5480669919569313e-05,
+      "loss": 0.0181,
+      "step": 127
+    },
+    {
+      "epoch": 1.0078740157480315,
+      "grad_norm": 0.10985995829105377,
+      "learning_rate": 1.54096457988859e-05,
+      "loss": 0.0049,
+      "step": 128
+    },
+    {
+      "epoch": 1.015748031496063,
+      "grad_norm": 0.12780147790908813,
+      "learning_rate": 1.533823377964791e-05,
+      "loss": 0.0026,
+      "step": 129
+    },
+    {
+      "epoch": 1.015748031496063,
+      "eval_loss": 0.33064374327659607,
+      "eval_runtime": 6.9286,
+      "eval_samples_per_second": 23.526,
+      "eval_steps_per_second": 3.031,
+      "step": 129
+    },
+    {
+      "epoch": 1.0236220472440944,
+      "grad_norm": 0.5142458081245422,
+      "learning_rate": 1.52664389824529e-05,
+      "loss": 0.0082,
+      "step": 130
+    },
+    {
+      "epoch": 1.031496062992126,
+      "grad_norm": 0.15617145597934723,
+      "learning_rate": 1.5194266555345505e-05,
+      "loss": 0.0016,
+      "step": 131
+    },
+    {
+      "epoch": 1.0393700787401574,
+      "grad_norm": 0.5782387852668762,
+      "learning_rate": 1.5121721673448319e-05,
+      "loss": 0.0117,
+      "step": 132
+    },
+    {
+      "epoch": 1.047244094488189,
+      "grad_norm": 0.08414836972951889,
+      "learning_rate": 1.5048809538590789e-05,
+      "loss": 0.0021,
+      "step": 133
+    },
+    {
+      "epoch": 1.0551181102362204,
+      "grad_norm": 0.28253939747810364,
+      "learning_rate": 1.4975535378936228e-05,
+      "loss": 0.0055,
+      "step": 134
+    },
+    {
+      "epoch": 1.0629921259842519,
+      "grad_norm": 0.47917842864990234,
+      "learning_rate": 1.490190444860694e-05,
+      "loss": 0.0046,
+      "step": 135
+    },
+    {
+      "epoch": 1.0708661417322836,
+      "grad_norm": 0.1895662248134613,
+      "learning_rate": 1.482792202730745e-05,
+      "loss": 0.006,
+      "step": 136
+    },
+    {
+      "epoch": 1.078740157480315,
+      "grad_norm": 0.13722768425941467,
+      "learning_rate": 1.475359341994595e-05,
+      "loss": 0.0031,
+      "step": 137
+    },
+    {
+      "epoch": 1.0866141732283465,
+      "grad_norm": 0.10731153190135956,
+      "learning_rate": 1.4678923956253894e-05,
+      "loss": 0.0005,
+      "step": 138
+    },
+    {
+      "epoch": 1.094488188976378,
+      "grad_norm": 0.12261265516281128,
+      "learning_rate": 1.460391899040383e-05,
+      "loss": 0.0031,
+      "step": 139
+    },
+    {
+      "epoch": 1.1023622047244095,
+      "grad_norm": 0.0038245893083512783,
+      "learning_rate": 1.4528583900625481e-05,
+      "loss": 0.0,
+      "step": 140
+    },
+    {
+      "epoch": 1.110236220472441,
+      "grad_norm": 0.28762558102607727,
+      "learning_rate": 1.4452924088820101e-05,
+      "loss": 0.004,
+      "step": 141
+    },
+    {
+      "epoch": 1.1181102362204725,
+      "grad_norm": 0.17267552018165588,
+      "learning_rate": 1.4376944980173138e-05,
+      "loss": 0.0002,
+      "step": 142
+    },
+    {
+      "epoch": 1.125984251968504,
+      "grad_norm": 0.12727122008800507,
+      "learning_rate": 1.4300652022765207e-05,
+      "loss": 0.0029,
+      "step": 143
+    },
+    {
+      "epoch": 1.1338582677165354,
+      "grad_norm": 0.25049135088920593,
+      "learning_rate": 1.4224050687181442e-05,
+      "loss": 0.0108,
+      "step": 144
+    },
+    {
+      "epoch": 1.141732283464567,
+      "grad_norm": 0.16092728078365326,
+      "learning_rate": 1.4147146466119235e-05,
+      "loss": 0.0024,
+      "step": 145
+    },
+    {
+      "epoch": 1.1496062992125984,
+      "grad_norm": 0.13642658293247223,
+      "learning_rate": 1.406994487399437e-05,
+      "loss": 0.0037,
+      "step": 146
+    },
+    {
+      "epoch": 1.1574803149606299,
+      "grad_norm": 0.9029403328895569,
+      "learning_rate": 1.3992451446545624e-05,
+      "loss": 0.0034,
+      "step": 147
+    },
+    {
+      "epoch": 1.1653543307086613,
+      "grad_norm": 0.19518424570560455,
+      "learning_rate": 1.3914671740437811e-05,
+      "loss": 0.0057,
+      "step": 148
+    },
+    {
+      "epoch": 1.1732283464566928,
+      "grad_norm": 0.12140502035617828,
+      "learning_rate": 1.3836611332863356e-05,
+      "loss": 0.0041,
+      "step": 149
+    },
+    {
+      "epoch": 1.1811023622047245,
+      "grad_norm": 0.5148038864135742,
+      "learning_rate": 1.3758275821142382e-05,
+      "loss": 0.0026,
+      "step": 150
+    },
+    {
+      "epoch": 1.188976377952756,
+      "grad_norm": 1.828904390335083,
+      "learning_rate": 1.3679670822321347e-05,
+      "loss": 0.0024,
+      "step": 151
+    },
+    {
+      "epoch": 1.1968503937007875,
+      "grad_norm": 0.3571717143058777,
+      "learning_rate": 1.3600801972770272e-05,
+      "loss": 0.0106,
+      "step": 152
+    },
+    {
+      "epoch": 1.204724409448819,
+      "grad_norm": 0.051027003675699234,
+      "learning_rate": 1.3521674927778594e-05,
+      "loss": 0.0003,
+      "step": 153
+    },
+    {
+      "epoch": 1.2125984251968505,
+      "grad_norm": 0.6490982174873352,
+      "learning_rate": 1.3442295361149651e-05,
+      "loss": 0.0035,
+      "step": 154
+    },
+    {
+      "epoch": 1.220472440944882,
+      "grad_norm": 0.08408445864915848,
+      "learning_rate": 1.336266896479384e-05,
+      "loss": 0.0027,
+      "step": 155
+    },
+    {
+      "epoch": 1.2283464566929134,
+      "grad_norm": 0.09666562080383301,
+      "learning_rate": 1.328280144832047e-05,
+      "loss": 0.0019,
+      "step": 156
+    },
+    {
+      "epoch": 1.236220472440945,
+      "grad_norm": 0.03880690038204193,
+      "learning_rate": 1.3202698538628376e-05,
+      "loss": 0.0003,
+      "step": 157
+    },
+    {
+      "epoch": 1.2440944881889764,
+      "grad_norm": 0.11940775066614151,
+      "learning_rate": 1.3122365979495259e-05,
+      "loss": 0.0024,
+      "step": 158
+    },
+    {
+      "epoch": 1.2519685039370079,
+      "grad_norm": 0.1442880481481552,
+      "learning_rate": 1.3041809531165819e-05,
+      "loss": 0.0015,
+      "step": 159
+    },
+    {
+      "epoch": 1.2598425196850394,
+      "grad_norm": 0.1961939036846161,
+      "learning_rate": 1.2961034969938732e-05,
+      "loss": 0.0056,
+      "step": 160
+    },
+    {
+      "epoch": 1.2677165354330708,
+      "grad_norm": 0.26947638392448425,
+      "learning_rate": 1.288004808775246e-05,
+      "loss": 0.0028,
+      "step": 161
+    },
+    {
+      "epoch": 1.2755905511811023,
+      "grad_norm": 0.5154056549072266,
+      "learning_rate": 1.2798854691769927e-05,
+      "loss": 0.0037,
+      "step": 162
+    },
+    {
+      "epoch": 1.2834645669291338,
+      "grad_norm": 0.4292369782924652,
+      "learning_rate": 1.2717460603962132e-05,
+      "loss": 0.0029,
+      "step": 163
+    },
+    {
+      "epoch": 1.2913385826771653,
+      "grad_norm": 0.19139212369918823,
+      "learning_rate": 1.2635871660690677e-05,
+      "loss": 0.0061,
+      "step": 164
+    },
+    {
+      "epoch": 1.2992125984251968,
+      "grad_norm": 0.19960306584835052,
+      "learning_rate": 1.2554093712289267e-05,
+      "loss": 0.005,
+      "step": 165
+    },
+    {
+      "epoch": 1.3070866141732282,
+      "grad_norm": 0.4523830711841583,
+      "learning_rate": 1.2472132622644222e-05,
+      "loss": 0.0065,
+      "step": 166
+    },
+    {
+      "epoch": 1.3149606299212597,
+      "grad_norm": 0.49343299865722656,
+      "learning_rate": 1.2389994268773995e-05,
+      "loss": 0.0061,
+      "step": 167
+    },
+    {
+      "epoch": 1.3228346456692912,
+      "grad_norm": 0.01938088797032833,
+      "learning_rate": 1.2307684540407775e-05,
+      "loss": 0.0001,
+      "step": 168
+    },
+    {
+      "epoch": 1.330708661417323,
+      "grad_norm": 0.3082112669944763,
+      "learning_rate": 1.2225209339563144e-05,
+      "loss": 0.0053,
+      "step": 169
+    },
+    {
+      "epoch": 1.3385826771653544,
+      "grad_norm": 0.01982509344816208,
+      "learning_rate": 1.2142574580122903e-05,
+      "loss": 0.0001,
+      "step": 170
+    },
+    {
+      "epoch": 1.3464566929133859,
+      "grad_norm": 0.12388588488101959,
+      "learning_rate": 1.2059786187410984e-05,
+      "loss": 0.0049,
+      "step": 171
+    },
+    {
+      "epoch": 1.3543307086614174,
+      "grad_norm": 0.43759095668792725,
+      "learning_rate": 1.1976850097767598e-05,
+      "loss": 0.0128,
+      "step": 172
+    },
+    {
+      "epoch": 1.3543307086614174,
+      "eval_loss": 0.3166251480579376,
+      "eval_runtime": 6.9515,
+      "eval_samples_per_second": 23.448,
+      "eval_steps_per_second": 3.021,
+      "step": 172
+    },
+    {
+      "epoch": 1.3622047244094488,
+      "grad_norm": 0.46561670303344727,
+      "learning_rate": 1.1893772258123554e-05,
+      "loss": 0.008,
+      "step": 173
+    },
+    {
+      "epoch": 1.3700787401574803,
+      "grad_norm": 0.16612188518047333,
+      "learning_rate": 1.1810558625573856e-05,
+      "loss": 0.0024,
+      "step": 174
+    },
+    {
+      "epoch": 1.3779527559055118,
+      "grad_norm": 0.13628093898296356,
+      "learning_rate": 1.1727215166950519e-05,
+      "loss": 0.0045,
+      "step": 175
+    },
+    {
+      "epoch": 1.3858267716535433,
+      "grad_norm": 0.565229058265686,
+      "learning_rate": 1.1643747858394743e-05,
+      "loss": 0.0103,
+      "step": 176
+    },
+    {
+      "epoch": 1.3937007874015748,
+      "grad_norm": 0.14550763368606567,
+      "learning_rate": 1.156016268492839e-05,
+      "loss": 0.0028,
+      "step": 177
+    },
+    {
+      "epoch": 1.4015748031496063,
+      "grad_norm": 0.12460129708051682,
+      "learning_rate": 1.1476465640024814e-05,
+      "loss": 0.0031,
+      "step": 178
+    },
+    {
+      "epoch": 1.4094488188976377,
+      "grad_norm": 0.19089221954345703,
+      "learning_rate": 1.1392662725179114e-05,
+      "loss": 0.0035,
+      "step": 179
+    },
+    {
+      "epoch": 1.4173228346456692,
+      "grad_norm": 0.6106573343276978,
+      "learning_rate": 1.1308759949477786e-05,
+      "loss": 0.0088,
+      "step": 180
+    },
+    {
+      "epoch": 1.425196850393701,
+      "grad_norm": 0.20053207874298096,
+      "learning_rate": 1.1224763329167859e-05,
+      "loss": 0.0033,
+      "step": 181
+    },
+    {
+      "epoch": 1.4330708661417324,
+      "grad_norm": 0.1984691321849823,
+      "learning_rate": 1.1140678887225468e-05,
+      "loss": 0.0051,
+      "step": 182
+    },
+    {
+      "epoch": 1.4409448818897639,
+      "grad_norm": 0.19264858961105347,
+      "learning_rate": 1.1056512652924014e-05,
+      "loss": 0.0046,
+      "step": 183
+    },
+    {
+      "epoch": 1.4488188976377954,
+      "grad_norm": 0.10979076474905014,
+      "learning_rate": 1.0972270661401812e-05,
+      "loss": 0.0031,
+      "step": 184
+    },
+    {
+      "epoch": 1.4566929133858268,
+      "grad_norm": 0.1744084656238556,
+      "learning_rate": 1.0887958953229349e-05,
+      "loss": 0.0024,
+      "step": 185
+    },
+    {
+      "epoch": 1.4645669291338583,
+      "grad_norm": 0.20646224915981293,
+      "learning_rate": 1.0803583573976137e-05,
+      "loss": 0.008,
+      "step": 186
+    },
+    {
+      "epoch": 1.4724409448818898,
+      "grad_norm": 0.14391584694385529,
+      "learning_rate": 1.0719150573777226e-05,
+      "loss": 0.004,
+      "step": 187
+    },
+    {
+      "epoch": 1.4803149606299213,
+      "grad_norm": 0.36887863278388977,
+      "learning_rate": 1.0634666006899375e-05,
+      "loss": 0.0074,
+      "step": 188
+    },
+    {
+      "epoch": 1.4881889763779528,
+      "grad_norm": 0.21352627873420715,
+      "learning_rate": 1.055013593130693e-05,
+      "loss": 0.0082,
+      "step": 189
+    },
+    {
+      "epoch": 1.4960629921259843,
+      "grad_norm": 0.22443020343780518,
+      "learning_rate": 1.046556640822744e-05,
+      "loss": 0.0087,
+      "step": 190
+    },
+    {
+      "epoch": 1.5039370078740157,
+      "grad_norm": 0.4243764281272888,
+      "learning_rate": 1.0380963501717034e-05,
+      "loss": 0.0068,
+      "step": 191
+    },
+    {
+      "epoch": 1.5118110236220472,
+      "grad_norm": 0.17558562755584717,
+      "learning_rate": 1.0296333278225599e-05,
+      "loss": 0.0054,
+      "step": 192
+    },
+    {
+      "epoch": 1.5196850393700787,
+      "grad_norm": 0.14842620491981506,
+      "learning_rate": 1.0211681806161787e-05,
+      "loss": 0.0031,
+      "step": 193
+    },
+    {
+      "epoch": 1.5275590551181102,
+      "grad_norm": 0.09316081553697586,
+      "learning_rate": 1.0127015155457875e-05,
+      "loss": 0.0013,
+      "step": 194
+    },
+    {
+      "epoch": 1.5354330708661417,
+      "grad_norm": 0.19795025885105133,
+      "learning_rate": 1.0042339397134528e-05,
+      "loss": 0.0051,
+      "step": 195
+    },
+    {
+      "epoch": 1.5433070866141732,
+      "grad_norm": 0.21606990694999695,
+      "learning_rate": 9.957660602865477e-06,
+      "loss": 0.0041,
+      "step": 196
+    },
+    {
+      "epoch": 1.5511811023622046,
+      "grad_norm": 0.18036173284053802,
+      "learning_rate": 9.872984844542128e-06,
+      "loss": 0.0037,
+      "step": 197
+    },
+    {
+      "epoch": 1.5590551181102361,
+      "grad_norm": 0.18953870236873627,
+      "learning_rate": 9.788318193838218e-06,
+      "loss": 0.0041,
+      "step": 198
+    },
+    {
+      "epoch": 1.5669291338582676,
+      "grad_norm": 0.12346503138542175,
+      "learning_rate": 9.703666721774403e-06,
+      "loss": 0.0035,
+      "step": 199
+    },
+    {
+      "epoch": 1.574803149606299,
+      "grad_norm": 0.4576225280761719,
+      "learning_rate": 9.619036498282968e-06,
+      "loss": 0.0041,
+      "step": 200
+    },
+    {
+      "epoch": 1.5826771653543306,
+      "grad_norm": 0.10333681106567383,
+      "learning_rate": 9.534433591772562e-06,
+      "loss": 0.0011,
+      "step": 201
+    },
+    {
+      "epoch": 1.590551181102362,
+      "grad_norm": 0.19167865812778473,
+      "learning_rate": 9.449864068693072e-06,
+      "loss": 0.0062,
+      "step": 202
+    },
+    {
+      "epoch": 1.5984251968503937,
+      "grad_norm": 0.2258184254169464,
+      "learning_rate": 9.365333993100628e-06,
+      "loss": 0.003,
+      "step": 203
+    },
+    {
+      "epoch": 1.6062992125984252,
+      "grad_norm": 0.07945302873849869,
+      "learning_rate": 9.280849426222778e-06,
+      "loss": 0.0008,
+      "step": 204
+    },
+    {
+      "epoch": 1.6141732283464567,
+      "grad_norm": 0.17767398059368134,
+      "learning_rate": 9.196416426023868e-06,
+      "loss": 0.0053,
+      "step": 205
+    },
+    {
+      "epoch": 1.6220472440944882,
+      "grad_norm": 0.12704500555992126,
+      "learning_rate": 9.112041046770653e-06,
+      "loss": 0.0023,
+      "step": 206
+    },
+    {
+      "epoch": 1.6299212598425197,
+      "grad_norm": 0.4054742753505707,
+      "learning_rate": 9.027729338598188e-06,
+      "loss": 0.0045,
+      "step": 207
+    },
+    {
+      "epoch": 1.6377952755905512,
+      "grad_norm": 0.4463757574558258,
+      "learning_rate": 8.943487347075988e-06,
+      "loss": 0.007,
+      "step": 208
+    },
+    {
+      "epoch": 1.6456692913385826,
+      "grad_norm": 0.6517045497894287,
+      "learning_rate": 8.859321112774535e-06,
+      "loss": 0.0052,
+      "step": 209
+    },
+    {
+      "epoch": 1.6535433070866141,
+      "grad_norm": 0.1542089730501175,
+      "learning_rate": 8.775236670832146e-06,
+      "loss": 0.0047,
+      "step": 210
+    },
+    {
+      "epoch": 1.6614173228346458,
+      "grad_norm": 0.14716440439224243,
+      "learning_rate": 8.691240050522215e-06,
+      "loss": 0.0049,
+      "step": 211
+    },
+    {
+      "epoch": 1.6692913385826773,
+      "grad_norm": 0.2997347116470337,
+      "learning_rate": 8.607337274820888e-06,
+      "loss": 0.0076,
+      "step": 212
+    },
+    {
+      "epoch": 1.6771653543307088,
+      "grad_norm": 0.22548256814479828,
+      "learning_rate": 8.52353435997519e-06,
+      "loss": 0.0063,
+      "step": 213
+    },
+    {
+      "epoch": 1.6850393700787403,
+      "grad_norm": 0.7220733165740967,
+      "learning_rate": 8.439837315071612e-06,
+      "loss": 0.0089,
+      "step": 214
+    },
+    {
+      "epoch": 1.6929133858267718,
+      "grad_norm": 0.5101618766784668,
+      "learning_rate": 8.35625214160526e-06,
+      "loss": 0.0042,
+      "step": 215
+    },
+    {
+      "epoch": 1.6929133858267718,
+      "eval_loss": 0.3484288156032562,
+      "eval_runtime": 6.4482,
+      "eval_samples_per_second": 25.278,
+      "eval_steps_per_second": 3.257,
+      "step": 215
+    },
+    {
+      "epoch": 1.7007874015748032,
+      "grad_norm": 0.1698393076658249,
+      "learning_rate": 8.272784833049485e-06,
+      "loss": 0.0028,
+      "step": 216
+    },
+    {
+      "epoch": 1.7086614173228347,
+      "grad_norm": 0.5772718191146851,
+      "learning_rate": 8.18944137442615e-06,
+      "loss": 0.0082,
+      "step": 217
+    },
+    {
+      "epoch": 1.7165354330708662,
+      "grad_norm": 0.09606469422578812,
+      "learning_rate": 8.106227741876447e-06,
+      "loss": 0.0011,
+      "step": 218
+    },
+    {
+      "epoch": 1.7244094488188977,
+      "grad_norm": 0.14510361850261688,
+      "learning_rate": 8.023149902232404e-06,
+      "loss": 0.0015,
+      "step": 219
+    },
+    {
+      "epoch": 1.7322834645669292,
+      "grad_norm": 0.055804118514060974,
+      "learning_rate": 7.940213812589018e-06,
+      "loss": 0.0008,
+      "step": 220
+    },
+    {
+      "epoch": 1.7401574803149606,
+      "grad_norm": 0.13318321108818054,
+      "learning_rate": 7.857425419877097e-06,
+      "loss": 0.005,
+      "step": 221
+    },
+    {
+      "epoch": 1.7480314960629921,
+      "grad_norm": 0.23600782454013824,
+      "learning_rate": 7.774790660436857e-06,
+      "loss": 0.0063,
+      "step": 222
+    },
+    {
+      "epoch": 1.7559055118110236,
+      "grad_norm": 0.8483791351318359,
+      "learning_rate": 7.69231545959223e-06,
+      "loss": 0.0027,
+      "step": 223
+    },
+    {
+      "epoch": 1.763779527559055,
+      "grad_norm": 0.16536197066307068,
+      "learning_rate": 7.610005731226009e-06,
+      "loss": 0.0039,
+      "step": 224
+    },
+    {
+      "epoch": 1.7716535433070866,
+      "grad_norm": 0.14446765184402466,
+      "learning_rate": 7.52786737735578e-06,
+      "loss": 0.0036,
+      "step": 225
+    },
+    {
+      "epoch": 1.779527559055118,
+      "grad_norm": 0.8880365490913391,
+      "learning_rate": 7.445906287710733e-06,
+      "loss": 0.0061,
+      "step": 226
+    },
+    {
+      "epoch": 1.7874015748031495,
+      "grad_norm": 0.151743084192276,
+      "learning_rate": 7.364128339309326e-06,
+      "loss": 0.0028,
+      "step": 227
+    },
+    {
+      "epoch": 1.795275590551181,
+      "grad_norm": 0.1224551647901535,
+      "learning_rate": 7.282539396037868e-06,
+      "loss": 0.002,
+      "step": 228
+    },
+    {
+      "epoch": 1.8031496062992125,
+      "grad_norm": 0.4868486225605011,
+      "learning_rate": 7.201145308230075e-06,
+      "loss": 0.0031,
+      "step": 229
+    },
+    {
+      "epoch": 1.811023622047244,
+      "grad_norm": 0.2875569462776184,
+      "learning_rate": 7.119951912247545e-06,
+      "loss": 0.0082,
+      "step": 230
+    },
+    {
+      "epoch": 1.8188976377952755,
+      "grad_norm": 0.43524420261383057,
+      "learning_rate": 7.038965030061273e-06,
+      "loss": 0.0075,
+      "step": 231
+    },
+    {
+      "epoch": 1.826771653543307,
+      "grad_norm": 0.39634883403778076,
+      "learning_rate": 6.9581904688341854e-06,
+      "loss": 0.0032,
+      "step": 232
+    },
+    {
+      "epoch": 1.8346456692913384,
+      "grad_norm": 0.9809433817863464,
+      "learning_rate": 6.8776340205047446e-06,
+      "loss": 0.0085,
+      "step": 233
+    },
+    {
+      "epoch": 1.84251968503937,
+      "grad_norm": 0.20062875747680664,
+      "learning_rate": 6.797301461371626e-06,
+      "loss": 0.0043,
+      "step": 234
+    },
+    {
+      "epoch": 1.8503937007874016,
+      "grad_norm": 0.148948073387146,
+      "learning_rate": 6.7171985516795315e-06,
+      "loss": 0.0036,
+      "step": 235
+    },
+    {
+      "epoch": 1.858267716535433,
+      "grad_norm": 0.15658679604530334,
+      "learning_rate": 6.637331035206166e-06,
+      "loss": 0.0046,
+      "step": 236
+    },
+    {
+      "epoch": 1.8661417322834646,
+      "grad_norm": 0.22365815937519073,
+      "learning_rate": 6.557704638850352e-06,
+      "loss": 0.0081,
+      "step": 237
+    },
+    {
+      "epoch": 1.874015748031496,
+      "grad_norm": 0.10596666485071182,
+      "learning_rate": 6.4783250722214066e-06,
+      "loss": 0.0032,
+      "step": 238
+    },
+    {
+      "epoch": 1.8818897637795275,
+      "grad_norm": 0.2130754142999649,
+      "learning_rate": 6.399198027229732e-06,
+      "loss": 0.0056,
+      "step": 239
+    },
+    {
+      "epoch": 1.889763779527559,
+      "grad_norm": 0.05641167238354683,
+      "learning_rate": 6.320329177678656e-06,
+      "loss": 0.0008,
+      "step": 240
+    },
+    {
+      "epoch": 1.8976377952755905,
+      "grad_norm": 0.10349344462156296,
+      "learning_rate": 6.241724178857621e-06,
+      "loss": 0.0026,
+      "step": 241
+    },
+    {
+      "epoch": 1.905511811023622,
+      "grad_norm": 0.08451675623655319,
+      "learning_rate": 6.163388667136646e-06,
+      "loss": 0.0016,
+      "step": 242
+    },
+    {
+      "epoch": 1.9133858267716537,
+      "grad_norm": 0.13671623170375824,
+      "learning_rate": 6.085328259562195e-06,
+      "loss": 0.0034,
+      "step": 243
+    },
+    {
+      "epoch": 1.9212598425196852,
+      "grad_norm": 0.5500523447990417,
+      "learning_rate": 6.007548553454379e-06,
+      "loss": 0.0028,
+      "step": 244
+    },
+    {
+      "epoch": 1.9291338582677167,
+      "grad_norm": 0.06702329218387604,
+      "learning_rate": 5.93005512600563e-06,
+      "loss": 0.0009,
+      "step": 245
+    },
+    {
+      "epoch": 1.9370078740157481,
+      "grad_norm": 0.15156973898410797,
+      "learning_rate": 5.852853533880768e-06,
+      "loss": 0.0064,
+      "step": 246
+    },
+    {
+      "epoch": 1.9448818897637796,
+      "grad_norm": 0.2970314621925354,
+      "learning_rate": 5.7759493128185584e-06,
+      "loss": 0.0077,
+      "step": 247
+    },
+    {
+      "epoch": 1.952755905511811,
+      "grad_norm": 0.06406261771917343,
+      "learning_rate": 5.699347977234799e-06,
+      "loss": 0.0006,
+      "step": 248
+    },
+    {
+      "epoch": 1.9606299212598426,
+      "grad_norm": 0.2910393178462982,
+      "learning_rate": 5.623055019826862e-06,
+      "loss": 0.0036,
+      "step": 249
+    },
+    {
+      "epoch": 1.968503937007874,
+      "grad_norm": 0.6454993486404419,
+      "learning_rate": 5.547075911179902e-06,
+      "loss": 0.0084,
+      "step": 250
+    },
+    {
+      "epoch": 1.9763779527559056,
+      "grad_norm": 0.09460143744945526,
+      "learning_rate": 5.471416099374525e-06,
+      "loss": 0.0021,
+      "step": 251
+    },
+    {
+      "epoch": 1.984251968503937,
+      "grad_norm": 0.2024363875389099,
+      "learning_rate": 5.3960810095961705e-06,
+      "loss": 0.0052,
+      "step": 252
+    },
+    {
+      "epoch": 1.9921259842519685,
+      "grad_norm": 0.09423142671585083,
+      "learning_rate": 5.321076043746108e-06,
+      "loss": 0.0018,
+      "step": 253
+    },
+    {
+      "epoch": 2.0,
+      "grad_norm": 0.1085880920290947,
+      "learning_rate": 5.246406580054051e-06,
+      "loss": 0.0039,
+      "step": 254
+    },
+    {
+      "epoch": 2.0078740157480315,
+      "grad_norm": 0.20550444722175598,
+      "learning_rate": 5.172077972692553e-06,
+      "loss": 0.0006,
+      "step": 255
+    },
+    {
+      "epoch": 2.015748031496063,
+      "grad_norm": 0.0635254830121994,
+      "learning_rate": 5.098095551393066e-06,
+      "loss": 0.0008,
+      "step": 256
+    },
+    {
+      "epoch": 2.0236220472440944,
+      "grad_norm": 0.12593789398670197,
+      "learning_rate": 5.024464621063773e-06,
+      "loss": 0.0016,
+      "step": 257
+    },
+    {
+      "epoch": 2.031496062992126,
+      "grad_norm": 0.08928010612726212,
+      "learning_rate": 4.951190461409214e-06,
+      "loss": 0.0019,
+      "step": 258
+    },
+    {
+      "epoch": 2.031496062992126,
+      "eval_loss": 0.2930968105792999,
+      "eval_runtime": 7.0864,
+      "eval_samples_per_second": 23.002,
+      "eval_steps_per_second": 2.963,
+      "step": 258
+    },
+    {
+      "epoch": 2.0393700787401574,
+      "grad_norm": 0.11555846035480499,
+      "learning_rate": 4.878278326551682e-06,
+      "loss": 0.0036,
+      "step": 259
+    },
+    {
+      "epoch": 2.047244094488189,
+      "grad_norm": 0.11923055350780487,
+      "learning_rate": 4.805733444654496e-06,
+      "loss": 0.0011,
+      "step": 260
+    },
+    {
+      "epoch": 2.0551181102362204,
+      "grad_norm": 0.5410908460617065,
+      "learning_rate": 4.733561017547104e-06,
+      "loss": 0.0065,
+      "step": 261
+    },
+    {
+      "epoch": 2.062992125984252,
+      "grad_norm": 0.43598446249961853,
+      "learning_rate": 4.661766220352098e-06,
+      "loss": 0.004,
+      "step": 262
+    },
+    {
+      "epoch": 2.0708661417322833,
+      "grad_norm": 0.08221737295389175,
+      "learning_rate": 4.590354201114103e-06,
+      "loss": 0.0018,
+      "step": 263
+    },
+    {
+      "epoch": 2.078740157480315,
+      "grad_norm": 0.07835202664136887,
+      "learning_rate": 4.519330080430687e-06,
+      "loss": 0.0011,
+      "step": 264
+    },
+    {
+      "epoch": 2.0866141732283463,
+      "grad_norm": 0.1391119360923767,
+      "learning_rate": 4.448698951085143e-06,
+      "loss": 0.0018,
+      "step": 265
+    },
+    {
+      "epoch": 2.094488188976378,
+      "grad_norm": 0.10286661982536316,
+      "learning_rate": 4.378465877681317e-06,
+      "loss": 0.0021,
+      "step": 266
+    },
+    {
+      "epoch": 2.1023622047244093,
+      "grad_norm": 0.16050903499126434,
+      "learning_rate": 4.3086358962804885e-06,
+      "loss": 0.004,
+      "step": 267
+    },
+    {
+      "epoch": 2.1102362204724407,
+      "grad_norm": 0.1615462303161621,
+      "learning_rate": 4.2392140140401996e-06,
+      "loss": 0.0049,
+      "step": 268
+    },
+    {
+      "epoch": 2.1181102362204722,
+      "grad_norm": 0.12022113800048828,
+      "learning_rate": 4.170205208855281e-06,
+      "loss": 0.0021,
+      "step": 269
+    },
+    {
+      "epoch": 2.1259842519685037,
+      "grad_norm": 0.18673180043697357,
+      "learning_rate": 4.101614429000857e-06,
+      "loss": 0.0026,
+      "step": 270
+    },
+    {
+      "epoch": 2.1338582677165356,
+      "grad_norm": 0.13400611281394958,
+      "learning_rate": 4.033446592777558e-06,
+      "loss": 0.0045,
+      "step": 271
+    },
+    {
+      "epoch": 2.141732283464567,
+      "grad_norm": 0.08963260799646378,
+      "learning_rate": 3.965706588158865e-06,
+      "loss": 0.002,
+      "step": 272
+    },
+    {
+      "epoch": 2.1496062992125986,
+      "grad_norm": 0.07362519204616547,
+      "learning_rate": 3.89839927244058e-06,
+      "loss": 0.0008,
+      "step": 273
+    },
+    {
+      "epoch": 2.15748031496063,
+      "grad_norm": 0.12438540160655975,
+      "learning_rate": 3.8315294718925656e-06,
+      "loss": 0.0032,
+      "step": 274
+    },
+    {
+      "epoch": 2.1653543307086616,
+      "grad_norm": 0.07505560666322708,
+      "learning_rate": 3.7651019814126656e-06,
+      "loss": 0.0011,
+      "step": 275
+    },
+    {
+      "epoch": 2.173228346456693,
+      "grad_norm": 0.24100656807422638,
+      "learning_rate": 3.6991215641828903e-06,
+      "loss": 0.0039,
+      "step": 276
+    },
+    {
+      "epoch": 2.1811023622047245,
+      "grad_norm": 0.08774268627166748,
+      "learning_rate": 3.6335929513278667e-06,
+      "loss": 0.0021,
+      "step": 277
+    },
+    {
+      "epoch": 2.188976377952756,
+      "grad_norm": 0.06761056184768677,
+      "learning_rate": 3.568520841575601e-06,
+      "loss": 0.0004,
+      "step": 278
+    },
+    {
+      "epoch": 2.1968503937007875,
+      "grad_norm": 0.514453113079071,
+      "learning_rate": 3.5039099009205503e-06,
+      "loss": 0.002,
+      "step": 279
+    },
+    {
+      "epoch": 2.204724409448819,
+      "grad_norm": 0.1681102067232132,
+      "learning_rate": 3.439764762289051e-06,
+      "loss": 0.0049,
+      "step": 280
+    },
+    {
+      "epoch": 2.2125984251968505,
+      "grad_norm": 0.46447646617889404,
+      "learning_rate": 3.376090025207115e-06,
+      "loss": 0.0037,
+      "step": 281
+    },
+    {
+      "epoch": 2.220472440944882,
+      "grad_norm": 0.09738212823867798,
+      "learning_rate": 3.312890255470609e-06,
+      "loss": 0.0018,
+      "step": 282
+    },
+    {
+      "epoch": 2.2283464566929134,
+      "grad_norm": 0.12760388851165771,
+      "learning_rate": 3.250169984817897e-06,
+      "loss": 0.0022,
+      "step": 283
+    },
+    {
+      "epoch": 2.236220472440945,
+      "grad_norm": 0.05433168262243271,
+      "learning_rate": 3.187933710604847e-06,
+      "loss": 0.0005,
+      "step": 284
+    },
+    {
+      "epoch": 2.2440944881889764,
+      "grad_norm": 0.06812359392642975,
+      "learning_rate": 3.1261858954823798e-06,
+      "loss": 0.0007,
+      "step": 285
+    },
+    {
+      "epoch": 2.251968503937008,
+      "grad_norm": 0.44168326258659363,
+      "learning_rate": 3.064930967076477e-06,
+      "loss": 0.0052,
+      "step": 286
+    },
+    {
+      "epoch": 2.2598425196850394,
+      "grad_norm": 0.4508403241634369,
+      "learning_rate": 3.0041733176706668e-06,
+      "loss": 0.0049,
+      "step": 287
+    },
+    {
+      "epoch": 2.267716535433071,
+      "grad_norm": 0.00029889008146710694,
+      "learning_rate": 2.943917303891107e-06,
+      "loss": 0.0,
+      "step": 288
+    },
+    {
+      "epoch": 2.2755905511811023,
+      "grad_norm": 0.16293245553970337,
+      "learning_rate": 2.8841672463941827e-06,
+      "loss": 0.0052,
+      "step": 289
+    },
+    {
+      "epoch": 2.283464566929134,
+      "grad_norm": 0.0034355763345956802,
+      "learning_rate": 2.8249274295566863e-06,
+      "loss": 0.0,
+      "step": 290
+    },
+    {
+      "epoch": 2.2913385826771653,
+      "grad_norm": 0.41321080923080444,
+      "learning_rate": 2.766202101168628e-06,
+      "loss": 0.0042,
+      "step": 291
+    },
+    {
+      "epoch": 2.2992125984251968,
+      "grad_norm": 0.05302264913916588,
+      "learning_rate": 2.7079954721286108e-06,
+      "loss": 0.0008,
+      "step": 292
+    },
+    {
+      "epoch": 2.3070866141732282,
+      "grad_norm": 0.16997075080871582,
+      "learning_rate": 2.6503117161419246e-06,
+      "loss": 0.0049,
+      "step": 293
+    },
+    {
+      "epoch": 2.3149606299212597,
+      "grad_norm": 0.15489016473293304,
+      "learning_rate": 2.5931549694212545e-06,
+      "loss": 0.0029,
+      "step": 294
+    },
+    {
+      "epoch": 2.322834645669291,
+      "grad_norm": 0.040922824293375015,
+      "learning_rate": 2.536529330390095e-06,
+      "loss": 0.0003,
+      "step": 295
+    },
+    {
+      "epoch": 2.3307086614173227,
+      "grad_norm": 0.15096415579319,
+      "learning_rate": 2.480438859388873e-06,
+      "loss": 0.0037,
+      "step": 296
+    },
+    {
+      "epoch": 2.338582677165354,
+      "grad_norm": 0.05358278378844261,
+      "learning_rate": 2.424887578383799e-06,
+      "loss": 0.0004,
+      "step": 297
+    },
+    {
+      "epoch": 2.3464566929133857,
+      "grad_norm": 0.16193096339702606,
+      "learning_rate": 2.36987947067848e-06,
+      "loss": 0.0025,
+      "step": 298
+    },
+    {
+      "epoch": 2.354330708661417,
+      "grad_norm": 0.10353274643421173,
+      "learning_rate": 2.3154184806282863e-06,
+      "loss": 0.0021,
+      "step": 299
+    },
+    {
+      "epoch": 2.362204724409449,
+      "grad_norm": 0.10735179483890533,
+      "learning_rate": 2.261508513357532e-06,
+      "loss": 0.0035,
+      "step": 300
+    },
+    {
+      "epoch": 2.3700787401574805,
+      "grad_norm": 0.18752367794513702,
+      "learning_rate": 2.208153434479442e-06,
+      "loss": 0.0039,
+      "step": 301
+    },
+    {
+      "epoch": 2.3700787401574805,
+      "eval_loss": 0.30320534110069275,
+      "eval_runtime": 6.5784,
+      "eval_samples_per_second": 24.778,
+      "eval_steps_per_second": 3.192,
+      "step": 301
+    },
+    {
+      "epoch": 2.377952755905512,
+      "grad_norm": 0.13881297409534454,
+      "learning_rate": 2.155357069818995e-06,
+      "loss": 0.0032,
+      "step": 302
+    },
+    {
+      "epoch": 2.3858267716535435,
+      "grad_norm": 0.09920285642147064,
+      "learning_rate": 2.1031232051385606e-06,
+      "loss": 0.0021,
+      "step": 303
+    },
+    {
+      "epoch": 2.393700787401575,
+      "grad_norm": 0.37194201350212097,
+      "learning_rate": 2.0514555858664663e-06,
+      "loss": 0.0045,
+      "step": 304
+    },
+    {
+      "epoch": 2.4015748031496065,
+      "grad_norm": 0.10560385882854462,
+      "learning_rate": 2.000357916828428e-06,
+      "loss": 0.0011,
+      "step": 305
+    },
+    {
+      "epoch": 2.409448818897638,
+      "grad_norm": 0.33549824357032776,
+      "learning_rate": 1.949833861981877e-06,
+      "loss": 0.0039,
+      "step": 306
+    },
+    {
+      "epoch": 2.4173228346456694,
+      "grad_norm": 0.3969619870185852,
+      "learning_rate": 1.8998870441532569e-06,
+      "loss": 0.0027,
+      "step": 307
+    },
+    {
+      "epoch": 2.425196850393701,
+      "grad_norm": 0.081158846616745,
+      "learning_rate": 1.8505210447782418e-06,
+      "loss": 0.0011,
+      "step": 308
+    },
+    {
+      "epoch": 2.4330708661417324,
+      "grad_norm": 0.28652095794677734,
+      "learning_rate": 1.8017394036449276e-06,
+      "loss": 0.0038,
+      "step": 309
+    },
+    {
+      "epoch": 2.440944881889764,
+      "grad_norm": 0.0656951516866684,
+      "learning_rate": 1.7535456186400123e-06,
+      "loss": 0.001,
+      "step": 310
+    },
+    {
+      "epoch": 2.4488188976377954,
+      "grad_norm": 0.14871421456336975,
+      "learning_rate": 1.7059431454979825e-06,
+      "loss": 0.0027,
+      "step": 311
+    },
+    {
+      "epoch": 2.456692913385827,
+      "grad_norm": 0.25429457426071167,
+      "learning_rate": 1.6589353975533174e-06,
+      "loss": 0.0012,
+      "step": 312
+    },
+    {
+      "epoch": 2.4645669291338583,
+      "grad_norm": 0.06939385086297989,
+      "learning_rate": 1.6125257454957365e-06,
+      "loss": 0.0008,
+      "step": 313
+    },
+    {
+      "epoch": 2.47244094488189,
+      "grad_norm": 0.15781065821647644,
+      "learning_rate": 1.5667175171285054e-06,
+      "loss": 0.003,
+      "step": 314
+    },
+    {
+      "epoch": 2.4803149606299213,
+      "grad_norm": 0.08229056000709534,
+      "learning_rate": 1.5215139971298131e-06,
+      "loss": 0.0015,
+      "step": 315
+    },
+    {
+      "epoch": 2.4881889763779528,
+      "grad_norm": 0.16827985644340515,
+      "learning_rate": 1.4769184268172465e-06,
+      "loss": 0.0032,
+      "step": 316
+    },
+    {
+      "epoch": 2.4960629921259843,
+      "grad_norm": 0.12261717021465302,
+      "learning_rate": 1.4329340039153738e-06,
+      "loss": 0.0022,
+      "step": 317
+    },
+    {
+      "epoch": 2.5039370078740157,
+      "grad_norm": 0.1208304911851883,
+      "learning_rate": 1.3895638823264447e-06,
+      "loss": 0.002,
+      "step": 318
+    },
+    {
+      "epoch": 2.5118110236220472,
+      "grad_norm": 0.22991932928562164,
+      "learning_rate": 1.3468111719042497e-06,
+      "loss": 0.0027,
+      "step": 319
+    },
+    {
+      "epoch": 2.5196850393700787,
+      "grad_norm": 0.468462198972702,
+      "learning_rate": 1.3046789382311132e-06,
+      "loss": 0.0042,
+      "step": 320
+    },
+    {
+      "epoch": 2.52755905511811,
+      "grad_norm": 0.029908303171396255,
+      "learning_rate": 1.2631702023980997e-06,
+      "loss": 0.0002,
+      "step": 321
+    },
+    {
+      "epoch": 2.5354330708661417,
+      "grad_norm": 0.07678980380296707,
+      "learning_rate": 1.2222879407883592e-06,
+      "loss": 0.0014,
+      "step": 322
+    },
+    {
+      "epoch": 2.543307086614173,
+      "grad_norm": 0.13547496497631073,
+      "learning_rate": 1.182035084863724e-06,
+      "loss": 0.0017,
+      "step": 323
+    },
+    {
+      "epoch": 2.5511811023622046,
+      "grad_norm": 0.15075382590293884,
+      "learning_rate": 1.1424145209545079e-06,
+      "loss": 0.0059,
+      "step": 324
+    },
+    {
+      "epoch": 2.559055118110236,
+      "grad_norm": 0.1271948516368866,
+      "learning_rate": 1.1034290900525279e-06,
+      "loss": 0.0021,
+      "step": 325
+    },
+    {
+      "epoch": 2.5669291338582676,
+      "grad_norm": 0.11441997438669205,
+      "learning_rate": 1.065081587607406e-06,
+      "loss": 0.0022,
+      "step": 326
+    },
+    {
+      "epoch": 2.574803149606299,
+      "grad_norm": 0.13326182961463928,
+      "learning_rate": 1.0273747633261144e-06,
+      "loss": 0.004,
+      "step": 327
+    },
+    {
+      "epoch": 2.5826771653543306,
+      "grad_norm": 0.07804345339536667,
+      "learning_rate": 9.903113209758098e-07,
+      "loss": 0.0018,
+      "step": 328
+    },
+    {
+      "epoch": 2.590551181102362,
+      "grad_norm": 0.0012728713918477297,
+      "learning_rate": 9.538939181899565e-07,
+      "loss": 0.0,
+      "step": 329
+    },
+    {
+      "epoch": 2.5984251968503935,
+      "grad_norm": 0.06427028775215149,
+      "learning_rate": 9.181251662777668e-07,
+      "loss": 0.0007,
+      "step": 330
+    },
+    {
+      "epoch": 2.606299212598425,
+      "grad_norm": 0.1923428475856781,
+      "learning_rate": 8.830076300369517e-07,
+      "loss": 0.006,
+      "step": 331
+    },
+    {
+      "epoch": 2.6141732283464565,
+      "grad_norm": 0.33056169748306274,
+      "learning_rate": 8.485438275698154e-07,
+      "loss": 0.0024,
+      "step": 332
+    },
+    {
+      "epoch": 2.622047244094488,
+      "grad_norm": 0.13692541420459747,
+      "learning_rate": 8.14736230102694e-07,
+      "loss": 0.0019,
+      "step": 333
+    },
+    {
+      "epoch": 2.6299212598425195,
+      "grad_norm": 0.11543405055999756,
+      "learning_rate": 7.815872618087506e-07,
+      "loss": 0.003,
+      "step": 334
+    },
+    {
+      "epoch": 2.637795275590551,
+      "grad_norm": 0.20871274173259735,
+      "learning_rate": 7.490992996341662e-07,
+      "loss": 0.0022,
+      "step": 335
+    },
+    {
+      "epoch": 2.6456692913385824,
+      "grad_norm": 0.1506434828042984,
+      "learning_rate": 7.17274673127677e-07,
+      "loss": 0.0034,
+      "step": 336
+    },
+    {
+      "epoch": 2.653543307086614,
+      "grad_norm": 0.1000061109662056,
+      "learning_rate": 6.861156642735578e-07,
+      "loss": 0.0015,
+      "step": 337
+    },
+    {
+      "epoch": 2.661417322834646,
+      "grad_norm": 0.04730301722884178,
+      "learning_rate": 6.556245073279777e-07,
+      "loss": 0.0003,
+      "step": 338
+    },
+    {
+      "epoch": 2.6692913385826773,
+      "grad_norm": 0.07712409645318985,
+      "learning_rate": 6.258033886587911e-07,
+      "loss": 0.0006,
+      "step": 339
+    },
+    {
+      "epoch": 2.677165354330709,
+      "grad_norm": 0.12951001524925232,
+      "learning_rate": 5.966544465887803e-07,
+      "loss": 0.0022,
+      "step": 340
+    },
+    {
+      "epoch": 2.6850393700787403,
+      "grad_norm": 0.3450707495212555,
+      "learning_rate": 5.681797712423099e-07,
+      "loss": 0.0031,
+      "step": 341
+    },
+    {
+      "epoch": 2.6929133858267718,
+      "grad_norm": 0.11356323957443237,
+      "learning_rate": 5.403814043954592e-07,
+      "loss": 0.0016,
+      "step": 342
+    },
+    {
+      "epoch": 2.7007874015748032,
+      "grad_norm": 0.40962764620780945,
+      "learning_rate": 5.132613393296293e-07,
+      "loss": 0.0022,
+      "step": 343
+    },
+    {
+      "epoch": 2.7086614173228347,
+      "grad_norm": 0.0026160525158047676,
+      "learning_rate": 4.868215206885918e-07,
+      "loss": 0.0,
+      "step": 344
+    },
+    {
+      "epoch": 2.7086614173228347,
+      "eval_loss": 0.3102666437625885,
+      "eval_runtime": 7.3029,
+      "eval_samples_per_second": 22.32,
+      "eval_steps_per_second": 2.876,
+      "step": 344
+    },
+    {
+      "epoch": 2.716535433070866,
+      "grad_norm": 0.2460733950138092,
+      "learning_rate": 4.61063844339068e-07,
+      "loss": 0.0044,
+      "step": 345
+    },
+    {
+      "epoch": 2.7244094488188977,
+      "grad_norm": 0.11104279011487961,
+      "learning_rate": 4.359901572347758e-07,
+      "loss": 0.0031,
+      "step": 346
+    },
+    {
+      "epoch": 2.732283464566929,
+      "grad_norm": 0.288809210062027,
+      "learning_rate": 4.116022572839984e-07,
+      "loss": 0.0023,
+      "step": 347
+    },
+    {
+      "epoch": 2.7401574803149606,
+      "grad_norm": 0.2904239892959595,
+      "learning_rate": 3.879018932206624e-07,
+      "loss": 0.001,
+      "step": 348
+    },
+    {
+      "epoch": 2.748031496062992,
+      "grad_norm": 0.5172310471534729,
+      "learning_rate": 3.6489076447894456e-07,
+      "loss": 0.0023,
+      "step": 349
+    },
+    {
+      "epoch": 2.7559055118110236,
+      "grad_norm": 0.555241048336029,
+      "learning_rate": 3.425705210714192e-07,
+      "loss": 0.0026,
+      "step": 350
+    },
+    {
+      "epoch": 2.763779527559055,
+      "grad_norm": 0.12381427735090256,
+      "learning_rate": 3.2094276347073626e-07,
+      "loss": 0.002,
+      "step": 351
+    },
+    {
+      "epoch": 2.7716535433070866,
+      "grad_norm": 0.16744810342788696,
+      "learning_rate": 3.000090424948665e-07,
+      "loss": 0.0036,
+      "step": 352
+    },
+    {
+      "epoch": 2.779527559055118,
+      "grad_norm": 0.512416422367096,
+      "learning_rate": 2.7977085919589253e-07,
+      "loss": 0.0026,
+      "step": 353
+    },
+    {
+      "epoch": 2.7874015748031495,
+      "grad_norm": 0.18864978849887848,
+      "learning_rate": 2.602296647523861e-07,
+      "loss": 0.0025,
+      "step": 354
+    },
+    {
+      "epoch": 2.795275590551181,
+      "grad_norm": 0.012189923785626888,
+      "learning_rate": 2.413868603653413e-07,
+      "loss": 0.0001,
+      "step": 355
+    },
+    {
+      "epoch": 2.8031496062992125,
+      "grad_norm": 0.10027037560939789,
+      "learning_rate": 2.2324379715770728e-07,
+      "loss": 0.0011,
+      "step": 356
+    },
+    {
+      "epoch": 2.811023622047244,
+      "grad_norm": 0.10117685794830322,
+      "learning_rate": 2.0580177607750663e-07,
+      "loss": 0.0036,
+      "step": 357
+    },
+    {
+      "epoch": 2.8188976377952755,
+      "grad_norm": 0.1535252332687378,
+      "learning_rate": 1.890620478045435e-07,
+      "loss": 0.0044,
+      "step": 358
+    },
+    {
+      "epoch": 2.826771653543307,
+      "grad_norm": 0.39140409231185913,
+      "learning_rate": 1.7302581266073537e-07,
+      "loss": 0.0037,
+      "step": 359
+    },
+    {
+      "epoch": 2.8346456692913384,
+      "grad_norm": 0.18143348395824432,
+      "learning_rate": 1.5769422052403172e-07,
+      "loss": 0.0033,
+      "step": 360
+    },
+    {
+      "epoch": 2.84251968503937,
+      "grad_norm": 0.6282801032066345,
+      "learning_rate": 1.4306837074597235e-07,
+      "loss": 0.0096,
+      "step": 361
+    },
+    {
+      "epoch": 2.850393700787402,
+      "grad_norm": 0.3672868311405182,
+      "learning_rate": 1.2914931207285154e-07,
+      "loss": 0.0014,
+      "step": 362
+    },
+    {
+      "epoch": 2.8582677165354333,
+      "grad_norm": 0.13403712213039398,
+      "learning_rate": 1.1593804257052143e-07,
+      "loss": 0.0046,
+      "step": 363
+    },
+    {
+      "epoch": 2.866141732283465,
+      "grad_norm": 0.004047819878906012,
+      "learning_rate": 1.0343550955282278e-07,
+      "loss": 0.0,
+      "step": 364
+    },
+    {
+      "epoch": 2.8740157480314963,
+      "grad_norm": 0.3351942002773285,
+      "learning_rate": 9.164260951366021e-08,
+      "loss": 0.0024,
+      "step": 365
+    },
+    {
+      "epoch": 2.8818897637795278,
+      "grad_norm": 0.09759978204965591,
+      "learning_rate": 8.056018806271937e-08,
+      "loss": 0.002,
+      "step": 366
+    },
+    {
+      "epoch": 2.8897637795275593,
+      "grad_norm": 0.06213594600558281,
+      "learning_rate": 7.018903986483083e-08,
+      "loss": 0.0009,
+      "step": 367
+    },
+    {
+      "epoch": 2.8976377952755907,
+      "grad_norm": 0.07074209302663803,
+      "learning_rate": 6.052990858298801e-08,
+      "loss": 0.0009,
+      "step": 368
+    },
+    {
+      "epoch": 2.905511811023622,
+      "grad_norm": 0.271335631608963,
+      "learning_rate": 5.158348682502756e-08,
+      "loss": 0.0037,
+      "step": 369
+    },
+    {
+      "epoch": 2.9133858267716537,
+      "grad_norm": 0.09063868969678879,
+      "learning_rate": 4.335041609396018e-08,
+      "loss": 0.0014,
+      "step": 370
+    },
+    {
+      "epoch": 2.921259842519685,
+      "grad_norm": 0.818594753742218,
+      "learning_rate": 3.5831286741973006e-08,
+      "loss": 0.0033,
+      "step": 371
+    },
+    {
+      "epoch": 2.9291338582677167,
+      "grad_norm": 0.09543661028146744,
+      "learning_rate": 2.902663792810012e-08,
+      "loss": 0.0015,
+      "step": 372
+    },
+    {
+      "epoch": 2.937007874015748,
+      "grad_norm": 0.13098907470703125,
+      "learning_rate": 2.293695757956571e-08,
+      "loss": 0.0037,
+      "step": 373
+    },
+    {
+      "epoch": 2.9448818897637796,
+      "grad_norm": 0.5491423010826111,
+      "learning_rate": 1.7562682356786488e-08,
+      "loss": 0.004,
+      "step": 374
+    },
+    {
+      "epoch": 2.952755905511811,
+      "grad_norm": 0.08357255905866623,
+      "learning_rate": 1.290419762207007e-08,
+      "loss": 0.0015,
+      "step": 375
+    },
+    {
+      "epoch": 2.9606299212598426,
+      "grad_norm": 0.24269114434719086,
+      "learning_rate": 8.961837411982643e-09,
+      "loss": 0.0028,
+      "step": 376
+    },
+    {
+      "epoch": 2.968503937007874,
+      "grad_norm": 0.1084604412317276,
+      "learning_rate": 5.735884413391457e-09,
+      "loss": 0.0022,
+      "step": 377
+    },
+    {
+      "epoch": 2.9763779527559056,
+      "grad_norm": 0.09172981232404709,
+      "learning_rate": 3.226569943197699e-09,
+      "loss": 0.0022,
+      "step": 378
+    },
+    {
+      "epoch": 2.984251968503937,
+      "grad_norm": 0.1312946230173111,
+      "learning_rate": 1.4340739317497688e-09,
+      "loss": 0.002,
+      "step": 379
+    },
+    {
+      "epoch": 2.9921259842519685,
+      "grad_norm": 0.0002817026397679001,
+      "learning_rate": 3.585249099435917e-10,
+      "loss": 0.0,
+      "step": 380
+    },
+    {
+      "epoch": 3.0,
+      "grad_norm": 0.2318553477525711,
+      "learning_rate": 0.0,
+      "loss": 0.0007,
+      "step": 381
+    }
+  ],
+  "logging_steps": 1,
+  "max_steps": 381,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 3,
+  "save_steps": 127,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 3.9261813209667994e+17,
+  "train_batch_size": 128,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-381/training_args.bin

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