Upload folder using huggingface_hub

README.md

@@ -0,0 +1,1091 @@
+---
+library_name: peft
+base_model: google/gemma-2b
+---
+
+# Model Card for Model ID
+
+Fine-tuned on python
+
+
+## Model Details
+
+### Model Description
+
+
+### Model Sources [optional]
+Gemma-2b trained on python-oasst dataset
+
+## Uses
+
+
+## Training Details
+
+### Training Data
+
+{
+  "_timestamp": 1711018613.433522,
+  "train/grad_norm": 0.1240904619429259,
+  "train/global_step": 232,
+  "eval/steps_per_second": 2.894,
+  "_step": 232,
+  "_runtime": 2545.4226660728455,
+  "eval/loss": 1.189491629600525,
+  "eval/runtime": 1805.8574,
+  "train/learning_rate": 0.000014800637958532697,
+  "eval/samples_per_second": 23.152,
+  "_wandb.runtime": 2547,
+  "train/loss": 1.0436,
+  "train/epoch": 0.01
+}
+
+### Results
+
+
+#### Summary
+
+
+
+
+## Technical Specifications [optional]
+{
+  "bf16": {
+    "desc": null,
+    "value": true
+  },
+  "fp16": {
+    "desc": null,
+    "value": false
+  },
+  "fsdp": {
+    "desc": null,
+    "value": []
+  },
+  "seed": {
+    "desc": null,
+    "value": 42
+  },
+  "tf32": {
+    "desc": null,
+    "value": false
+  },
+  "debug": {
+    "desc": null,
+    "value": []
+  },
+  "optim": {
+    "desc": null,
+    "value": "adamw_bnb_8bit"
+  },
+  "qlora": {
+    "desc": null,
+    "value": true
+  },
+  "top_k": {
+    "desc": null,
+    "value": 50
+  },
+  "top_p": {
+    "desc": null,
+    "value": 1
+  },
+  "_wandb": {
+    "desc": null,
+    "value": {
+      "m": [
+        {
+          "1": "train/global_step",
+          "6": [
+            3
+          ]
+        },
+        {
+          "1": "train/loss",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "train/grad_norm",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "train/learning_rate",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "train/epoch",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "eval/loss",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "eval/runtime",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "eval/samples_per_second",
+          "5": 1,
+          "6": [
+            1
+          ]
+        },
+        {
+          "1": "eval/steps_per_second",
+          "5": 1,
+          "6": [
+            1
+          ]
+        }
+      ],
+      "t": {
+        "1": [
+          1,
+          5,
+          11,
+          49,
+          51,
+          53,
+          55,
+          71,
+          84,
+          98,
+          99,
+          100,
+          105
+        ],
+        "2": [
+          1,
+          5,
+          11,
+          49,
+          51,
+          53,
+          55,
+          71,
+          84,
+          98,
+          99,
+          100,
+          105
+        ],
+        "3": [
+          3,
+          7,
+          23
+        ],
+        "4": "3.10.13",
+        "5": "0.16.4",
+        "6": "4.39.0.dev0",
+        "8": [
+          5
+        ],
+        "9": {
+          "1": "transformers_trainer"
+        },
+        "13": "linux-x86_64"
+      },
+      "framework": "huggingface",
+      "start_time": 1711016068,
+      "cli_version": "0.16.4",
+      "is_jupyter_run": false,
+      "python_version": "3.10.13",
+      "is_kaggle_kernel": false,
+      "huggingface_version": "4.39.0.dev0"
+    }
+  },
+  "prefix": {
+    "desc": null,
+    "value": null
+  },
+  "do_eval": {
+    "desc": null,
+    "value": true
+  },
+  "no_cuda": {
+    "desc": null,
+    "value": false
+  },
+  "use_cpu": {
+    "desc": null,
+    "value": false
+  },
+  "do_train": {
+    "desc": null,
+    "value": false
+  },
+  "head_dim": {
+    "desc": null,
+    "value": 256
+  },
+  "id2label": {
+    "desc": null,
+    "value": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    }
+  },
+  "label2id": {
+    "desc": null,
+    "value": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    }
+  },
+  "run_name": {
+    "desc": null,
+    "value": "./out"
+  },
+  "use_ipex": {
+    "desc": null,
+    "value": false
+  },
+  "adafactor": {
+    "desc": null,
+    "value": false
+  },
+  "data_seed": {
+    "desc": null,
+    "value": null
+  },
+  "deepspeed": {
+    "desc": null,
+    "value": "deepspeed_configs/zero1.json"
+  },
+  "do_sample": {
+    "desc": null,
+    "value": false
+  },
+  "hub_token": {
+    "desc": null,
+  },
+  "log_level": {
+    "desc": null,
+    "value": "passive"
+  },
+  "max_steps": {
+    "desc": null,
+    "value": -1
+  },
+  "num_beams": {
+    "desc": null,
+    "value": 1
+  },
+  "ray_scope": {
+    "desc": null,
+    "value": "last"
+  },
+  "report_to": {
+    "desc": null,
+    "value": [
+      "wandb"
+    ]
+  },
+  "typical_p": {
+    "desc": null,
+    "value": 1
+  },
+  "use_cache": {
+    "desc": null,
+    "value": false
+  },
+  "adam_beta1": {
+    "desc": null,
+    "value": 0.9
+  },
+  "adam_beta2": {
+    "desc": null,
+    "value": 0.999
+  },
+  "do_predict": {
+    "desc": null,
+    "value": false
+  },
+  "eval_delay": {
+    "desc": null,
+    "value": 0
+  },
+  "eval_steps": {
+    "desc": null,
+    "value": 0.03125
+  },
+  "hidden_act": {
+    "desc": null,
+    "value": "gelu"
+  },
+  "is_decoder": {
+    "desc": null,
+    "value": false
+  },
+  "local_rank": {
+    "desc": null,
+    "value": 0
+  },
+  "max_length": {
+    "desc": null,
+    "value": 20
+  },
+  "min_length": {
+    "desc": null,
+    "value": 0
+  },
+  "model_type": {
+    "desc": null,
+    "value": "gemma"
+  },
+  "optim_args": {
+    "desc": null,
+    "value": null
+  },
+  "orpo_alpha": {
+    "desc": null,
+    "value": null
+  },
+  "output_dir": {
+    "desc": null,
+    "value": "./out"
+  },
+  "past_index": {
+    "desc": null,
+    "value": -1
+  },
+  "rope_theta": {
+    "desc": null,
+    "value": 10000
+  },
+  "save_steps": {
+    "desc": null,
+    "value": 0.125
+  },
+  "vocab_size": {
+    "desc": null,
+    "value": 256000
+  },
+  "bench_split": {
+    "desc": null,
+    "value": "eval"
+  },
+  "ddp_backend": {
+    "desc": null,
+    "value": null
+  },
+  "ddp_timeout": {
+    "desc": null,
+    "value": 1800
+  },
+  "fsdp_config": {
+    "desc": null,
+    "value": {
+      "xla": false,
+      "xla_fsdp_v2": false,
+      "min_num_params": 0,
+      "xla_fsdp_grad_ckpt": false
+    }
+  },
+  "hidden_size": {
+    "desc": null,
+    "value": 2048
+  },
+  "label_names": {
+    "desc": null,
+    "value": null
+  },
+  "logging_dir": {
+    "desc": null,
+    "value": "./out/runs/Mar21_10-14-24_8205afe3ecd2"
+  },
+  "pretraining": {
+    "desc": null,
+    "value": false
+  },
+  "push_to_hub": {
+    "desc": null,
+    "value": false
+  },
+  "return_dict": {
+    "desc": null,
+    "value": true
+  },
+  "temperature": {
+    "desc": null,
+    "value": 1
+  },
+  "torch_dtype": {
+    "desc": null,
+    "value": "bfloat16"
+  },
+  "torchdynamo": {
+    "desc": null,
+    "value": null
+  },
+  "torchscript": {
+    "desc": null,
+    "value": false
+  },
+  "adam_epsilon": {
+    "desc": null,
+    "value": 1e-8
+  },
+  "bos_token_id": {
+    "desc": null,
+    "value": 2
+  },
+  "disable_tqdm": {
+    "desc": null,
+    "value": false
+  },
+  "eos_token_id": {
+    "desc": null,
+    "value": 1
+  },
+  "fp16_backend": {
+    "desc": null,
+    "value": "auto"
+  },
+  "hub_model_id": {
+    "desc": null,
+    "value": null
+  },
+  "hub_strategy": {
+    "desc": null,
+    "value": "every_save"
+  },
+  "pad_token_id": {
+    "desc": null,
+    "value": 0
+  },
+  "problem_type": {
+    "desc": null,
+    "value": null
+  },
+  "pruned_heads": {
+    "desc": null,
+    "value": {}
+  },
+  "relora_steps": {
+    "desc": null,
+    "value": null
+  },
+  "rms_norm_eps": {
+    "desc": null,
+    "value": 0.000001
+  },
+  "rope_scaling": {
+    "desc": null,
+    "value": null
+  },
+  "sep_token_id": {
+    "desc": null,
+    "value": null
+  },
+  "use_bfloat16": {
+    "desc": null,
+    "value": false
+  },
+  "warmup_ratio": {
+    "desc": null,
+    "value": 0
+  },
+  "warmup_steps": {
+    "desc": null,
+    "value": 3135
+  },
+  "weight_decay": {
+    "desc": null,
+    "value": 0
+  },
+  "_name_or_path": {
+    "desc": null,
+    "value": "dvdmrs09/gemma2b-train"
+  },
+  "architectures": {
+    "desc": null,
+    "value": [
+      "GemmaForCausalLM"
+    ]
+  },
+  "bad_words_ids": {
+    "desc": null,
+    "value": null
+  },
+  "bench_dataset": {
+    "desc": null,
+    "value": "pharaouk/dharma-1/dharma_1_mini.json"
+  },
+  "do_bench_eval": {
+    "desc": null,
+    "value": false
+  },
+  "jit_mode_eval": {
+    "desc": null,
+    "value": false
+  },
+  "learning_rate": {
+    "desc": null,
+    "value": 0.0002
+  },
+  "logging_steps": {
+    "desc": null,
+    "value": 1
+  },
+  "max_grad_norm": {
+    "desc": null,
+    "value": 1
+  },
+  "mp_parameters": {
+    "desc": null,
+    "value": ""
+  },
+  "output_scores": {
+    "desc": null,
+    "value": false
+  },
+  "save_strategy": {
+    "desc": null,
+    "value": "steps"
+  },
+  "split_batches": {
+    "desc": null,
+    "value": null
+  },
+  "torch_compile": {
+    "desc": null,
+    "value": false
+  },
+  "tpu_num_cores": {
+    "desc": null,
+    "value": null
+  },
+  "attention_bias": {
+    "desc": null,
+    "value": false
+  },
+  "bf16_full_eval": {
+    "desc": null,
+    "value": false
+  },
+  "early_stopping": {
+    "desc": null,
+    "value": false
+  },
+  "fp16_full_eval": {
+    "desc": null,
+    "value": false
+  },
+  "fp16_opt_level": {
+    "desc": null,
+    "value": "O1"
+  },
+  "length_penalty": {
+    "desc": null,
+    "value": 1
+  },
+  "max_seq_length": {
+    "desc": null,
+    "value": 4096
+  },
+  "sample_packing": {
+    "desc": null,
+    "value": false
+  },
+  "tf_legacy_loss": {
+    "desc": null,
+    "value": false
+  },
+  "use_mps_device": {
+    "desc": null,
+    "value": false
+  },
+  "finetuning_task": {
+    "desc": null,
+    "value": null
+  },
+  "group_by_length": {
+    "desc": null,
+    "value": false
+  },
+  "hub_always_push": {
+    "desc": null,
+    "value": false
+  },
+  "num_beam_groups": {
+    "desc": null,
+    "value": 1
+  },
+  "save_only_model": {
+    "desc": null,
+    "value": false
+  },
+  "suppress_tokens": {
+    "desc": null,
+    "value": null
+  },
+  "tokenizer_class": {
+    "desc": null,
+    "value": null
+  },
+  "dispatch_batches": {
+    "desc": null,
+    "value": null
+  },
+  "full_determinism": {
+    "desc": null,
+    "value": false
+  },
+  "hub_private_repo": {
+    "desc": null,
+    "value": false
+  },
+  "ignore_data_skip": {
+    "desc": null,
+    "value": false
+  },
+  "log_on_each_node": {
+    "desc": null,
+    "value": true
+  },
+  "logging_strategy": {
+    "desc": null,
+    "value": "steps"
+  },
+  "num_train_epochs": {
+    "desc": null,
+    "value": 8
+  },
+  "save_safetensors": {
+    "desc": null,
+    "value": true
+  },
+  "save_total_limit": {
+    "desc": null,
+    "value": 4
+  },
+  "attention_dropout": {
+    "desc": null,
+    "value": 0
+  },
+  "ddp_bucket_cap_mb": {
+    "desc": null,
+    "value": null
+  },
+  "diversity_penalty": {
+    "desc": null,
+    "value": 0
+  },
+  "do_causal_lm_eval": {
+    "desc": null,
+    "value": false
+  },
+  "greater_is_better": {
+    "desc": null,
+    "value": false
+  },
+  "initializer_range": {
+    "desc": null,
+    "value": 0.02
+  },
+  "intermediate_size": {
+    "desc": null,
+    "value": 16384
+  },
+  "log_level_replica": {
+    "desc": null,
+    "value": "warning"
+  },
+  "loraplus_lr_ratio": {
+    "desc": null,
+    "value": null
+  },
+  "lr_scheduler_type": {
+    "desc": null,
+    "value": "cosine"
+  },
+  "max_bench_samples": {
+    "desc": null,
+    "value": null
+  },
+  "num_hidden_layers": {
+    "desc": null,
+    "value": 18
+  },
+  "output_attentions": {
+    "desc": null,
+    "value": false
+  },
+  "push_to_hub_token": {
+    "desc": null,
+    "value": "<PUSH_TO_HUB_TOKEN>"
+  },
+  "save_on_each_node": {
+    "desc": null,
+    "value": false
+  },
+  "tpu_metrics_debug": {
+    "desc": null,
+    "value": false
+  },
+  "accelerator_config": {
+    "desc": null,
+    "value": {
+      "even_batches": true,
+      "split_batches": false,
+      "dispatch_batches": null,
+      "use_seedable_sampler": true
+    }
+  },
+  "is_encoder_decoder": {
+    "desc": null,
+    "value": false
+  },
+  "length_column_name": {
+    "desc": null,
+    "value": "length"
+  },
+  "logging_first_step": {
+    "desc": null,
+    "value": false
+  },
+  "relora_prune_ratio": {
+    "desc": null,
+    "value": 0.9
+  },
+  "repetition_penalty": {
+    "desc": null,
+    "value": 1
+  },
+  "torch_compile_mode": {
+    "desc": null,
+    "value": null
+  },
+  "add_cross_attention": {
+    "desc": null,
+    "value": false
+  },
+  "cosine_min_lr_ratio": {
+    "desc": null,
+    "value": null
+  },
+  "eval_sample_packing": {
+    "desc": null,
+    "value": false
+  },
+  "evaluation_strategy": {
+    "desc": null,
+    "value": "steps"
+  },
+  "forced_bos_token_id": {
+    "desc": null,
+    "value": null
+  },
+  "forced_eos_token_id": {
+    "desc": null,
+    "value": null
+  },
+  "fsdp_min_num_params": {
+    "desc": null,
+    "value": 0
+  },
+  "lr_quadratic_warmup": {
+    "desc": null,
+    "value": false
+  },
+  "lr_scheduler_kwargs": {
+    "desc": null,
+    "value": {}
+  },
+  "neftune_noise_alpha": {
+    "desc": null,
+    "value": null
+  },
+  "num_attention_heads": {
+    "desc": null,
+    "value": 8
+  },
+  "num_key_value_heads": {
+    "desc": null,
+    "value": 1
+  },
+  "quantization_config": {
+    "desc": null,
+    "value": {
+      "load_in_4bit": true,
+      "load_in_8bit": false,
+      "quant_method": "QuantizationMethod.BITS_AND_BYTES",
+      "_load_in_4bit": true,
+      "_load_in_8bit": false,
+      "llm_int8_threshold": 6,
+      "bnb_4bit_quant_type": "nf4",
+      "llm_int8_skip_modules": null,
+      "bnb_4bit_compute_dtype": "bfloat16",
+      "bnb_4bit_quant_storage": "uint8",
+      "llm_int8_has_fp16_weight": false,
+      "bnb_4bit_use_double_quant": true,
+      "llm_int8_enable_fp32_cpu_offload": false
+    }
+  },
+  "relora_anneal_steps": {
+    "desc": null,
+    "value": null
+  },
+  "relora_warmup_steps": {
+    "desc": null,
+    "value": null
+  },
+  "skip_memory_metrics": {
+    "desc": null,
+    "value": true
+  },
+  "tie_encoder_decoder": {
+    "desc": null,
+    "value": false
+  },
+  "tie_word_embeddings": {
+    "desc": null,
+    "value": true
+  },
+  "auto_find_batch_size": {
+    "desc": null,
+    "value": false
+  },
+  "bench_source_max_len": {
+    "desc": null,
+    "value": 2048
+  },
+  "dataloader_drop_last": {
+    "desc": null,
+    "value": false
+  },
+  "no_repeat_ngram_size": {
+    "desc": null,
+    "value": 0
+  },
+  "num_return_sequences": {
+    "desc": null,
+    "value": 1
+  },
+  "optim_target_modules": {
+    "desc": null,
+    "value": null
+  },
+  "output_hidden_states": {
+    "desc": null,
+    "value": false
+  },
+  "overwrite_output_dir": {
+    "desc": null,
+    "value": false
+  },
+  "prediction_loss_only": {
+    "desc": null,
+    "value": false
+  },
+  "push_to_hub_model_id": {
+    "desc": null,
+    "value": null
+  },
+  "task_specific_params": {
+    "desc": null,
+    "value": null
+  },
+  "transformers_version": {
+    "desc": null,
+    "value": "4.39.0.dev0"
+  },
+  "begin_suppress_tokens": {
+    "desc": null,
+    "value": null
+  },
+  "dataloader_pin_memory": {
+    "desc": null,
+    "value": true
+  },
+  "ddp_broadcast_buffers": {
+    "desc": null,
+    "value": null
+  },
+  "loraplus_lr_embedding": {
+    "desc": null,
+    "value": null
+  },
+  "metric_for_best_model": {
+    "desc": null,
+    "value": "loss"
+  },
+  "remove_invalid_values": {
+    "desc": null,
+    "value": false
+  },
+  "remove_unused_columns": {
+    "desc": null,
+    "value": true
+  },
+  "torch_compile_backend": {
+    "desc": null,
+    "value": null
+  },
+  "dataloader_num_workers": {
+    "desc": null,
+    "value": 0
+  },
+  "decoder_start_token_id": {
+    "desc": null,
+    "value": null
+  },
+  "gradient_checkpointing": {
+    "desc": null,
+    "value": true
+  },
+  "half_precision_backend": {
+    "desc": null,
+    "value": "auto"
+  },
+  "label_smoothing_factor": {
+    "desc": null,
+    "value": 0
+  },
+  "load_best_model_at_end": {
+    "desc": null,
+    "value": true
+  },
+  "logging_nan_inf_filter": {
+    "desc": null,
+    "value": true
+  },
+  "multipack_real_batches": {
+    "desc": null,
+    "value": false
+  },
+  "resume_from_checkpoint": {
+    "desc": null,
+    "value": null
+  },
+  "chunk_size_feed_forward": {
+    "desc": null,
+    "value": 0
+  },
+  "eval_accumulation_steps": {
+    "desc": null,
+    "value": 3
+  },
+  "max_position_embeddings": {
+    "desc": null,
+    "value": 8192
+  },
+  "per_gpu_eval_batch_size": {
+    "desc": null,
+    "value": null
+  },
+  "return_dict_in_generate": {
+    "desc": null,
+    "value": false
+  },
+  "cosine_constant_lr_ratio": {
+    "desc": null,
+    "value": null
+  },
+  "per_gpu_train_batch_size": {
+    "desc": null,
+    "value": null
+  },
+  "push_to_hub_organization": {
+    "desc": null,
+    "value": null
+  },
+  "include_tokens_per_second": {
+    "desc": null,
+    "value": false
+  },
+  "sample_packing_efficiency": {
+    "desc": null,
+    "value": 1
+  },
+  "dataloader_prefetch_factor": {
+    "desc": null,
+    "value": null
+  },
+  "ddp_find_unused_parameters": {
+    "desc": null,
+    "value": false
+  },
+  "include_inputs_for_metrics": {
+    "desc": null,
+    "value": false
+  },
+  "per_device_eval_batch_size": {
+    "desc": null,
+    "value": 2
+  },
+  "use_legacy_prediction_loop": {
+    "desc": null,
+    "value": false
+  },
+  "cross_attention_hidden_size": {
+    "desc": null,
+    "value": null
+  },
+  "gradient_accumulation_steps": {
+    "desc": null,
+    "value": 3
+  },
+  "per_device_train_batch_size": {
+    "desc": null,
+    "value": 2
+  },
+  "encoder_no_repeat_ngram_size": {
+    "desc": null,
+    "value": 0
+  },
+  "dataloader_persistent_workers": {
+    "desc": null,
+    "value": false
+  },
+  "gradient_checkpointing_kwargs": {
+    "desc": null,
+    "value": {
+      "use_reentrant": true
+    }
+  },
+  "include_num_input_tokens_seen": {
+    "desc": null,
+    "value": false
+  },
+  "exponential_decay_length_penalty": {
+    "desc": null,
+    "value": null
+  },
+  "sample_packing_seq_len_multiplier": {
+    "desc": null,
+    "value": 2
+  },
+  "fsdp_transformer_layer_cls_to_wrap": {
+    "desc": null,
+    "value": null
+  }
+}
+### Model Architecture and Objective
+
+## Citation [optional]
+
+
+## Glossary [optional]
+### Framework versions
+
+- PEFT 0.9.0

adapter_config.json

@@ -0,0 +1,33 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "google/gemma-2b",
+  "bias": "none",
+  "fan_in_fan_out": null,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "up_proj",
+    "q_proj",
+    "v_proj",
+    "down_proj",
+    "gate_proj",
+    "k_proj",
+    "o_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

adapter_model.safetensors

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

global_step448/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt

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

global_step448/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt

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

global_step448/bf16_zero_pp_rank_2_mp_rank_00_optim_states.pt

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

global_step448/bf16_zero_pp_rank_3_mp_rank_00_optim_states.pt

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

global_step448/mp_rank_00_model_states.pt

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

latest

@@ -0,0 +1 @@
+global_step448

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

rng_state_3.pth

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

scheduler.pt

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

trainer_state.json

@@ -0,0 +1,3293 @@
+{
+  "best_metric": 1.203959345817566,
+  "best_model_checkpoint": "./out/checkpoint-112",
+  "epoch": 3.9881305637982196,
+  "eval_steps": 28,
+  "global_step": 448,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.01,
+      "grad_norm": 4.313233023002325,
+      "learning_rate": 1.8181818181818182e-05,
+      "loss": 1.9528,
+      "step": 1
+    },
+    {
+      "epoch": 0.01,
+      "eval_loss": 2.1875686645507812,
+      "eval_runtime": 12.8608,
+      "eval_samples_per_second": 23.327,
+      "eval_steps_per_second": 2.955,
+      "step": 1
+    },
+    {
+      "epoch": 0.02,
+      "grad_norm": 4.039172290955229,
+      "learning_rate": 3.6363636363636364e-05,
+      "loss": 1.8358,
+      "step": 2
+    },
+    {
+      "epoch": 0.03,
+      "grad_norm": 4.504705512003857,
+      "learning_rate": 5.4545454545454546e-05,
+      "loss": 2.0207,
+      "step": 3
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 4.591862504847867,
+      "learning_rate": 7.272727272727273e-05,
+      "loss": 1.979,
+      "step": 4
+    },
+    {
+      "epoch": 0.04,
+      "grad_norm": 3.812893581399005,
+      "learning_rate": 9.090909090909092e-05,
+      "loss": 1.8356,
+      "step": 5
+    },
+    {
+      "epoch": 0.05,
+      "grad_norm": 0.42886752872747064,
+      "learning_rate": 0.00010909090909090909,
+      "loss": 1.6722,
+      "step": 6
+    },
+    {
+      "epoch": 0.06,
+      "grad_norm": 0.22497294481851865,
+      "learning_rate": 0.00012727272727272728,
+      "loss": 1.6711,
+      "step": 7
+    },
+    {
+      "epoch": 0.07,
+      "grad_norm": 0.20955259847301927,
+      "learning_rate": 0.00014545454545454546,
+      "loss": 1.8546,
+      "step": 8
+    },
+    {
+      "epoch": 0.08,
+      "grad_norm": 0.2200095325539683,
+      "learning_rate": 0.00016363636363636366,
+      "loss": 1.7538,
+      "step": 9
+    },
+    {
+      "epoch": 0.09,
+      "grad_norm": 0.19187339879899318,
+      "learning_rate": 0.00018181818181818183,
+      "loss": 1.6137,
+      "step": 10
+    },
+    {
+      "epoch": 0.1,
+      "grad_norm": 0.2113395673717837,
+      "learning_rate": 0.0002,
+      "loss": 1.5225,
+      "step": 11
+    },
+    {
+      "epoch": 0.11,
+      "grad_norm": 0.17673768408382828,
+      "learning_rate": 0.00019999741592564903,
+      "loss": 1.5303,
+      "step": 12
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 0.24120852820548402,
+      "learning_rate": 0.00019998966383614488,
+      "loss": 1.5089,
+      "step": 13
+    },
+    {
+      "epoch": 0.12,
+      "grad_norm": 0.3089489160535682,
+      "learning_rate": 0.00019997674413212708,
+      "loss": 1.4525,
+      "step": 14
+    },
+    {
+      "epoch": 0.13,
+      "grad_norm": 0.2656143410731927,
+      "learning_rate": 0.00019995865748130516,
+      "loss": 1.4648,
+      "step": 15
+    },
+    {
+      "epoch": 0.14,
+      "grad_norm": 3.769410316227205,
+      "learning_rate": 0.0001999354048184241,
+      "loss": 1.3439,
+      "step": 16
+    },
+    {
+      "epoch": 0.15,
+      "grad_norm": 0.32102180658823753,
+      "learning_rate": 0.00019990698734521613,
+      "loss": 1.4644,
+      "step": 17
+    },
+    {
+      "epoch": 0.16,
+      "grad_norm": 0.22094428128919438,
+      "learning_rate": 0.0001998734065303385,
+      "loss": 1.1927,
+      "step": 18
+    },
+    {
+      "epoch": 0.17,
+      "grad_norm": 0.22344487218098863,
+      "learning_rate": 0.00019983466410929764,
+      "loss": 1.2916,
+      "step": 19
+    },
+    {
+      "epoch": 0.18,
+      "grad_norm": 0.25036262498479456,
+      "learning_rate": 0.0001997907620843595,
+      "loss": 1.2982,
+      "step": 20
+    },
+    {
+      "epoch": 0.19,
+      "grad_norm": 0.22671119151539426,
+      "learning_rate": 0.00019974170272444604,
+      "loss": 1.2146,
+      "step": 21
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 0.259249080403425,
+      "learning_rate": 0.00019968748856501788,
+      "loss": 1.2072,
+      "step": 22
+    },
+    {
+      "epoch": 0.2,
+      "grad_norm": 0.23538477651406017,
+      "learning_rate": 0.00019962812240794343,
+      "loss": 1.3281,
+      "step": 23
+    },
+    {
+      "epoch": 0.21,
+      "grad_norm": 0.2659115087625978,
+      "learning_rate": 0.000199563607321354,
+      "loss": 1.1396,
+      "step": 24
+    },
+    {
+      "epoch": 0.22,
+      "grad_norm": 0.23617264858854836,
+      "learning_rate": 0.0001994939466394851,
+      "loss": 1.1389,
+      "step": 25
+    },
+    {
+      "epoch": 0.23,
+      "grad_norm": 0.20514227454180176,
+      "learning_rate": 0.00019941914396250446,
+      "loss": 1.249,
+      "step": 26
+    },
+    {
+      "epoch": 0.24,
+      "grad_norm": 0.19660894225830144,
+      "learning_rate": 0.00019933920315632557,
+      "loss": 1.1776,
+      "step": 27
+    },
+    {
+      "epoch": 0.25,
+      "grad_norm": 0.2067663909729571,
+      "learning_rate": 0.00019925412835240826,
+      "loss": 1.1327,
+      "step": 28
+    },
+    {
+      "epoch": 0.25,
+      "eval_loss": 1.2991960048675537,
+      "eval_runtime": 13.153,
+      "eval_samples_per_second": 22.808,
+      "eval_steps_per_second": 2.889,
+      "step": 28
+    },
+    {
+      "epoch": 0.26,
+      "grad_norm": 0.1816588361901526,
+      "learning_rate": 0.0001991639239475448,
+      "loss": 1.1247,
+      "step": 29
+    },
+    {
+      "epoch": 0.27,
+      "grad_norm": 0.19626955153633807,
+      "learning_rate": 0.00019906859460363307,
+      "loss": 1.1212,
+      "step": 30
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 0.21084275590405852,
+      "learning_rate": 0.00019896814524743528,
+      "loss": 0.9927,
+      "step": 31
+    },
+    {
+      "epoch": 0.28,
+      "grad_norm": 0.16560054949456768,
+      "learning_rate": 0.0001988625810703235,
+      "loss": 1.1249,
+      "step": 32
+    },
+    {
+      "epoch": 0.29,
+      "grad_norm": 0.14950879528294536,
+      "learning_rate": 0.0001987519075280114,
+      "loss": 1.1401,
+      "step": 33
+    },
+    {
+      "epoch": 0.3,
+      "grad_norm": 0.1777966882651237,
+      "learning_rate": 0.00019863613034027224,
+      "loss": 1.0769,
+      "step": 34
+    },
+    {
+      "epoch": 0.31,
+      "grad_norm": 0.1480537272052743,
+      "learning_rate": 0.00019851525549064323,
+      "loss": 1.0686,
+      "step": 35
+    },
+    {
+      "epoch": 0.32,
+      "grad_norm": 0.16911906750319078,
+      "learning_rate": 0.00019838928922611632,
+      "loss": 1.0253,
+      "step": 36
+    },
+    {
+      "epoch": 0.33,
+      "grad_norm": 0.15987682972555176,
+      "learning_rate": 0.00019825823805681543,
+      "loss": 1.0609,
+      "step": 37
+    },
+    {
+      "epoch": 0.34,
+      "grad_norm": 0.15757332939676763,
+      "learning_rate": 0.0001981221087556598,
+      "loss": 1.1086,
+      "step": 38
+    },
+    {
+      "epoch": 0.35,
+      "grad_norm": 0.13201845744757537,
+      "learning_rate": 0.00019798090835801418,
+      "loss": 1.073,
+      "step": 39
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 0.12544508015984754,
+      "learning_rate": 0.00019783464416132506,
+      "loss": 1.0633,
+      "step": 40
+    },
+    {
+      "epoch": 0.36,
+      "grad_norm": 0.14645820383886451,
+      "learning_rate": 0.00019768332372474366,
+      "loss": 1.0653,
+      "step": 41
+    },
+    {
+      "epoch": 0.37,
+      "grad_norm": 0.14814101902137117,
+      "learning_rate": 0.00019752695486873517,
+      "loss": 1.0937,
+      "step": 42
+    },
+    {
+      "epoch": 0.38,
+      "grad_norm": 0.13888915595055443,
+      "learning_rate": 0.00019736554567467452,
+      "loss": 1.0462,
+      "step": 43
+    },
+    {
+      "epoch": 0.39,
+      "grad_norm": 0.13185349806639524,
+      "learning_rate": 0.00019719910448442893,
+      "loss": 1.2177,
+      "step": 44
+    },
+    {
+      "epoch": 0.4,
+      "grad_norm": 0.15271046712350847,
+      "learning_rate": 0.00019702763989992662,
+      "loss": 1.0237,
+      "step": 45
+    },
+    {
+      "epoch": 0.41,
+      "grad_norm": 0.17053588557430902,
+      "learning_rate": 0.00019685116078271223,
+      "loss": 1.0038,
+      "step": 46
+    },
+    {
+      "epoch": 0.42,
+      "grad_norm": 0.15641087356577812,
+      "learning_rate": 0.00019666967625348906,
+      "loss": 1.0886,
+      "step": 47
+    },
+    {
+      "epoch": 0.43,
+      "grad_norm": 0.1544028594191567,
+      "learning_rate": 0.00019648319569164736,
+      "loss": 1.1378,
+      "step": 48
+    },
+    {
+      "epoch": 0.44,
+      "grad_norm": 0.14794885994140625,
+      "learning_rate": 0.00019629172873477995,
+      "loss": 1.1495,
+      "step": 49
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 0.1577684884028266,
+      "learning_rate": 0.0001960952852781838,
+      "loss": 1.0782,
+      "step": 50
+    },
+    {
+      "epoch": 0.45,
+      "grad_norm": 0.15961044045091288,
+      "learning_rate": 0.0001958938754743489,
+      "loss": 1.0107,
+      "step": 51
+    },
+    {
+      "epoch": 0.46,
+      "grad_norm": 0.14486696586022083,
+      "learning_rate": 0.0001956875097324334,
+      "loss": 1.0494,
+      "step": 52
+    },
+    {
+      "epoch": 0.47,
+      "grad_norm": 0.14250413725518896,
+      "learning_rate": 0.00019547619871772574,
+      "loss": 1.039,
+      "step": 53
+    },
+    {
+      "epoch": 0.48,
+      "grad_norm": 0.1196720279125328,
+      "learning_rate": 0.00019525995335109334,
+      "loss": 1.0966,
+      "step": 54
+    },
+    {
+      "epoch": 0.49,
+      "grad_norm": 0.14984795891635327,
+      "learning_rate": 0.0001950387848084183,
+      "loss": 1.0874,
+      "step": 55
+    },
+    {
+      "epoch": 0.5,
+      "grad_norm": 0.14891088442480416,
+      "learning_rate": 0.00019481270452001987,
+      "loss": 1.097,
+      "step": 56
+    },
+    {
+      "epoch": 0.5,
+      "eval_loss": 1.2264304161071777,
+      "eval_runtime": 13.2279,
+      "eval_samples_per_second": 22.679,
+      "eval_steps_per_second": 2.873,
+      "step": 56
+    },
+    {
+      "epoch": 0.51,
+      "grad_norm": 0.17814266552244534,
+      "learning_rate": 0.00019458172417006347,
+      "loss": 1.1372,
+      "step": 57
+    },
+    {
+      "epoch": 0.52,
+      "grad_norm": 0.16125636132578247,
+      "learning_rate": 0.00019434585569595708,
+      "loss": 1.0623,
+      "step": 58
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 0.15203437202125702,
+      "learning_rate": 0.00019410511128773418,
+      "loss": 1.0399,
+      "step": 59
+    },
+    {
+      "epoch": 0.53,
+      "grad_norm": 0.1677461135605213,
+      "learning_rate": 0.0001938595033874238,
+      "loss": 1.0884,
+      "step": 60
+    },
+    {
+      "epoch": 0.54,
+      "grad_norm": 0.13564559875683407,
+      "learning_rate": 0.0001936090446884074,
+      "loss": 1.0176,
+      "step": 61
+    },
+    {
+      "epoch": 0.55,
+      "grad_norm": 0.1521886500642157,
+      "learning_rate": 0.00019335374813476302,
+      "loss": 1.0146,
+      "step": 62
+    },
+    {
+      "epoch": 0.56,
+      "grad_norm": 0.1410132122625916,
+      "learning_rate": 0.00019309362692059617,
+      "loss": 1.044,
+      "step": 63
+    },
+    {
+      "epoch": 0.57,
+      "grad_norm": 0.15237848179385577,
+      "learning_rate": 0.00019282869448935798,
+      "loss": 1.0354,
+      "step": 64
+    },
+    {
+      "epoch": 0.58,
+      "grad_norm": 0.13871660988504514,
+      "learning_rate": 0.00019255896453315052,
+      "loss": 1.0189,
+      "step": 65
+    },
+    {
+      "epoch": 0.59,
+      "grad_norm": 0.14863047478901453,
+      "learning_rate": 0.000192284450992019,
+      "loss": 1.0704,
+      "step": 66
+    },
+    {
+      "epoch": 0.6,
+      "grad_norm": 0.13794806124403974,
+      "learning_rate": 0.0001920051680532314,
+      "loss": 1.0996,
+      "step": 67
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 0.13030507705779365,
+      "learning_rate": 0.00019172113015054532,
+      "loss": 1.0015,
+      "step": 68
+    },
+    {
+      "epoch": 0.61,
+      "grad_norm": 0.15092494718902358,
+      "learning_rate": 0.0001914323519634619,
+      "loss": 1.0822,
+      "step": 69
+    },
+    {
+      "epoch": 0.62,
+      "grad_norm": 0.1350212989006066,
+      "learning_rate": 0.00019113884841646736,
+      "loss": 1.0197,
+      "step": 70
+    },
+    {
+      "epoch": 0.63,
+      "grad_norm": 0.18991168066586347,
+      "learning_rate": 0.00019084063467826137,
+      "loss": 1.046,
+      "step": 71
+    },
+    {
+      "epoch": 0.64,
+      "grad_norm": 0.14884381774710187,
+      "learning_rate": 0.00019053772616097337,
+      "loss": 1.0346,
+      "step": 72
+    },
+    {
+      "epoch": 0.65,
+      "grad_norm": 0.15579311209945296,
+      "learning_rate": 0.000190230138519366,
+      "loss": 1.0505,
+      "step": 73
+    },
+    {
+      "epoch": 0.66,
+      "grad_norm": 0.16015337150592127,
+      "learning_rate": 0.000189917887650026,
+      "loss": 1.0504,
+      "step": 74
+    },
+    {
+      "epoch": 0.67,
+      "grad_norm": 0.1443969321518926,
+      "learning_rate": 0.00018960098969054255,
+      "loss": 1.0755,
+      "step": 75
+    },
+    {
+      "epoch": 0.68,
+      "grad_norm": 0.15722162227095848,
+      "learning_rate": 0.00018927946101867347,
+      "loss": 1.0541,
+      "step": 76
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.17009697584926559,
+      "learning_rate": 0.0001889533182514986,
+      "loss": 1.0231,
+      "step": 77
+    },
+    {
+      "epoch": 0.69,
+      "grad_norm": 0.1256822726781221,
+      "learning_rate": 0.0001886225782445612,
+      "loss": 0.8814,
+      "step": 78
+    },
+    {
+      "epoch": 0.7,
+      "grad_norm": 0.14019958069756655,
+      "learning_rate": 0.00018828725809099655,
+      "loss": 1.0277,
+      "step": 79
+    },
+    {
+      "epoch": 0.71,
+      "grad_norm": 0.17159459150063183,
+      "learning_rate": 0.0001879473751206489,
+      "loss": 0.9495,
+      "step": 80
+    },
+    {
+      "epoch": 0.72,
+      "grad_norm": 0.146430011834186,
+      "learning_rate": 0.00018760294689917553,
+      "loss": 1.0598,
+      "step": 81
+    },
+    {
+      "epoch": 0.73,
+      "grad_norm": 0.16834256802992476,
+      "learning_rate": 0.00018725399122713912,
+      "loss": 1.0237,
+      "step": 82
+    },
+    {
+      "epoch": 0.74,
+      "grad_norm": 0.15663699267164208,
+      "learning_rate": 0.00018690052613908772,
+      "loss": 0.939,
+      "step": 83
+    },
+    {
+      "epoch": 0.75,
+      "grad_norm": 0.15655985150409854,
+      "learning_rate": 0.0001865425699026226,
+      "loss": 1.0302,
+      "step": 84
+    },
+    {
+      "epoch": 0.75,
+      "eval_loss": 1.2143030166625977,
+      "eval_runtime": 13.2387,
+      "eval_samples_per_second": 22.661,
+      "eval_steps_per_second": 2.87,
+      "step": 84
+    },
+    {
+      "epoch": 0.76,
+      "grad_norm": 0.15273470110260864,
+      "learning_rate": 0.00018618014101745442,
+      "loss": 1.0127,
+      "step": 85
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 0.1723243680259614,
+      "learning_rate": 0.0001858132582144469,
+      "loss": 0.9306,
+      "step": 86
+    },
+    {
+      "epoch": 0.77,
+      "grad_norm": 0.14747098547446996,
+      "learning_rate": 0.00018544194045464886,
+      "loss": 1.0073,
+      "step": 87
+    },
+    {
+      "epoch": 0.78,
+      "grad_norm": 0.17208333285514918,
+      "learning_rate": 0.00018506620692831428,
+      "loss": 1.0328,
+      "step": 88
+    },
+    {
+      "epoch": 0.79,
+      "grad_norm": 0.14918051024971962,
+      "learning_rate": 0.0001846860770539105,
+      "loss": 1.0022,
+      "step": 89
+    },
+    {
+      "epoch": 0.8,
+      "grad_norm": 0.156315164090714,
+      "learning_rate": 0.00018430157047711474,
+      "loss": 1.0293,
+      "step": 90
+    },
+    {
+      "epoch": 0.81,
+      "grad_norm": 0.2013424548288477,
+      "learning_rate": 0.00018391270706979862,
+      "loss": 0.9395,
+      "step": 91
+    },
+    {
+      "epoch": 0.82,
+      "grad_norm": 0.17909726353002614,
+      "learning_rate": 0.00018351950692900126,
+      "loss": 0.9756,
+      "step": 92
+    },
+    {
+      "epoch": 0.83,
+      "grad_norm": 0.16939245158726288,
+      "learning_rate": 0.00018312199037589068,
+      "loss": 0.9576,
+      "step": 93
+    },
+    {
+      "epoch": 0.84,
+      "grad_norm": 0.14685720680893694,
+      "learning_rate": 0.00018272017795471345,
+      "loss": 1.0045,
+      "step": 94
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.17464839085505987,
+      "learning_rate": 0.000182314090431733,
+      "loss": 0.9862,
+      "step": 95
+    },
+    {
+      "epoch": 0.85,
+      "grad_norm": 0.16060904136932572,
+      "learning_rate": 0.00018190374879415632,
+      "loss": 1.0022,
+      "step": 96
+    },
+    {
+      "epoch": 0.86,
+      "grad_norm": 0.18715193350083867,
+      "learning_rate": 0.00018148917424904953,
+      "loss": 1.042,
+      "step": 97
+    },
+    {
+      "epoch": 0.87,
+      "grad_norm": 0.1675573400576595,
+      "learning_rate": 0.0001810703882222415,
+      "loss": 1.0047,
+      "step": 98
+    },
+    {
+      "epoch": 0.88,
+      "grad_norm": 0.1871466286989249,
+      "learning_rate": 0.00018064741235721687,
+      "loss": 0.9834,
+      "step": 99
+    },
+    {
+      "epoch": 0.89,
+      "grad_norm": 0.17453934867565302,
+      "learning_rate": 0.00018022026851399737,
+      "loss": 0.9649,
+      "step": 100
+    },
+    {
+      "epoch": 0.9,
+      "grad_norm": 0.15960631507184767,
+      "learning_rate": 0.0001797889787680119,
+      "loss": 0.9673,
+      "step": 101
+    },
+    {
+      "epoch": 0.91,
+      "grad_norm": 0.17844936635366368,
+      "learning_rate": 0.00017935356540895597,
+      "loss": 1.0951,
+      "step": 102
+    },
+    {
+      "epoch": 0.92,
+      "grad_norm": 0.16733018789000254,
+      "learning_rate": 0.00017891405093963938,
+      "loss": 0.9954,
+      "step": 103
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.17305556075296993,
+      "learning_rate": 0.00017847045807482345,
+      "loss": 0.892,
+      "step": 104
+    },
+    {
+      "epoch": 0.93,
+      "grad_norm": 0.17197614099805034,
+      "learning_rate": 0.00017802280974004716,
+      "loss": 1.0494,
+      "step": 105
+    },
+    {
+      "epoch": 0.94,
+      "grad_norm": 0.18063836817127235,
+      "learning_rate": 0.000177571129070442,
+      "loss": 1.0264,
+      "step": 106
+    },
+    {
+      "epoch": 0.95,
+      "grad_norm": 0.14597707005699143,
+      "learning_rate": 0.00017711543940953668,
+      "loss": 0.9532,
+      "step": 107
+    },
+    {
+      "epoch": 0.96,
+      "grad_norm": 0.1422048149465345,
+      "learning_rate": 0.00017665576430805053,
+      "loss": 0.97,
+      "step": 108
+    },
+    {
+      "epoch": 0.97,
+      "grad_norm": 0.18313914688655572,
+      "learning_rate": 0.0001761921275226763,
+      "loss": 0.9282,
+      "step": 109
+    },
+    {
+      "epoch": 0.98,
+      "grad_norm": 0.200679751171441,
+      "learning_rate": 0.00017572455301485249,
+      "loss": 1.0,
+      "step": 110
+    },
+    {
+      "epoch": 0.99,
+      "grad_norm": 0.17700985594898055,
+      "learning_rate": 0.00017525306494952498,
+      "loss": 1.0165,
+      "step": 111
+    },
+    {
+      "epoch": 1.0,
+      "grad_norm": 0.19925777202726191,
+      "learning_rate": 0.0001747776876938981,
+      "loss": 1.0346,
+      "step": 112
+    },
+    {
+      "epoch": 1.0,
+      "eval_loss": 1.203959345817566,
+      "eval_runtime": 13.2547,
+      "eval_samples_per_second": 22.634,
+      "eval_steps_per_second": 2.867,
+      "step": 112
+    },
+    {
+      "epoch": 1.01,
+      "grad_norm": 0.1606469603473709,
+      "learning_rate": 0.00017429844581617532,
+      "loss": 0.9832,
+      "step": 113
+    },
+    {
+      "epoch": 1.01,
+      "grad_norm": 0.16403912763780054,
+      "learning_rate": 0.00017381536408428948,
+      "loss": 0.9346,
+      "step": 114
+    },
+    {
+      "epoch": 1.02,
+      "grad_norm": 0.1936046893744468,
+      "learning_rate": 0.00017332846746462288,
+      "loss": 0.9382,
+      "step": 115
+    },
+    {
+      "epoch": 1.03,
+      "grad_norm": 0.14250769247239573,
+      "learning_rate": 0.0001728377811207168,
+      "loss": 0.8914,
+      "step": 116
+    },
+    {
+      "epoch": 1.04,
+      "grad_norm": 0.17889563599797687,
+      "learning_rate": 0.00017234333041197126,
+      "loss": 0.9736,
+      "step": 117
+    },
+    {
+      "epoch": 1.05,
+      "grad_norm": 0.20288960866045594,
+      "learning_rate": 0.00017184514089233405,
+      "loss": 0.8477,
+      "step": 118
+    },
+    {
+      "epoch": 1.06,
+      "grad_norm": 0.20926349930533472,
+      "learning_rate": 0.00017134323830898037,
+      "loss": 0.9933,
+      "step": 119
+    },
+    {
+      "epoch": 1.07,
+      "grad_norm": 0.21316934416499642,
+      "learning_rate": 0.00017083764860098205,
+      "loss": 0.9168,
+      "step": 120
+    },
+    {
+      "epoch": 1.08,
+      "grad_norm": 0.21654320387312692,
+      "learning_rate": 0.0001703283978979671,
+      "loss": 0.9584,
+      "step": 121
+    },
+    {
+      "epoch": 1.09,
+      "grad_norm": 0.23789742308175463,
+      "learning_rate": 0.00016981551251876904,
+      "loss": 1.0298,
+      "step": 122
+    },
+    {
+      "epoch": 1.09,
+      "grad_norm": 0.16433271793469648,
+      "learning_rate": 0.00016929901897006698,
+      "loss": 0.8833,
+      "step": 123
+    },
+    {
+      "epoch": 1.1,
+      "grad_norm": 0.16908727866207868,
+      "learning_rate": 0.0001687789439450156,
+      "loss": 1.0675,
+      "step": 124
+    },
+    {
+      "epoch": 1.11,
+      "grad_norm": 0.1670067931363302,
+      "learning_rate": 0.00016825531432186543,
+      "loss": 0.9515,
+      "step": 125
+    },
+    {
+      "epoch": 1.12,
+      "grad_norm": 0.17777465531550865,
+      "learning_rate": 0.00016772815716257412,
+      "loss": 0.8929,
+      "step": 126
+    },
+    {
+      "epoch": 1.13,
+      "grad_norm": 0.18442783204919333,
+      "learning_rate": 0.00016719749971140754,
+      "loss": 0.8388,
+      "step": 127
+    },
+    {
+      "epoch": 1.14,
+      "grad_norm": 0.19073362304284272,
+      "learning_rate": 0.0001666633693935319,
+      "loss": 0.9584,
+      "step": 128
+    },
+    {
+      "epoch": 1.15,
+      "grad_norm": 0.20189563405135308,
+      "learning_rate": 0.00016612579381359622,
+      "loss": 1.0264,
+      "step": 129
+    },
+    {
+      "epoch": 1.16,
+      "grad_norm": 0.1694138210313381,
+      "learning_rate": 0.00016558480075430594,
+      "loss": 0.9592,
+      "step": 130
+    },
+    {
+      "epoch": 1.17,
+      "grad_norm": 0.19195382946787184,
+      "learning_rate": 0.00016504041817498678,
+      "loss": 0.974,
+      "step": 131
+    },
+    {
+      "epoch": 1.18,
+      "grad_norm": 0.20684215619155688,
+      "learning_rate": 0.00016449267421013994,
+      "loss": 0.8499,
+      "step": 132
+    },
+    {
+      "epoch": 1.18,
+      "grad_norm": 0.22003490429847744,
+      "learning_rate": 0.00016394159716798807,
+      "loss": 0.9659,
+      "step": 133
+    },
+    {
+      "epoch": 1.19,
+      "grad_norm": 0.21977918206745437,
+      "learning_rate": 0.00016338721552901212,
+      "loss": 0.9213,
+      "step": 134
+    },
+    {
+      "epoch": 1.2,
+      "grad_norm": 0.2076993903333204,
+      "learning_rate": 0.0001628295579444796,
+      "loss": 0.8119,
+      "step": 135
+    },
+    {
+      "epoch": 1.21,
+      "grad_norm": 0.2001771499954729,
+      "learning_rate": 0.0001622686532349637,
+      "loss": 0.9183,
+      "step": 136
+    },
+    {
+      "epoch": 1.22,
+      "grad_norm": 0.18671550149366203,
+      "learning_rate": 0.00016170453038885394,
+      "loss": 0.8836,
+      "step": 137
+    },
+    {
+      "epoch": 1.23,
+      "grad_norm": 0.20867427207572573,
+      "learning_rate": 0.0001611372185608578,
+      "loss": 0.9964,
+      "step": 138
+    },
+    {
+      "epoch": 1.24,
+      "grad_norm": 0.20035138443113176,
+      "learning_rate": 0.0001605667470704942,
+      "loss": 0.9209,
+      "step": 139
+    },
+    {
+      "epoch": 1.25,
+      "grad_norm": 0.22696612020505577,
+      "learning_rate": 0.0001599931454005781,
+      "loss": 1.0162,
+      "step": 140
+    },
+    {
+      "epoch": 1.25,
+      "eval_loss": 1.2188584804534912,
+      "eval_runtime": 13.249,
+      "eval_samples_per_second": 22.643,
+      "eval_steps_per_second": 2.868,
+      "step": 140
+    },
+    {
+      "epoch": 1.26,
+      "grad_norm": 0.21554353495018647,
+      "learning_rate": 0.00015941644319569665,
+      "loss": 1.0487,
+      "step": 141
+    },
+    {
+      "epoch": 1.26,
+      "grad_norm": 0.22894492131909072,
+      "learning_rate": 0.00015883667026067745,
+      "loss": 0.9352,
+      "step": 142
+    },
+    {
+      "epoch": 1.27,
+      "grad_norm": 0.19145184577172686,
+      "learning_rate": 0.00015825385655904788,
+      "loss": 0.8878,
+      "step": 143
+    },
+    {
+      "epoch": 1.28,
+      "grad_norm": 0.22544664152936575,
+      "learning_rate": 0.00015766803221148673,
+      "loss": 1.0,
+      "step": 144
+    },
+    {
+      "epoch": 1.29,
+      "grad_norm": 0.26000661355557114,
+      "learning_rate": 0.00015707922749426737,
+      "loss": 0.9339,
+      "step": 145
+    },
+    {
+      "epoch": 1.3,
+      "grad_norm": 0.24433845134512236,
+      "learning_rate": 0.00015648747283769317,
+      "loss": 0.9474,
+      "step": 146
+    },
+    {
+      "epoch": 1.31,
+      "grad_norm": 0.21973931169609887,
+      "learning_rate": 0.00015589279882452476,
+      "loss": 0.9357,
+      "step": 147
+    },
+    {
+      "epoch": 1.32,
+      "grad_norm": 0.23929008733305812,
+      "learning_rate": 0.0001552952361883994,
+      "loss": 0.9985,
+      "step": 148
+    },
+    {
+      "epoch": 1.33,
+      "grad_norm": 0.23431856747573573,
+      "learning_rate": 0.00015469481581224272,
+      "loss": 0.8913,
+      "step": 149
+    },
+    {
+      "epoch": 1.34,
+      "grad_norm": 0.2233543327912565,
+      "learning_rate": 0.00015409156872667258,
+      "loss": 0.9877,
+      "step": 150
+    },
+    {
+      "epoch": 1.34,
+      "grad_norm": 0.21281207674183256,
+      "learning_rate": 0.0001534855261083954,
+      "loss": 0.9071,
+      "step": 151
+    },
+    {
+      "epoch": 1.35,
+      "grad_norm": 0.20314832700152685,
+      "learning_rate": 0.00015287671927859494,
+      "loss": 0.9373,
+      "step": 152
+    },
+    {
+      "epoch": 1.36,
+      "grad_norm": 0.19648565819019825,
+      "learning_rate": 0.00015226517970131343,
+      "loss": 0.9469,
+      "step": 153
+    },
+    {
+      "epoch": 1.37,
+      "grad_norm": 0.2262428264639853,
+      "learning_rate": 0.00015165093898182562,
+      "loss": 1.0066,
+      "step": 154
+    },
+    {
+      "epoch": 1.38,
+      "grad_norm": 0.22253433035020442,
+      "learning_rate": 0.00015103402886500525,
+      "loss": 0.8875,
+      "step": 155
+    },
+    {
+      "epoch": 1.39,
+      "grad_norm": 0.181161648904613,
+      "learning_rate": 0.00015041448123368455,
+      "loss": 0.9004,
+      "step": 156
+    },
+    {
+      "epoch": 1.4,
+      "grad_norm": 0.20968483802367816,
+      "learning_rate": 0.00014979232810700637,
+      "loss": 0.9133,
+      "step": 157
+    },
+    {
+      "epoch": 1.41,
+      "grad_norm": 0.20540509271288435,
+      "learning_rate": 0.0001491676016387694,
+      "loss": 0.8876,
+      "step": 158
+    },
+    {
+      "epoch": 1.42,
+      "grad_norm": 0.18762795731312454,
+      "learning_rate": 0.00014854033411576659,
+      "loss": 0.933,
+      "step": 159
+    },
+    {
+      "epoch": 1.42,
+      "grad_norm": 0.23223345997338857,
+      "learning_rate": 0.00014791055795611624,
+      "loss": 0.9182,
+      "step": 160
+    },
+    {
+      "epoch": 1.43,
+      "grad_norm": 0.21932384461027146,
+      "learning_rate": 0.00014727830570758678,
+      "loss": 0.9514,
+      "step": 161
+    },
+    {
+      "epoch": 1.44,
+      "grad_norm": 0.21819663730951108,
+      "learning_rate": 0.0001466436100459146,
+      "loss": 0.9162,
+      "step": 162
+    },
+    {
+      "epoch": 1.45,
+      "grad_norm": 0.2325813323476676,
+      "learning_rate": 0.00014600650377311522,
+      "loss": 0.9308,
+      "step": 163
+    },
+    {
+      "epoch": 1.46,
+      "grad_norm": 0.2568337182939043,
+      "learning_rate": 0.0001453670198157883,
+      "loss": 0.9995,
+      "step": 164
+    },
+    {
+      "epoch": 1.47,
+      "grad_norm": 0.22578454460723413,
+      "learning_rate": 0.00014472519122341566,
+      "loss": 0.9052,
+      "step": 165
+    },
+    {
+      "epoch": 1.48,
+      "grad_norm": 0.23564258958796755,
+      "learning_rate": 0.00014408105116665336,
+      "loss": 0.9714,
+      "step": 166
+    },
+    {
+      "epoch": 1.49,
+      "grad_norm": 0.24266133562839415,
+      "learning_rate": 0.00014343463293561734,
+      "loss": 0.9219,
+      "step": 167
+    },
+    {
+      "epoch": 1.5,
+      "grad_norm": 0.23472454708184465,
+      "learning_rate": 0.00014278596993816308,
+      "loss": 0.8762,
+      "step": 168
+    },
+    {
+      "epoch": 1.5,
+      "eval_loss": 1.2197421789169312,
+      "eval_runtime": 13.2616,
+      "eval_samples_per_second": 22.622,
+      "eval_steps_per_second": 2.865,
+      "step": 168
+    },
+    {
+      "epoch": 1.5,
+      "grad_norm": 0.23623633375452713,
+      "learning_rate": 0.00014213509569815884,
+      "loss": 0.8809,
+      "step": 169
+    },
+    {
+      "epoch": 1.51,
+      "grad_norm": 0.25344275204523486,
+      "learning_rate": 0.00014148204385375321,
+      "loss": 0.7972,
+      "step": 170
+    },
+    {
+      "epoch": 1.52,
+      "grad_norm": 0.23111396119549557,
+      "learning_rate": 0.0001408268481556366,
+      "loss": 0.8228,
+      "step": 171
+    },
+    {
+      "epoch": 1.53,
+      "grad_norm": 0.2510618369255398,
+      "learning_rate": 0.00014016954246529696,
+      "loss": 0.8849,
+      "step": 172
+    },
+    {
+      "epoch": 1.54,
+      "grad_norm": 0.2764366116622668,
+      "learning_rate": 0.0001395101607532698,
+      "loss": 0.8936,
+      "step": 173
+    },
+    {
+      "epoch": 1.55,
+      "grad_norm": 0.24325811719582827,
+      "learning_rate": 0.00013884873709738257,
+      "loss": 0.8602,
+      "step": 174
+    },
+    {
+      "epoch": 1.56,
+      "grad_norm": 0.213781513838486,
+      "learning_rate": 0.00013818530568099327,
+      "loss": 0.9492,
+      "step": 175
+    },
+    {
+      "epoch": 1.57,
+      "grad_norm": 0.2397396374239057,
+      "learning_rate": 0.00013751990079122412,
+      "loss": 1.0499,
+      "step": 176
+    },
+    {
+      "epoch": 1.58,
+      "grad_norm": 0.21579907170368723,
+      "learning_rate": 0.00013685255681718922,
+      "loss": 0.9438,
+      "step": 177
+    },
+    {
+      "epoch": 1.58,
+      "grad_norm": 0.2359312681928786,
+      "learning_rate": 0.0001361833082482175,
+      "loss": 0.9289,
+      "step": 178
+    },
+    {
+      "epoch": 1.59,
+      "grad_norm": 0.2618189093396496,
+      "learning_rate": 0.0001355121896720703,
+      "loss": 0.981,
+      "step": 179
+    },
+    {
+      "epoch": 1.6,
+      "grad_norm": 0.20876513773174135,
+      "learning_rate": 0.00013483923577315348,
+      "loss": 0.82,
+      "step": 180
+    },
+    {
+      "epoch": 1.61,
+      "grad_norm": 0.22162748553995645,
+      "learning_rate": 0.00013416448133072526,
+      "loss": 1.0131,
+      "step": 181
+    },
+    {
+      "epoch": 1.62,
+      "grad_norm": 0.20975549982451164,
+      "learning_rate": 0.00013348796121709862,
+      "loss": 0.8763,
+      "step": 182
+    },
+    {
+      "epoch": 1.63,
+      "grad_norm": 0.22840397707525473,
+      "learning_rate": 0.00013280971039583906,
+      "loss": 0.949,
+      "step": 183
+    },
+    {
+      "epoch": 1.64,
+      "grad_norm": 0.23384636230161737,
+      "learning_rate": 0.0001321297639199575,
+      "loss": 0.9567,
+      "step": 184
+    },
+    {
+      "epoch": 1.65,
+      "grad_norm": 0.22905979409902957,
+      "learning_rate": 0.000131448156930099,
+      "loss": 0.9153,
+      "step": 185
+    },
+    {
+      "epoch": 1.66,
+      "grad_norm": 0.27620894683694563,
+      "learning_rate": 0.0001307649246527263,
+      "loss": 0.8246,
+      "step": 186
+    },
+    {
+      "epoch": 1.66,
+      "grad_norm": 0.23004170633106227,
+      "learning_rate": 0.0001300801023982995,
+      "loss": 1.0181,
+      "step": 187
+    },
+    {
+      "epoch": 1.67,
+      "grad_norm": 0.2219849136264378,
+      "learning_rate": 0.00012939372555945112,
+      "loss": 0.9535,
+      "step": 188
+    },
+    {
+      "epoch": 1.68,
+      "grad_norm": 0.24458750452490116,
+      "learning_rate": 0.0001287058296091567,
+      "loss": 0.8968,
+      "step": 189
+    },
+    {
+      "epoch": 1.69,
+      "grad_norm": 0.2564337740159555,
+      "learning_rate": 0.00012801645009890195,
+      "loss": 0.7955,
+      "step": 190
+    },
+    {
+      "epoch": 1.7,
+      "grad_norm": 0.24100850371438767,
+      "learning_rate": 0.0001273256226568451,
+      "loss": 0.9235,
+      "step": 191
+    },
+    {
+      "epoch": 1.71,
+      "grad_norm": 0.24757089527873732,
+      "learning_rate": 0.00012663338298597563,
+      "loss": 1.007,
+      "step": 192
+    },
+    {
+      "epoch": 1.72,
+      "grad_norm": 0.24701038583742888,
+      "learning_rate": 0.00012593976686226904,
+      "loss": 0.9885,
+      "step": 193
+    },
+    {
+      "epoch": 1.73,
+      "grad_norm": 0.26373721125634964,
+      "learning_rate": 0.0001252448101328381,
+      "loss": 0.8785,
+      "step": 194
+    },
+    {
+      "epoch": 1.74,
+      "grad_norm": 0.2227761464470136,
+      "learning_rate": 0.00012454854871407994,
+      "loss": 0.8806,
+      "step": 195
+    },
+    {
+      "epoch": 1.74,
+      "grad_norm": 0.2283950634350429,
+      "learning_rate": 0.00012385101858982005,
+      "loss": 0.9053,
+      "step": 196
+    },
+    {
+      "epoch": 1.74,
+      "eval_loss": 1.2198154926300049,
+      "eval_runtime": 13.2208,
+      "eval_samples_per_second": 22.692,
+      "eval_steps_per_second": 2.874,
+      "step": 196
+    },
+    {
+      "epoch": 1.75,
+      "grad_norm": 0.23406423788354982,
+      "learning_rate": 0.00012315225580945252,
+      "loss": 0.9397,
+      "step": 197
+    },
+    {
+      "epoch": 1.76,
+      "grad_norm": 0.23807045727443327,
+      "learning_rate": 0.0001224522964860769,
+      "loss": 0.9712,
+      "step": 198
+    },
+    {
+      "epoch": 1.77,
+      "grad_norm": 0.2463614808838948,
+      "learning_rate": 0.00012175117679463187,
+      "loss": 0.8558,
+      "step": 199
+    },
+    {
+      "epoch": 1.78,
+      "grad_norm": 0.24737417059302014,
+      "learning_rate": 0.00012104893297002567,
+      "loss": 0.9723,
+      "step": 200
+    },
+    {
+      "epoch": 1.79,
+      "grad_norm": 0.243750688050595,
+      "learning_rate": 0.0001203456013052634,
+      "loss": 0.964,
+      "step": 201
+    },
+    {
+      "epoch": 1.8,
+      "grad_norm": 0.24572059557106538,
+      "learning_rate": 0.00011964121814957137,
+      "loss": 0.9109,
+      "step": 202
+    },
+    {
+      "epoch": 1.81,
+      "grad_norm": 0.24044117903962453,
+      "learning_rate": 0.00011893581990651848,
+      "loss": 1.0019,
+      "step": 203
+    },
+    {
+      "epoch": 1.82,
+      "grad_norm": 0.2737568489071465,
+      "learning_rate": 0.00011822944303213486,
+      "loss": 0.8893,
+      "step": 204
+    },
+    {
+      "epoch": 1.82,
+      "grad_norm": 0.24122455882790084,
+      "learning_rate": 0.00011752212403302784,
+      "loss": 0.9162,
+      "step": 205
+    },
+    {
+      "epoch": 1.83,
+      "grad_norm": 0.28991871401626856,
+      "learning_rate": 0.00011681389946449504,
+      "loss": 0.8555,
+      "step": 206
+    },
+    {
+      "epoch": 1.84,
+      "grad_norm": 0.23767408810646548,
+      "learning_rate": 0.00011610480592863531,
+      "loss": 0.9936,
+      "step": 207
+    },
+    {
+      "epoch": 1.85,
+      "grad_norm": 0.22614733706173062,
+      "learning_rate": 0.00011539488007245702,
+      "loss": 0.916,
+      "step": 208
+    },
+    {
+      "epoch": 1.86,
+      "grad_norm": 0.22471992425846515,
+      "learning_rate": 0.00011468415858598411,
+      "loss": 0.8872,
+      "step": 209
+    },
+    {
+      "epoch": 1.87,
+      "grad_norm": 0.22675717145909688,
+      "learning_rate": 0.00011397267820035986,
+      "loss": 0.8393,
+      "step": 210
+    },
+    {
+      "epoch": 1.88,
+      "grad_norm": 0.2727459336483823,
+      "learning_rate": 0.00011326047568594851,
+      "loss": 0.8265,
+      "step": 211
+    },
+    {
+      "epoch": 1.89,
+      "grad_norm": 0.25216778031670767,
+      "learning_rate": 0.00011254758785043515,
+      "loss": 0.9939,
+      "step": 212
+    },
+    {
+      "epoch": 1.9,
+      "grad_norm": 0.269147378424304,
+      "learning_rate": 0.0001118340515369232,
+      "loss": 0.9102,
+      "step": 213
+    },
+    {
+      "epoch": 1.91,
+      "grad_norm": 0.2216178370833471,
+      "learning_rate": 0.00011111990362203033,
+      "loss": 0.8778,
+      "step": 214
+    },
+    {
+      "epoch": 1.91,
+      "grad_norm": 0.2602474934716497,
+      "learning_rate": 0.00011040518101398276,
+      "loss": 0.9454,
+      "step": 215
+    },
+    {
+      "epoch": 1.92,
+      "grad_norm": 0.2658635078442998,
+      "learning_rate": 0.00010968992065070769,
+      "loss": 0.8098,
+      "step": 216
+    },
+    {
+      "epoch": 1.93,
+      "grad_norm": 0.20997905209488962,
+      "learning_rate": 0.00010897415949792427,
+      "loss": 0.9318,
+      "step": 217
+    },
+    {
+      "epoch": 1.94,
+      "grad_norm": 0.24752453752221557,
+      "learning_rate": 0.00010825793454723325,
+      "loss": 0.949,
+      "step": 218
+    },
+    {
+      "epoch": 1.95,
+      "grad_norm": 0.255579569750529,
+      "learning_rate": 0.0001075412828142051,
+      "loss": 0.915,
+      "step": 219
+    },
+    {
+      "epoch": 1.96,
+      "grad_norm": 0.23186981930561867,
+      "learning_rate": 0.0001068242413364671,
+      "loss": 0.9132,
+      "step": 220
+    },
+    {
+      "epoch": 1.97,
+      "grad_norm": 0.35685140391438824,
+      "learning_rate": 0.00010610684717178905,
+      "loss": 0.9398,
+      "step": 221
+    },
+    {
+      "epoch": 1.98,
+      "grad_norm": 0.27320389987223703,
+      "learning_rate": 0.00010538913739616816,
+      "loss": 0.857,
+      "step": 222
+    },
+    {
+      "epoch": 1.99,
+      "grad_norm": 0.2324276771141761,
+      "learning_rate": 0.00010467114910191289,
+      "loss": 0.8546,
+      "step": 223
+    },
+    {
+      "epoch": 1.99,
+      "grad_norm": 0.22820341349854167,
+      "learning_rate": 0.00010395291939572593,
+      "loss": 0.9301,
+      "step": 224
+    },
+    {
+      "epoch": 1.99,
+      "eval_loss": 1.2246263027191162,
+      "eval_runtime": 13.1981,
+      "eval_samples_per_second": 22.731,
+      "eval_steps_per_second": 2.879,
+      "step": 224
+    },
+    {
+      "epoch": 2.0,
+      "grad_norm": 0.2289800489154315,
+      "learning_rate": 0.00010323448539678653,
+      "loss": 0.9922,
+      "step": 225
+    },
+    {
+      "epoch": 2.01,
+      "grad_norm": 0.2673353778680862,
+      "learning_rate": 0.00010251588423483205,
+      "loss": 0.7779,
+      "step": 226
+    },
+    {
+      "epoch": 2.02,
+      "grad_norm": 0.2420933678952559,
+      "learning_rate": 0.0001017971530482392,
+      "loss": 0.8044,
+      "step": 227
+    },
+    {
+      "epoch": 2.03,
+      "grad_norm": 0.21799264660625498,
+      "learning_rate": 0.00010107832898210439,
+      "loss": 0.8773,
+      "step": 228
+    },
+    {
+      "epoch": 2.04,
+      "grad_norm": 0.21443255695871016,
+      "learning_rate": 0.00010035944918632429,
+      "loss": 0.9031,
+      "step": 229
+    },
+    {
+      "epoch": 2.05,
+      "grad_norm": 0.23983734165788242,
+      "learning_rate": 9.96405508136757e-05,
+      "loss": 0.9014,
+      "step": 230
+    },
+    {
+      "epoch": 2.06,
+      "grad_norm": 0.27915481475799336,
+      "learning_rate": 9.892167101789564e-05,
+      "loss": 0.8853,
+      "step": 231
+    },
+    {
+      "epoch": 2.07,
+      "grad_norm": 0.2688949371564916,
+      "learning_rate": 9.820284695176082e-05,
+      "loss": 0.8452,
+      "step": 232
+    },
+    {
+      "epoch": 2.07,
+      "grad_norm": 0.2623278518867105,
+      "learning_rate": 9.748411576516794e-05,
+      "loss": 0.8612,
+      "step": 233
+    },
+    {
+      "epoch": 2.08,
+      "grad_norm": 0.2710502639103885,
+      "learning_rate": 9.676551460321349e-05,
+      "loss": 0.8108,
+      "step": 234
+    },
+    {
+      "epoch": 2.09,
+      "grad_norm": 0.282572880285737,
+      "learning_rate": 9.60470806042741e-05,
+      "loss": 0.7866,
+      "step": 235
+    },
+    {
+      "epoch": 2.1,
+      "grad_norm": 0.2829396962922612,
+      "learning_rate": 9.532885089808713e-05,
+      "loss": 0.8557,
+      "step": 236
+    },
+    {
+      "epoch": 2.11,
+      "grad_norm": 0.2721172338857335,
+      "learning_rate": 9.461086260383187e-05,
+      "loss": 0.7933,
+      "step": 237
+    },
+    {
+      "epoch": 2.12,
+      "grad_norm": 0.29736638811364446,
+      "learning_rate": 9.389315282821097e-05,
+      "loss": 0.7674,
+      "step": 238
+    },
+    {
+      "epoch": 2.13,
+      "grad_norm": 0.28571679920981263,
+      "learning_rate": 9.317575866353292e-05,
+      "loss": 0.7442,
+      "step": 239
+    },
+    {
+      "epoch": 2.14,
+      "grad_norm": 0.264545167150173,
+      "learning_rate": 9.245871718579491e-05,
+      "loss": 0.8505,
+      "step": 240
+    },
+    {
+      "epoch": 2.15,
+      "grad_norm": 0.30691085134027757,
+      "learning_rate": 9.174206545276677e-05,
+      "loss": 0.7898,
+      "step": 241
+    },
+    {
+      "epoch": 2.15,
+      "grad_norm": 0.31375028121981235,
+      "learning_rate": 9.102584050207578e-05,
+      "loss": 0.7661,
+      "step": 242
+    },
+    {
+      "epoch": 2.16,
+      "grad_norm": 0.28421530221837016,
+      "learning_rate": 9.031007934929236e-05,
+      "loss": 0.8328,
+      "step": 243
+    },
+    {
+      "epoch": 2.17,
+      "grad_norm": 0.25601367811173414,
+      "learning_rate": 8.959481898601728e-05,
+      "loss": 0.8281,
+      "step": 244
+    },
+    {
+      "epoch": 2.18,
+      "grad_norm": 0.2983724947729522,
+      "learning_rate": 8.888009637796968e-05,
+      "loss": 0.8567,
+      "step": 245
+    },
+    {
+      "epoch": 2.19,
+      "grad_norm": 0.2545616786933236,
+      "learning_rate": 8.81659484630768e-05,
+      "loss": 0.9151,
+      "step": 246
+    },
+    {
+      "epoch": 2.2,
+      "grad_norm": 0.23873712362647942,
+      "learning_rate": 8.745241214956483e-05,
+      "loss": 0.8818,
+      "step": 247
+    },
+    {
+      "epoch": 2.21,
+      "grad_norm": 0.285331972404065,
+      "learning_rate": 8.673952431405148e-05,
+      "loss": 0.7983,
+      "step": 248
+    },
+    {
+      "epoch": 2.22,
+      "grad_norm": 0.23897707291689843,
+      "learning_rate": 8.602732179964017e-05,
+      "loss": 0.8758,
+      "step": 249
+    },
+    {
+      "epoch": 2.23,
+      "grad_norm": 0.2830966091447457,
+      "learning_rate": 8.531584141401591e-05,
+      "loss": 0.8714,
+      "step": 250
+    },
+    {
+      "epoch": 2.23,
+      "grad_norm": 0.28872599217076506,
+      "learning_rate": 8.4605119927543e-05,
+      "loss": 0.8387,
+      "step": 251
+    },
+    {
+      "epoch": 2.24,
+      "grad_norm": 0.2652236346400331,
+      "learning_rate": 8.38951940713647e-05,
+      "loss": 0.8232,
+      "step": 252
+    },
+    {
+      "epoch": 2.24,
+      "eval_loss": 1.2432794570922852,
+      "eval_runtime": 13.2405,
+      "eval_samples_per_second": 22.658,
+      "eval_steps_per_second": 2.87,
+      "step": 252
+    },
+    {
+      "epoch": 2.25,
+      "grad_norm": 0.299978013524394,
+      "learning_rate": 8.318610053550497e-05,
+      "loss": 0.7321,
+      "step": 253
+    },
+    {
+      "epoch": 2.26,
+      "grad_norm": 0.2740002835117391,
+      "learning_rate": 8.247787596697218e-05,
+      "loss": 0.7605,
+      "step": 254
+    },
+    {
+      "epoch": 2.27,
+      "grad_norm": 0.2848366030132808,
+      "learning_rate": 8.177055696786516e-05,
+      "loss": 0.8485,
+      "step": 255
+    },
+    {
+      "epoch": 2.28,
+      "grad_norm": 0.24847418856075218,
+      "learning_rate": 8.106418009348157e-05,
+      "loss": 0.7557,
+      "step": 256
+    },
+    {
+      "epoch": 2.29,
+      "grad_norm": 0.33515508602624905,
+      "learning_rate": 8.035878185042868e-05,
+      "loss": 0.8015,
+      "step": 257
+    },
+    {
+      "epoch": 2.3,
+      "grad_norm": 0.2905943721096322,
+      "learning_rate": 7.965439869473664e-05,
+      "loss": 0.8457,
+      "step": 258
+    },
+    {
+      "epoch": 2.31,
+      "grad_norm": 0.3140679719552616,
+      "learning_rate": 7.895106702997437e-05,
+      "loss": 0.8559,
+      "step": 259
+    },
+    {
+      "epoch": 2.31,
+      "grad_norm": 0.29745105018138573,
+      "learning_rate": 7.824882320536814e-05,
+      "loss": 0.7453,
+      "step": 260
+    },
+    {
+      "epoch": 2.32,
+      "grad_norm": 0.29818631731197365,
+      "learning_rate": 7.754770351392311e-05,
+      "loss": 0.8354,
+      "step": 261
+    },
+    {
+      "epoch": 2.33,
+      "grad_norm": 0.24721488944366407,
+      "learning_rate": 7.684774419054747e-05,
+      "loss": 0.7755,
+      "step": 262
+    },
+    {
+      "epoch": 2.34,
+      "grad_norm": 0.31210442779019465,
+      "learning_rate": 7.614898141017996e-05,
+      "loss": 0.7208,
+      "step": 263
+    },
+    {
+      "epoch": 2.35,
+      "grad_norm": 0.2873220240109992,
+      "learning_rate": 7.54514512859201e-05,
+      "loss": 0.7548,
+      "step": 264
+    },
+    {
+      "epoch": 2.36,
+      "grad_norm": 0.3006634171776217,
+      "learning_rate": 7.475518986716194e-05,
+      "loss": 0.7566,
+      "step": 265
+    },
+    {
+      "epoch": 2.37,
+      "grad_norm": 0.2799417613336026,
+      "learning_rate": 7.406023313773097e-05,
+      "loss": 0.727,
+      "step": 266
+    },
+    {
+      "epoch": 2.38,
+      "grad_norm": 0.2451761866231664,
+      "learning_rate": 7.336661701402439e-05,
+      "loss": 0.9641,
+      "step": 267
+    },
+    {
+      "epoch": 2.39,
+      "grad_norm": 0.305202611125298,
+      "learning_rate": 7.267437734315492e-05,
+      "loss": 0.7891,
+      "step": 268
+    },
+    {
+      "epoch": 2.39,
+      "grad_norm": 0.29107717848747816,
+      "learning_rate": 7.198354990109805e-05,
+      "loss": 0.9032,
+      "step": 269
+    },
+    {
+      "epoch": 2.4,
+      "grad_norm": 0.2688898665176787,
+      "learning_rate": 7.129417039084333e-05,
+      "loss": 0.8416,
+      "step": 270
+    },
+    {
+      "epoch": 2.41,
+      "grad_norm": 0.2814206029778395,
+      "learning_rate": 7.060627444054893e-05,
+      "loss": 0.8443,
+      "step": 271
+    },
+    {
+      "epoch": 2.42,
+      "grad_norm": 0.2862094867555512,
+      "learning_rate": 6.99198976017005e-05,
+      "loss": 0.8271,
+      "step": 272
+    },
+    {
+      "epoch": 2.43,
+      "grad_norm": 0.3214647340394826,
+      "learning_rate": 6.923507534727373e-05,
+      "loss": 0.7793,
+      "step": 273
+    },
+    {
+      "epoch": 2.44,
+      "grad_norm": 0.3033659714564417,
+      "learning_rate": 6.855184306990106e-05,
+      "loss": 0.7856,
+      "step": 274
+    },
+    {
+      "epoch": 2.45,
+      "grad_norm": 0.3024382342577774,
+      "learning_rate": 6.78702360800425e-05,
+      "loss": 0.8633,
+      "step": 275
+    },
+    {
+      "epoch": 2.46,
+      "grad_norm": 0.25803598196729505,
+      "learning_rate": 6.719028960416098e-05,
+      "loss": 0.8428,
+      "step": 276
+    },
+    {
+      "epoch": 2.47,
+      "grad_norm": 0.35469202971401803,
+      "learning_rate": 6.651203878290139e-05,
+      "loss": 0.8665,
+      "step": 277
+    },
+    {
+      "epoch": 2.47,
+      "grad_norm": 0.3122516837597691,
+      "learning_rate": 6.583551866927475e-05,
+      "loss": 0.8787,
+      "step": 278
+    },
+    {
+      "epoch": 2.48,
+      "grad_norm": 0.3305470786367901,
+      "learning_rate": 6.516076422684654e-05,
+      "loss": 0.8765,
+      "step": 279
+    },
+    {
+      "epoch": 2.49,
+      "grad_norm": 0.3324622666488467,
+      "learning_rate": 6.448781032792972e-05,
+      "loss": 0.8318,
+      "step": 280
+    },
+    {
+      "epoch": 2.49,
+      "eval_loss": 1.2546111345291138,
+      "eval_runtime": 13.2379,
+      "eval_samples_per_second": 22.662,
+      "eval_steps_per_second": 2.871,
+      "step": 280
+    },
+    {
+      "epoch": 2.5,
+      "grad_norm": 0.342341713579355,
+      "learning_rate": 6.381669175178248e-05,
+      "loss": 0.9517,
+      "step": 281
+    },
+    {
+      "epoch": 2.51,
+      "grad_norm": 0.33913458352374665,
+      "learning_rate": 6.31474431828108e-05,
+      "loss": 0.8564,
+      "step": 282
+    },
+    {
+      "epoch": 2.52,
+      "grad_norm": 0.30528689383480295,
+      "learning_rate": 6.248009920877592e-05,
+      "loss": 0.8199,
+      "step": 283
+    },
+    {
+      "epoch": 2.53,
+      "grad_norm": 0.29698648367254743,
+      "learning_rate": 6.181469431900672e-05,
+      "loss": 0.785,
+      "step": 284
+    },
+    {
+      "epoch": 2.54,
+      "grad_norm": 0.32239262939282626,
+      "learning_rate": 6.115126290261745e-05,
+      "loss": 0.7794,
+      "step": 285
+    },
+    {
+      "epoch": 2.55,
+      "grad_norm": 0.2694595905080167,
+      "learning_rate": 6.048983924673022e-05,
+      "loss": 0.8056,
+      "step": 286
+    },
+    {
+      "epoch": 2.55,
+      "grad_norm": 0.3045496751154443,
+      "learning_rate": 5.983045753470308e-05,
+      "loss": 0.8164,
+      "step": 287
+    },
+    {
+      "epoch": 2.56,
+      "grad_norm": 0.2927868214627918,
+      "learning_rate": 5.917315184436345e-05,
+      "loss": 0.8358,
+      "step": 288
+    },
+    {
+      "epoch": 2.57,
+      "grad_norm": 0.2931914055644858,
+      "learning_rate": 5.851795614624682e-05,
+      "loss": 0.8011,
+      "step": 289
+    },
+    {
+      "epoch": 2.58,
+      "grad_norm": 0.3158716819379082,
+      "learning_rate": 5.786490430184115e-05,
+      "loss": 0.8332,
+      "step": 290
+    },
+    {
+      "epoch": 2.59,
+      "grad_norm": 0.3482519147352008,
+      "learning_rate": 5.72140300618369e-05,
+      "loss": 0.7621,
+      "step": 291
+    },
+    {
+      "epoch": 2.6,
+      "grad_norm": 0.28652801822050894,
+      "learning_rate": 5.656536706438267e-05,
+      "loss": 0.77,
+      "step": 292
+    },
+    {
+      "epoch": 2.61,
+      "grad_norm": 0.29691290613407717,
+      "learning_rate": 5.591894883334667e-05,
+      "loss": 0.9394,
+      "step": 293
+    },
+    {
+      "epoch": 2.62,
+      "grad_norm": 0.26699581966985203,
+      "learning_rate": 5.5274808776584367e-05,
+      "loss": 0.7918,
+      "step": 294
+    },
+    {
+      "epoch": 2.63,
+      "grad_norm": 0.2926923719762685,
+      "learning_rate": 5.463298018421171e-05,
+      "loss": 0.8723,
+      "step": 295
+    },
+    {
+      "epoch": 2.64,
+      "grad_norm": 0.3403087263187063,
+      "learning_rate": 5.399349622688479e-05,
+      "loss": 0.8097,
+      "step": 296
+    },
+    {
+      "epoch": 2.64,
+      "grad_norm": 0.34261233464532476,
+      "learning_rate": 5.335638995408545e-05,
+      "loss": 0.9032,
+      "step": 297
+    },
+    {
+      "epoch": 2.65,
+      "grad_norm": 0.31315234759634086,
+      "learning_rate": 5.272169429241325e-05,
+      "loss": 0.82,
+      "step": 298
+    },
+    {
+      "epoch": 2.66,
+      "grad_norm": 0.3179759425444047,
+      "learning_rate": 5.208944204388377e-05,
+      "loss": 0.8864,
+      "step": 299
+    },
+    {
+      "epoch": 2.67,
+      "grad_norm": 0.3121296356843828,
+      "learning_rate": 5.145966588423341e-05,
+      "loss": 0.8258,
+      "step": 300
+    },
+    {
+      "epoch": 2.68,
+      "grad_norm": 0.268436849924173,
+      "learning_rate": 5.0832398361230596e-05,
+      "loss": 0.8906,
+      "step": 301
+    },
+    {
+      "epoch": 2.69,
+      "grad_norm": 0.2961161602467319,
+      "learning_rate": 5.020767189299369e-05,
+      "loss": 0.8828,
+      "step": 302
+    },
+    {
+      "epoch": 2.7,
+      "grad_norm": 0.27743957099992345,
+      "learning_rate": 4.9585518766315496e-05,
+      "loss": 0.8251,
+      "step": 303
+    },
+    {
+      "epoch": 2.71,
+      "grad_norm": 0.2949909861852426,
+      "learning_rate": 4.896597113499479e-05,
+      "loss": 0.7911,
+      "step": 304
+    },
+    {
+      "epoch": 2.72,
+      "grad_norm": 0.3161115451278363,
+      "learning_rate": 4.834906101817438e-05,
+      "loss": 0.8157,
+      "step": 305
+    },
+    {
+      "epoch": 2.72,
+      "grad_norm": 0.28720077046065867,
+      "learning_rate": 4.773482029868657e-05,
+      "loss": 0.82,
+      "step": 306
+    },
+    {
+      "epoch": 2.73,
+      "grad_norm": 0.4045319357608716,
+      "learning_rate": 4.712328072140505e-05,
+      "loss": 0.8414,
+      "step": 307
+    },
+    {
+      "epoch": 2.74,
+      "grad_norm": 0.3070232288390269,
+      "learning_rate": 4.651447389160458e-05,
+      "loss": 0.8427,
+      "step": 308
+    },
+    {
+      "epoch": 2.74,
+      "eval_loss": 1.2574400901794434,
+      "eval_runtime": 13.2473,
+      "eval_samples_per_second": 22.646,
+      "eval_steps_per_second": 2.869,
+      "step": 308
+    },
+    {
+      "epoch": 2.75,
+      "grad_norm": 0.3214782806968351,
+      "learning_rate": 4.5908431273327436e-05,
+      "loss": 0.8469,
+      "step": 309
+    },
+    {
+      "epoch": 2.76,
+      "grad_norm": 0.24241410698156174,
+      "learning_rate": 4.530518418775733e-05,
+      "loss": 0.8346,
+      "step": 310
+    },
+    {
+      "epoch": 2.77,
+      "grad_norm": 0.3303263594210879,
+      "learning_rate": 4.470476381160065e-05,
+      "loss": 0.8298,
+      "step": 311
+    },
+    {
+      "epoch": 2.78,
+      "grad_norm": 0.30711900849760865,
+      "learning_rate": 4.4107201175475275e-05,
+      "loss": 0.789,
+      "step": 312
+    },
+    {
+      "epoch": 2.79,
+      "grad_norm": 0.2954465859389713,
+      "learning_rate": 4.351252716230685e-05,
+      "loss": 0.8029,
+      "step": 313
+    },
+    {
+      "epoch": 2.8,
+      "grad_norm": 0.29925087091531116,
+      "learning_rate": 4.292077250573266e-05,
+      "loss": 0.8633,
+      "step": 314
+    },
+    {
+      "epoch": 2.8,
+      "grad_norm": 0.3177611223775825,
+      "learning_rate": 4.2331967788513295e-05,
+      "loss": 0.76,
+      "step": 315
+    },
+    {
+      "epoch": 2.81,
+      "grad_norm": 0.28642407848269513,
+      "learning_rate": 4.174614344095213e-05,
+      "loss": 0.823,
+      "step": 316
+    },
+    {
+      "epoch": 2.82,
+      "grad_norm": 0.3243224656005062,
+      "learning_rate": 4.116332973932256e-05,
+      "loss": 0.7831,
+      "step": 317
+    },
+    {
+      "epoch": 2.83,
+      "grad_norm": 0.34877334027822726,
+      "learning_rate": 4.058355680430337e-05,
+      "loss": 0.899,
+      "step": 318
+    },
+    {
+      "epoch": 2.84,
+      "grad_norm": 0.28640325479143114,
+      "learning_rate": 4.0006854599421926e-05,
+      "loss": 0.8292,
+      "step": 319
+    },
+    {
+      "epoch": 2.85,
+      "grad_norm": 0.3135316628014535,
+      "learning_rate": 3.943325292950579e-05,
+      "loss": 0.8731,
+      "step": 320
+    },
+    {
+      "epoch": 2.86,
+      "grad_norm": 0.2949970604257085,
+      "learning_rate": 3.886278143914219e-05,
+      "loss": 0.8402,
+      "step": 321
+    },
+    {
+      "epoch": 2.87,
+      "grad_norm": 0.30057896586780075,
+      "learning_rate": 3.829546961114607e-05,
+      "loss": 0.7713,
+      "step": 322
+    },
+    {
+      "epoch": 2.88,
+      "grad_norm": 0.3558574270285126,
+      "learning_rate": 3.773134676503629e-05,
+      "loss": 0.8435,
+      "step": 323
+    },
+    {
+      "epoch": 2.88,
+      "grad_norm": 0.29115288332943334,
+      "learning_rate": 3.7170442055520415e-05,
+      "loss": 0.9022,
+      "step": 324
+    },
+    {
+      "epoch": 2.89,
+      "grad_norm": 0.3192074718527619,
+      "learning_rate": 3.661278447098789e-05,
+      "loss": 0.7662,
+      "step": 325
+    },
+    {
+      "epoch": 2.9,
+      "grad_norm": 0.33335742888185405,
+      "learning_rate": 3.605840283201195e-05,
+      "loss": 0.8111,
+      "step": 326
+    },
+    {
+      "epoch": 2.91,
+      "grad_norm": 0.29748212071395186,
+      "learning_rate": 3.550732578986006e-05,
+      "loss": 0.7543,
+      "step": 327
+    },
+    {
+      "epoch": 2.92,
+      "grad_norm": 0.3680409192627914,
+      "learning_rate": 3.495958182501325e-05,
+      "loss": 0.8124,
+      "step": 328
+    },
+    {
+      "epoch": 2.93,
+      "grad_norm": 0.27807302364345643,
+      "learning_rate": 3.441519924569408e-05,
+      "loss": 0.7856,
+      "step": 329
+    },
+    {
+      "epoch": 2.94,
+      "grad_norm": 0.3050855823733691,
+      "learning_rate": 3.387420618640379e-05,
+      "loss": 0.8506,
+      "step": 330
+    },
+    {
+      "epoch": 2.95,
+      "grad_norm": 0.3322620263029238,
+      "learning_rate": 3.3336630606468134e-05,
+      "loss": 0.8771,
+      "step": 331
+    },
+    {
+      "epoch": 2.96,
+      "grad_norm": 0.3112008867427982,
+      "learning_rate": 3.280250028859248e-05,
+      "loss": 0.7785,
+      "step": 332
+    },
+    {
+      "epoch": 2.96,
+      "grad_norm": 0.2839548329095365,
+      "learning_rate": 3.227184283742591e-05,
+      "loss": 0.9153,
+      "step": 333
+    },
+    {
+      "epoch": 2.97,
+      "grad_norm": 0.34615397822650606,
+      "learning_rate": 3.174468567813461e-05,
+      "loss": 0.7753,
+      "step": 334
+    },
+    {
+      "epoch": 2.98,
+      "grad_norm": 0.34691866307772695,
+      "learning_rate": 3.122105605498442e-05,
+      "loss": 0.851,
+      "step": 335
+    },
+    {
+      "epoch": 2.99,
+      "grad_norm": 0.296369624391198,
+      "learning_rate": 3.070098102993302e-05,
+      "loss": 0.8572,
+      "step": 336
+    },
+    {
+      "epoch": 2.99,
+      "eval_loss": 1.2511259317398071,
+      "eval_runtime": 13.2202,
+      "eval_samples_per_second": 22.692,
+      "eval_steps_per_second": 2.874,
+      "step": 336
+    },
+    {
+      "epoch": 3.0,
+      "grad_norm": 0.2960426809914107,
+      "learning_rate": 3.018448748123097e-05,
+      "loss": 0.8664,
+      "step": 337
+    },
+    {
+      "epoch": 3.01,
+      "grad_norm": 0.3213684656412904,
+      "learning_rate": 2.9671602102032926e-05,
+      "loss": 0.7387,
+      "step": 338
+    },
+    {
+      "epoch": 3.02,
+      "grad_norm": 0.3057515564517862,
+      "learning_rate": 2.9162351399017963e-05,
+      "loss": 0.767,
+      "step": 339
+    },
+    {
+      "epoch": 3.03,
+      "grad_norm": 0.31163087624473346,
+      "learning_rate": 2.8656761691019673e-05,
+      "loss": 0.8084,
+      "step": 340
+    },
+    {
+      "epoch": 3.04,
+      "grad_norm": 0.29945884165692627,
+      "learning_rate": 2.8154859107665987e-05,
+      "loss": 0.9033,
+      "step": 341
+    },
+    {
+      "epoch": 3.04,
+      "grad_norm": 0.32235895597138414,
+      "learning_rate": 2.7656669588028762e-05,
+      "loss": 0.7144,
+      "step": 342
+    },
+    {
+      "epoch": 3.05,
+      "grad_norm": 0.3458958643468077,
+      "learning_rate": 2.7162218879283176e-05,
+      "loss": 0.7426,
+      "step": 343
+    },
+    {
+      "epoch": 3.06,
+      "grad_norm": 0.3073811670129139,
+      "learning_rate": 2.667153253537713e-05,
+      "loss": 0.723,
+      "step": 344
+    },
+    {
+      "epoch": 3.07,
+      "grad_norm": 0.3223615792503137,
+      "learning_rate": 2.618463591571052e-05,
+      "loss": 0.6815,
+      "step": 345
+    },
+    {
+      "epoch": 3.08,
+      "grad_norm": 0.357176712253696,
+      "learning_rate": 2.570155418382473e-05,
+      "loss": 0.7286,
+      "step": 346
+    },
+    {
+      "epoch": 3.09,
+      "grad_norm": 0.29547503050650586,
+      "learning_rate": 2.5222312306101925e-05,
+      "loss": 0.7895,
+      "step": 347
+    },
+    {
+      "epoch": 3.1,
+      "grad_norm": 0.37115835330190844,
+      "learning_rate": 2.474693505047504e-05,
+      "loss": 0.8722,
+      "step": 348
+    },
+    {
+      "epoch": 3.11,
+      "grad_norm": 0.3840299049618982,
+      "learning_rate": 2.427544698514753e-05,
+      "loss": 0.7022,
+      "step": 349
+    },
+    {
+      "epoch": 3.12,
+      "grad_norm": 0.33738326420366743,
+      "learning_rate": 2.3807872477323733e-05,
+      "loss": 0.7731,
+      "step": 350
+    },
+    {
+      "epoch": 3.12,
+      "grad_norm": 0.35407291134427865,
+      "learning_rate": 2.334423569194948e-05,
+      "loss": 0.7407,
+      "step": 351
+    },
+    {
+      "epoch": 3.13,
+      "grad_norm": 0.2589965276614251,
+      "learning_rate": 2.288456059046331e-05,
+      "loss": 0.865,
+      "step": 352
+    },
+    {
+      "epoch": 3.14,
+      "grad_norm": 0.3521332875577516,
+      "learning_rate": 2.242887092955801e-05,
+      "loss": 0.655,
+      "step": 353
+    },
+    {
+      "epoch": 3.15,
+      "grad_norm": 0.33409347759889046,
+      "learning_rate": 2.1977190259952883e-05,
+      "loss": 0.7536,
+      "step": 354
+    },
+    {
+      "epoch": 3.16,
+      "grad_norm": 0.38883294070164465,
+      "learning_rate": 2.1529541925176555e-05,
+      "loss": 0.7435,
+      "step": 355
+    },
+    {
+      "epoch": 3.17,
+      "grad_norm": 0.3513324514694806,
+      "learning_rate": 2.1085949060360654e-05,
+      "loss": 0.7405,
+      "step": 356
+    },
+    {
+      "epoch": 3.18,
+      "grad_norm": 0.3391182488874308,
+      "learning_rate": 2.064643459104405e-05,
+      "loss": 0.7737,
+      "step": 357
+    },
+    {
+      "epoch": 3.19,
+      "grad_norm": 0.3647814711811327,
+      "learning_rate": 2.0211021231988102e-05,
+      "loss": 0.7095,
+      "step": 358
+    },
+    {
+      "epoch": 3.2,
+      "grad_norm": 0.35539668264402535,
+      "learning_rate": 1.977973148600266e-05,
+      "loss": 0.7543,
+      "step": 359
+    },
+    {
+      "epoch": 3.2,
+      "grad_norm": 0.30454971181182305,
+      "learning_rate": 1.935258764278314e-05,
+      "loss": 0.7897,
+      "step": 360
+    },
+    {
+      "epoch": 3.21,
+      "grad_norm": 0.33290039075033256,
+      "learning_rate": 1.8929611777758526e-05,
+      "loss": 0.6976,
+      "step": 361
+    },
+    {
+      "epoch": 3.22,
+      "grad_norm": 0.28252566634530724,
+      "learning_rate": 1.851082575095051e-05,
+      "loss": 0.8688,
+      "step": 362
+    },
+    {
+      "epoch": 3.23,
+      "grad_norm": 0.38585709417657277,
+      "learning_rate": 1.8096251205843684e-05,
+      "loss": 0.8361,
+      "step": 363
+    },
+    {
+      "epoch": 3.24,
+      "grad_norm": 0.3372165483048113,
+      "learning_rate": 1.7685909568267033e-05,
+      "loss": 0.7818,
+      "step": 364
+    },
+    {
+      "epoch": 3.24,
+      "eval_loss": 1.2727638483047485,
+      "eval_runtime": 13.2462,
+      "eval_samples_per_second": 22.648,
+      "eval_steps_per_second": 2.869,
+      "step": 364
+    },
+    {
+      "epoch": 3.25,
+      "grad_norm": 0.3250536223337791,
+      "learning_rate": 1.7279822045286576e-05,
+      "loss": 0.7973,
+      "step": 365
+    },
+    {
+      "epoch": 3.26,
+      "grad_norm": 0.3593434237000891,
+      "learning_rate": 1.6878009624109313e-05,
+      "loss": 0.7015,
+      "step": 366
+    },
+    {
+      "epoch": 3.27,
+      "grad_norm": 0.3197427893489943,
+      "learning_rate": 1.648049307099874e-05,
+      "loss": 0.7536,
+      "step": 367
+    },
+    {
+      "epoch": 3.28,
+      "grad_norm": 0.36529609852202527,
+      "learning_rate": 1.6087292930201394e-05,
+      "loss": 0.7487,
+      "step": 368
+    },
+    {
+      "epoch": 3.28,
+      "grad_norm": 0.25007649085986594,
+      "learning_rate": 1.569842952288527e-05,
+      "loss": 0.8906,
+      "step": 369
+    },
+    {
+      "epoch": 3.29,
+      "grad_norm": 0.36329388340138336,
+      "learning_rate": 1.5313922946089486e-05,
+      "loss": 0.7645,
+      "step": 370
+    },
+    {
+      "epoch": 3.3,
+      "grad_norm": 0.35241947629057724,
+      "learning_rate": 1.4933793071685732e-05,
+      "loss": 0.8134,
+      "step": 371
+    },
+    {
+      "epoch": 3.31,
+      "grad_norm": 0.3664249862728154,
+      "learning_rate": 1.4558059545351143e-05,
+      "loss": 0.7328,
+      "step": 372
+    },
+    {
+      "epoch": 3.32,
+      "grad_norm": 0.29196433505411595,
+      "learning_rate": 1.4186741785553115e-05,
+      "loss": 0.7608,
+      "step": 373
+    },
+    {
+      "epoch": 3.33,
+      "grad_norm": 0.2923688341161706,
+      "learning_rate": 1.3819858982545598e-05,
+      "loss": 0.7696,
+      "step": 374
+    },
+    {
+      "epoch": 3.34,
+      "grad_norm": 0.3520285590111714,
+      "learning_rate": 1.3457430097377421e-05,
+      "loss": 0.7871,
+      "step": 375
+    },
+    {
+      "epoch": 3.35,
+      "grad_norm": 0.30838315755571466,
+      "learning_rate": 1.3099473860912326e-05,
+      "loss": 0.7398,
+      "step": 376
+    },
+    {
+      "epoch": 3.36,
+      "grad_norm": 0.327647793123604,
+      "learning_rate": 1.2746008772860884e-05,
+      "loss": 0.7974,
+      "step": 377
+    },
+    {
+      "epoch": 3.36,
+      "grad_norm": 0.29911511172677985,
+      "learning_rate": 1.2397053100824463e-05,
+      "loss": 0.8535,
+      "step": 378
+    },
+    {
+      "epoch": 3.37,
+      "grad_norm": 0.34600894138601895,
+      "learning_rate": 1.2052624879351104e-05,
+      "loss": 0.6872,
+      "step": 379
+    },
+    {
+      "epoch": 3.38,
+      "grad_norm": 0.2993988371658933,
+      "learning_rate": 1.1712741909003444e-05,
+      "loss": 0.8554,
+      "step": 380
+    },
+    {
+      "epoch": 3.39,
+      "grad_norm": 0.34130407502289717,
+      "learning_rate": 1.1377421755438832e-05,
+      "loss": 0.9014,
+      "step": 381
+    },
+    {
+      "epoch": 3.4,
+      "grad_norm": 0.33654746414289566,
+      "learning_rate": 1.1046681748501408e-05,
+      "loss": 0.853,
+      "step": 382
+    },
+    {
+      "epoch": 3.41,
+      "grad_norm": 0.32172878518370435,
+      "learning_rate": 1.0720538981326556e-05,
+      "loss": 0.7594,
+      "step": 383
+    },
+    {
+      "epoch": 3.42,
+      "grad_norm": 0.3211202670840339,
+      "learning_rate": 1.0399010309457457e-05,
+      "loss": 0.8017,
+      "step": 384
+    },
+    {
+      "epoch": 3.43,
+      "grad_norm": 0.31506066245322806,
+      "learning_rate": 1.0082112349974016e-05,
+      "loss": 0.7279,
+      "step": 385
+    },
+    {
+      "epoch": 3.44,
+      "grad_norm": 0.33655604261545546,
+      "learning_rate": 9.76986148063398e-06,
+      "loss": 0.8117,
+      "step": 386
+    },
+    {
+      "epoch": 3.45,
+      "grad_norm": 0.34046018532794836,
+      "learning_rate": 9.462273839026624e-06,
+      "loss": 0.7242,
+      "step": 387
+    },
+    {
+      "epoch": 3.45,
+      "grad_norm": 0.36645383790153846,
+      "learning_rate": 9.159365321738655e-06,
+      "loss": 0.6773,
+      "step": 388
+    },
+    {
+      "epoch": 3.46,
+      "grad_norm": 0.29937594196593476,
+      "learning_rate": 8.861151583532656e-06,
+      "loss": 0.8557,
+      "step": 389
+    },
+    {
+      "epoch": 3.47,
+      "grad_norm": 0.382087527071458,
+      "learning_rate": 8.56764803653809e-06,
+      "loss": 0.8144,
+      "step": 390
+    },
+    {
+      "epoch": 3.48,
+      "grad_norm": 0.37165100019020914,
+      "learning_rate": 8.278869849454718e-06,
+      "loss": 0.8668,
+      "step": 391
+    },
+    {
+      "epoch": 3.49,
+      "grad_norm": 0.29949583615336506,
+      "learning_rate": 7.994831946768622e-06,
+      "loss": 0.6793,
+      "step": 392
+    },
+    {
+      "epoch": 3.49,
+      "eval_loss": 1.2786916494369507,
+      "eval_runtime": 13.2527,
+      "eval_samples_per_second": 22.637,
+      "eval_steps_per_second": 2.867,
+      "step": 392
+    },
+    {
+      "epoch": 3.5,
+      "grad_norm": 0.3128750696018338,
+      "learning_rate": 7.715549007981027e-06,
+      "loss": 0.7723,
+      "step": 393
+    },
+    {
+      "epoch": 3.51,
+      "grad_norm": 0.3036786519151417,
+      "learning_rate": 7.441035466849489e-06,
+      "loss": 0.7668,
+      "step": 394
+    },
+    {
+      "epoch": 3.52,
+      "grad_norm": 0.3723258995422837,
+      "learning_rate": 7.171305510642023e-06,
+      "loss": 0.7493,
+      "step": 395
+    },
+    {
+      "epoch": 3.53,
+      "grad_norm": 0.34640959843456987,
+      "learning_rate": 6.906373079403849e-06,
+      "loss": 0.7326,
+      "step": 396
+    },
+    {
+      "epoch": 3.53,
+      "grad_norm": 0.3708769315457807,
+      "learning_rate": 6.646251865236997e-06,
+      "loss": 0.7532,
+      "step": 397
+    },
+    {
+      "epoch": 3.54,
+      "grad_norm": 0.3204533685892251,
+      "learning_rate": 6.390955311592617e-06,
+      "loss": 0.7287,
+      "step": 398
+    },
+    {
+      "epoch": 3.55,
+      "grad_norm": 0.32625233919759344,
+      "learning_rate": 6.140496612576241e-06,
+      "loss": 0.7138,
+      "step": 399
+    },
+    {
+      "epoch": 3.56,
+      "grad_norm": 0.30531644707337335,
+      "learning_rate": 5.8948887122658335e-06,
+      "loss": 0.7123,
+      "step": 400
+    },
+    {
+      "epoch": 3.57,
+      "grad_norm": 0.33767502582039044,
+      "learning_rate": 5.65414430404293e-06,
+      "loss": 0.797,
+      "step": 401
+    },
+    {
+      "epoch": 3.58,
+      "grad_norm": 0.3282428030443849,
+      "learning_rate": 5.418275829936537e-06,
+      "loss": 0.8456,
+      "step": 402
+    },
+    {
+      "epoch": 3.59,
+      "grad_norm": 0.3532395206693705,
+      "learning_rate": 5.187295479980136e-06,
+      "loss": 0.7484,
+      "step": 403
+    },
+    {
+      "epoch": 3.6,
+      "grad_norm": 0.32556237019334716,
+      "learning_rate": 4.961215191581692e-06,
+      "loss": 0.7909,
+      "step": 404
+    },
+    {
+      "epoch": 3.61,
+      "grad_norm": 0.34045978047642994,
+      "learning_rate": 4.740046648906682e-06,
+      "loss": 0.7618,
+      "step": 405
+    },
+    {
+      "epoch": 3.61,
+      "grad_norm": 0.28937387449733615,
+      "learning_rate": 4.523801282274287e-06,
+      "loss": 0.8292,
+      "step": 406
+    },
+    {
+      "epoch": 3.62,
+      "grad_norm": 0.3160933371902765,
+      "learning_rate": 4.312490267566616e-06,
+      "loss": 0.7755,
+      "step": 407
+    },
+    {
+      "epoch": 3.63,
+      "grad_norm": 0.31618578023192107,
+      "learning_rate": 4.106124525651123e-06,
+      "loss": 0.8033,
+      "step": 408
+    },
+    {
+      "epoch": 3.64,
+      "grad_norm": 0.3105026069614601,
+      "learning_rate": 3.904714721816228e-06,
+      "loss": 0.7873,
+      "step": 409
+    },
+    {
+      "epoch": 3.65,
+      "grad_norm": 0.34709823487491576,
+      "learning_rate": 3.7082712652200867e-06,
+      "loss": 0.7882,
+      "step": 410
+    },
+    {
+      "epoch": 3.66,
+      "grad_norm": 0.3349684333518626,
+      "learning_rate": 3.516804308352628e-06,
+      "loss": 0.7841,
+      "step": 411
+    },
+    {
+      "epoch": 3.67,
+      "grad_norm": 0.3467993408892318,
+      "learning_rate": 3.330323746510955e-06,
+      "loss": 0.7605,
+      "step": 412
+    },
+    {
+      "epoch": 3.68,
+      "grad_norm": 0.37371307324588393,
+      "learning_rate": 3.148839217287758e-06,
+      "loss": 0.7229,
+      "step": 413
+    },
+    {
+      "epoch": 3.69,
+      "grad_norm": 0.32271590708905057,
+      "learning_rate": 2.9723601000734016e-06,
+      "loss": 0.7227,
+      "step": 414
+    },
+    {
+      "epoch": 3.69,
+      "grad_norm": 0.3080800887837192,
+      "learning_rate": 2.800895515571078e-06,
+      "loss": 0.7322,
+      "step": 415
+    },
+    {
+      "epoch": 3.7,
+      "grad_norm": 0.3563938469507623,
+      "learning_rate": 2.634454325325497e-06,
+      "loss": 0.7951,
+      "step": 416
+    },
+    {
+      "epoch": 3.71,
+      "grad_norm": 0.3544288347449977,
+      "learning_rate": 2.473045131264862e-06,
+      "loss": 0.7204,
+      "step": 417
+    },
+    {
+      "epoch": 3.72,
+      "grad_norm": 0.37683581963069984,
+      "learning_rate": 2.3166762752563463e-06,
+      "loss": 0.7817,
+      "step": 418
+    },
+    {
+      "epoch": 3.73,
+      "grad_norm": 0.33996733249074595,
+      "learning_rate": 2.1653558386749427e-06,
+      "loss": 0.68,
+      "step": 419
+    },
+    {
+      "epoch": 3.74,
+      "grad_norm": 0.3069725548353936,
+      "learning_rate": 2.0190916419858484e-06,
+      "loss": 0.798,
+      "step": 420
+    },
+    {
+      "epoch": 3.74,
+      "eval_loss": 1.280185580253601,
+      "eval_runtime": 13.1908,
+      "eval_samples_per_second": 22.743,
+      "eval_steps_per_second": 2.881,
+      "step": 420
+    },
+    {
+      "epoch": 3.75,
+      "grad_norm": 0.3417039799474373,
+      "learning_rate": 1.8778912443402242e-06,
+      "loss": 0.7972,
+      "step": 421
+    },
+    {
+      "epoch": 3.76,
+      "grad_norm": 0.36875255309045446,
+      "learning_rate": 1.7417619431845944e-06,
+      "loss": 0.6653,
+      "step": 422
+    },
+    {
+      "epoch": 3.77,
+      "grad_norm": 0.37727682021595976,
+      "learning_rate": 1.6107107738836835e-06,
+      "loss": 0.7994,
+      "step": 423
+    },
+    {
+      "epoch": 3.77,
+      "grad_norm": 0.3680295622641893,
+      "learning_rate": 1.4847445093567836e-06,
+      "loss": 0.7551,
+      "step": 424
+    },
+    {
+      "epoch": 3.78,
+      "grad_norm": 0.34533281482301054,
+      "learning_rate": 1.3638696597277679e-06,
+      "loss": 0.7416,
+      "step": 425
+    },
+    {
+      "epoch": 3.79,
+      "grad_norm": 0.3908807108284709,
+      "learning_rate": 1.2480924719885934e-06,
+      "loss": 0.7188,
+      "step": 426
+    },
+    {
+      "epoch": 3.8,
+      "grad_norm": 0.31916637760320243,
+      "learning_rate": 1.1374189296765037e-06,
+      "loss": 0.7594,
+      "step": 427
+    },
+    {
+      "epoch": 3.81,
+      "grad_norm": 0.33316597860051206,
+      "learning_rate": 1.0318547525647315e-06,
+      "loss": 0.7249,
+      "step": 428
+    },
+    {
+      "epoch": 3.82,
+      "grad_norm": 0.3416859037907314,
+      "learning_rate": 9.314053963669245e-07,
+      "loss": 0.8137,
+      "step": 429
+    },
+    {
+      "epoch": 3.83,
+      "grad_norm": 0.3166447199987983,
+      "learning_rate": 8.360760524551814e-07,
+      "loss": 0.7548,
+      "step": 430
+    },
+    {
+      "epoch": 3.84,
+      "grad_norm": 0.3678155762471668,
+      "learning_rate": 7.458716475917559e-07,
+      "loss": 0.6848,
+      "step": 431
+    },
+    {
+      "epoch": 3.85,
+      "grad_norm": 0.3564576656343474,
+      "learning_rate": 6.607968436744272e-07,
+      "loss": 0.7188,
+      "step": 432
+    },
+    {
+      "epoch": 3.85,
+      "grad_norm": 0.31739334088982646,
+      "learning_rate": 5.808560374955585e-07,
+      "loss": 0.6899,
+      "step": 433
+    },
+    {
+      "epoch": 3.86,
+      "grad_norm": 0.3517404847481434,
+      "learning_rate": 5.06053360514902e-07,
+      "loss": 0.7085,
+      "step": 434
+    },
+    {
+      "epoch": 3.87,
+      "grad_norm": 0.32714795923909773,
+      "learning_rate": 4.363926786460359e-07,
+      "loss": 0.8216,
+      "step": 435
+    },
+    {
+      "epoch": 3.88,
+      "grad_norm": 0.3635146713982181,
+      "learning_rate": 3.718775920565687e-07,
+      "loss": 0.7684,
+      "step": 436
+    },
+    {
+      "epoch": 3.89,
+      "grad_norm": 0.3419976870684463,
+      "learning_rate": 3.125114349821212e-07,
+      "loss": 0.7825,
+      "step": 437
+    },
+    {
+      "epoch": 3.9,
+      "grad_norm": 0.3872335456149991,
+      "learning_rate": 2.5829727555397584e-07,
+      "loss": 0.8252,
+      "step": 438
+    },
+    {
+      "epoch": 3.91,
+      "grad_norm": 0.3918700023496713,
+      "learning_rate": 2.0923791564050333e-07,
+      "loss": 0.6846,
+      "step": 439
+    },
+    {
+      "epoch": 3.92,
+      "grad_norm": 0.3039893808886515,
+      "learning_rate": 1.6533589070237832e-07,
+      "loss": 0.6893,
+      "step": 440
+    },
+    {
+      "epoch": 3.93,
+      "grad_norm": 0.2983095993470017,
+      "learning_rate": 1.2659346966152895e-07,
+      "loss": 0.7482,
+      "step": 441
+    },
+    {
+      "epoch": 3.93,
+      "grad_norm": 0.29173707821708994,
+      "learning_rate": 9.30126547838861e-08,
+      "loss": 0.8286,
+      "step": 442
+    },
+    {
+      "epoch": 3.94,
+      "grad_norm": 0.3410941283765189,
+      "learning_rate": 6.459518157591048e-08,
+      "loss": 0.7608,
+      "step": 443
+    },
+    {
+      "epoch": 3.95,
+      "grad_norm": 0.3551291216161622,
+      "learning_rate": 4.1342518694853506e-08,
+      "loss": 0.8346,
+      "step": 444
+    },
+    {
+      "epoch": 3.96,
+      "grad_norm": 0.340830105966026,
+      "learning_rate": 2.3255867872928882e-08,
+      "loss": 0.7388,
+      "step": 445
+    },
+    {
+      "epoch": 3.97,
+      "grad_norm": 0.3072115448848053,
+      "learning_rate": 1.0336163855129144e-08,
+      "loss": 0.842,
+      "step": 446
+    },
+    {
+      "epoch": 3.98,
+      "grad_norm": 0.3339675653700237,
+      "learning_rate": 2.5840743509863983e-09,
+      "loss": 0.7289,
+      "step": 447
+    },
+    {
+      "epoch": 3.99,
+      "grad_norm": 0.29581508154823644,
+      "learning_rate": 0.0,
+      "loss": 0.7461,
+      "step": 448
+    },
+    {
+      "epoch": 3.99,
+      "eval_loss": 1.2807621955871582,
+      "eval_runtime": 13.2308,
+      "eval_samples_per_second": 22.674,
+      "eval_steps_per_second": 2.872,
+      "step": 448
+    }
+  ],
+  "logging_steps": 1,
+  "max_steps": 448,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 4,
+  "save_steps": 112,
+  "total_flos": 5.340598301349642e+17,
+  "train_batch_size": 2,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

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

zero_to_fp32.py

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