Super-squash branch 'main' using huggingface_hub

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

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

config.json

@@ -0,0 +1,30 @@
+{
+  "_name_or_path": "/root/autodl-tmp/data/Qwen2.5-1.5B",
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 1536,
+  "initializer_range": 0.02,
+  "intermediate_size": 8960,
+  "max_position_embeddings": 131072,
+  "max_window_layers": 28,
+  "model_type": "qwen2",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.2",
+  "use_cache": true,
+  "use_mrope": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "max_new_tokens": 2048,
+  "transformers_version": "4.48.2"
+}

latest

@@ -0,0 +1 @@
+global_step186

model.safetensors

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

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

scheduler.pt

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

special_tokens_map.json

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

tokenizer.json

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

tokenizer_config.json

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

trainer_state.json

@@ -0,0 +1,2637 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.992,
+  "eval_steps": 10,
+  "global_step": 186,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 444.66667556762695,
+      "epoch": 0.010666666666666666,
+      "grad_norm": 0.13749803602695465,
+      "kl": 0.0,
+      "learning_rate": 1.0714285714285714e-07,
+      "loss": -0.0439,
+      "reward": 0.10416666977107525,
+      "reward_std": 0.1801304928958416,
+      "rewards/accuracy_reward": 0.10416666977107525,
+      "rewards/format_reward": 0.0,
+      "step": 1
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 438.5625114440918,
+      "epoch": 0.021333333333333333,
+      "grad_norm": 0.10789892077445984,
+      "kl": 0.0,
+      "learning_rate": 2.1428571428571428e-07,
+      "loss": 0.0387,
+      "reward": 0.0833333358168602,
+      "reward_std": 0.20412414520978928,
+      "rewards/accuracy_reward": 0.06250000186264515,
+      "rewards/format_reward": 0.02083333395421505,
+      "step": 2
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 440.7083549499512,
+      "epoch": 0.032,
+      "grad_norm": 0.1401524692773819,
+      "kl": 0.00017404556274414062,
+      "learning_rate": 3.214285714285714e-07,
+      "loss": -0.0602,
+      "reward": 0.1250000037252903,
+      "reward_std": 0.23116153478622437,
+      "rewards/accuracy_reward": 0.1250000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 3
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 588.8125152587891,
+      "epoch": 0.042666666666666665,
+      "grad_norm": 0.06724399328231812,
+      "kl": 0.00015866756439208984,
+      "learning_rate": 4.2857142857142857e-07,
+      "loss": -0.0222,
+      "reward": 0.0833333358168602,
+      "reward_std": 0.12909945845603943,
+      "rewards/accuracy_reward": 0.0833333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 4
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 495.20834732055664,
+      "epoch": 0.05333333333333334,
+      "grad_norm": 0.08458305150270462,
+      "kl": 0.00021141767501831055,
+      "learning_rate": 5.357142857142858e-07,
+      "loss": -0.0194,
+      "reward": 0.14583333767950535,
+      "reward_std": 0.28219256177544594,
+      "rewards/accuracy_reward": 0.14583333767950535,
+      "rewards/format_reward": 0.0,
+      "step": 5
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 476.8333435058594,
+      "epoch": 0.064,
+      "grad_norm": 0.10098902881145477,
+      "kl": 0.00014531612396240234,
+      "learning_rate": 6.428571428571428e-07,
+      "loss": -0.0519,
+      "reward": 0.06250000186264515,
+      "reward_std": 0.1530931107699871,
+      "rewards/accuracy_reward": 0.06250000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 6
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 556.3125152587891,
+      "epoch": 0.07466666666666667,
+      "grad_norm": 0.08065751940011978,
+      "kl": 0.00015625357627868652,
+      "learning_rate": 7.5e-07,
+      "loss": 0.0062,
+      "reward": 0.1250000037252903,
+      "reward_std": 0.23116153106093407,
+      "rewards/accuracy_reward": 0.1250000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 7
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 487.5000114440918,
+      "epoch": 0.08533333333333333,
+      "grad_norm": 0.07614883035421371,
+      "kl": 0.00019866228103637695,
+      "learning_rate": 8.571428571428571e-07,
+      "loss": -0.0199,
+      "reward": 0.08333333395421505,
+      "reward_std": 0.11949635669589043,
+      "rewards/accuracy_reward": 0.08333333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 8
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 511.70834732055664,
+      "epoch": 0.096,
+      "grad_norm": 0.09461379051208496,
+      "kl": 0.00015351176261901855,
+      "learning_rate": 9.642857142857145e-07,
+      "loss": 0.0476,
+      "reward": 0.1458333358168602,
+      "reward_std": 0.2350771278142929,
+      "rewards/accuracy_reward": 0.1458333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 9
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 441.1250114440918,
+      "epoch": 0.10666666666666667,
+      "grad_norm": 0.09255574643611908,
+      "kl": 0.00015854835510253906,
+      "learning_rate": 1.0714285714285716e-06,
+      "loss": -0.0865,
+      "reward": 0.14583333767950535,
+      "reward_std": 0.2446802258491516,
+      "rewards/accuracy_reward": 0.14583333767950535,
+      "rewards/format_reward": 0.0,
+      "step": 10
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 540.2916831970215,
+      "epoch": 0.11733333333333333,
+      "grad_norm": 0.12472786009311676,
+      "kl": 0.00023055076599121094,
+      "learning_rate": 1.1785714285714285e-06,
+      "loss": 0.0267,
+      "reward": 0.2083333395421505,
+      "reward_std": 0.3602609857916832,
+      "rewards/accuracy_reward": 0.2083333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 11
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 487.60418128967285,
+      "epoch": 0.128,
+      "grad_norm": 0.029079366475343704,
+      "kl": 0.00019365549087524414,
+      "learning_rate": 1.2857142857142856e-06,
+      "loss": 0.0049,
+      "reward": 0.02083333395421505,
+      "reward_std": 0.05103103443980217,
+      "rewards/accuracy_reward": 0.02083333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 12
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 469.1250114440918,
+      "epoch": 0.13866666666666666,
+      "grad_norm": 0.12546207010746002,
+      "kl": 0.00028073787689208984,
+      "learning_rate": 1.392857142857143e-06,
+      "loss": -0.0073,
+      "reward": 0.1875000037252903,
+      "reward_std": 0.2996268533170223,
+      "rewards/accuracy_reward": 0.16666666977107525,
+      "rewards/format_reward": 0.02083333395421505,
+      "step": 13
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 413.56250762939453,
+      "epoch": 0.14933333333333335,
+      "grad_norm": 0.19048571586608887,
+      "kl": 0.0003783702850341797,
+      "learning_rate": 1.5e-06,
+      "loss": -0.0698,
+      "reward": 0.25000000558793545,
+      "reward_std": 0.4152076207101345,
+      "rewards/accuracy_reward": 0.25000000558793545,
+      "rewards/format_reward": 0.0,
+      "step": 14
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 498.3958511352539,
+      "epoch": 0.16,
+      "grad_norm": 0.08490169793367386,
+      "kl": 0.0004260540008544922,
+      "learning_rate": 1.6071428571428572e-06,
+      "loss": -0.0216,
+      "reward": 0.06250000186264515,
+      "reward_std": 0.1530931107699871,
+      "rewards/accuracy_reward": 0.06250000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 15
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 496.1250114440918,
+      "epoch": 0.17066666666666666,
+      "grad_norm": 0.07948283106088638,
+      "kl": 0.0011513233184814453,
+      "learning_rate": 1.7142857142857143e-06,
+      "loss": 0.0709,
+      "reward": 0.1458333358168602,
+      "reward_std": 0.2350771240890026,
+      "rewards/accuracy_reward": 0.1458333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 16
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 489.4583568572998,
+      "epoch": 0.18133333333333335,
+      "grad_norm": 0.11512169986963272,
+      "kl": 0.000606536865234375,
+      "learning_rate": 1.8214285714285714e-06,
+      "loss": -0.0082,
+      "reward": 0.18750000186264515,
+      "reward_std": 0.2900237552821636,
+      "rewards/accuracy_reward": 0.18750000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 17
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 510.60418701171875,
+      "epoch": 0.192,
+      "grad_norm": 0.1316482424736023,
+      "kl": 0.0010819435119628906,
+      "learning_rate": 1.928571428571429e-06,
+      "loss": 0.0295,
+      "reward": 0.1458333358168602,
+      "reward_std": 0.2350771203637123,
+      "rewards/accuracy_reward": 0.1458333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 18
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 464.7708435058594,
+      "epoch": 0.20266666666666666,
+      "grad_norm": 0.12288669496774673,
+      "kl": 0.0012531280517578125,
+      "learning_rate": 2.035714285714286e-06,
+      "loss": 0.0104,
+      "reward": 0.27083333767950535,
+      "reward_std": 0.41912321001291275,
+      "rewards/accuracy_reward": 0.27083333767950535,
+      "rewards/format_reward": 0.0,
+      "step": 19
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 443.0208435058594,
+      "epoch": 0.21333333333333335,
+      "grad_norm": 0.114227794110775,
+      "kl": 0.0019593238830566406,
+      "learning_rate": 2.142857142857143e-06,
+      "loss": -0.0525,
+      "reward": 0.2083333358168602,
+      "reward_std": 0.16661179810762405,
+      "rewards/accuracy_reward": 0.2083333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 20
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 532.5833473205566,
+      "epoch": 0.224,
+      "grad_norm": 0.10986621677875519,
+      "kl": 0.0016994476318359375,
+      "learning_rate": 2.25e-06,
+      "loss": -0.0249,
+      "reward": 0.14583333767950535,
+      "reward_std": 0.31970490515232086,
+      "rewards/accuracy_reward": 0.14583333767950535,
+      "rewards/format_reward": 0.0,
+      "step": 21
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 492.1875057220459,
+      "epoch": 0.23466666666666666,
+      "grad_norm": 0.0868026539683342,
+      "kl": 0.0018286705017089844,
+      "learning_rate": 2.357142857142857e-06,
+      "loss": -0.0081,
+      "reward": 0.16666666977107525,
+      "reward_std": 0.2861081548035145,
+      "rewards/accuracy_reward": 0.16666666977107525,
+      "rewards/format_reward": 0.0,
+      "step": 22
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 492.72918701171875,
+      "epoch": 0.24533333333333332,
+      "grad_norm": 0.10867961496114731,
+      "kl": 0.0030241012573242188,
+      "learning_rate": 2.4642857142857143e-06,
+      "loss": -0.0295,
+      "reward": 0.2916666753590107,
+      "reward_std": 0.3506578840315342,
+      "rewards/accuracy_reward": 0.2916666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 23
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 555.1875228881836,
+      "epoch": 0.256,
+      "grad_norm": 0.08439142256975174,
+      "kl": 0.0024747848510742188,
+      "learning_rate": 2.571428571428571e-06,
+      "loss": -0.0001,
+      "reward": 0.1875000037252903,
+      "reward_std": 0.2996268533170223,
+      "rewards/accuracy_reward": 0.1875000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 24
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 569.3125152587891,
+      "epoch": 0.26666666666666666,
+      "grad_norm": 0.10544775426387787,
+      "kl": 0.0042171478271484375,
+      "learning_rate": 2.678571428571429e-06,
+      "loss": -0.0134,
+      "reward": 0.1875000037252903,
+      "reward_std": 0.2350771203637123,
+      "rewards/accuracy_reward": 0.1875000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 25
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 481.45834732055664,
+      "epoch": 0.2773333333333333,
+      "grad_norm": 0.10559644550085068,
+      "kl": 0.0049877166748046875,
+      "learning_rate": 2.785714285714286e-06,
+      "loss": -0.0202,
+      "reward": 0.3333333395421505,
+      "reward_std": 0.39512956142425537,
+      "rewards/accuracy_reward": 0.3333333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 26
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 549.7916870117188,
+      "epoch": 0.288,
+      "grad_norm": 0.09459959715604782,
+      "kl": 0.00540924072265625,
+      "learning_rate": 2.892857142857143e-06,
+      "loss": 0.0228,
+      "reward": 0.1875000074505806,
+      "reward_std": 0.21764283999800682,
+      "rewards/accuracy_reward": 0.1875000074505806,
+      "rewards/format_reward": 0.0,
+      "step": 27
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 565.3958473205566,
+      "epoch": 0.2986666666666667,
+      "grad_norm": 0.09324613213539124,
+      "kl": 0.00745844841003418,
+      "learning_rate": 3e-06,
+      "loss": 0.0506,
+      "reward": 0.25000000931322575,
+      "reward_std": 0.3332235999405384,
+      "rewards/accuracy_reward": 0.2291666753590107,
+      "rewards/format_reward": 0.02083333395421505,
+      "step": 28
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 515.1041831970215,
+      "epoch": 0.30933333333333335,
+      "grad_norm": 0.07702303677797318,
+      "kl": 0.009542465209960938,
+      "learning_rate": 2.999882508108334e-06,
+      "loss": 0.0352,
+      "reward": 0.3333333432674408,
+      "reward_std": 0.24859581515192986,
+      "rewards/accuracy_reward": 0.3333333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 29
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 558.8125190734863,
+      "epoch": 0.32,
+      "grad_norm": 0.14078065752983093,
+      "kl": 0.009675979614257812,
+      "learning_rate": 2.9995300508391286e-06,
+      "loss": 0.0497,
+      "reward": 0.29166666977107525,
+      "reward_std": 0.4230388030409813,
+      "rewards/accuracy_reward": 0.29166666977107525,
+      "rewards/format_reward": 0.0,
+      "step": 30
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 514.6875114440918,
+      "epoch": 0.33066666666666666,
+      "grad_norm": 0.10777795314788818,
+      "kl": 0.00920867919921875,
+      "learning_rate": 2.998942683406879e-06,
+      "loss": -0.0089,
+      "reward": 0.25000000558793545,
+      "reward_std": 0.31314554437994957,
+      "rewards/accuracy_reward": 0.25000000558793545,
+      "rewards/format_reward": 0.0,
+      "step": 31
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 541.7708511352539,
+      "epoch": 0.3413333333333333,
+      "grad_norm": 0.07520294934511185,
+      "kl": 0.008222579956054688,
+      "learning_rate": 2.998120497826133e-06,
+      "loss": -0.0195,
+      "reward": 0.35416667349636555,
+      "reward_std": 0.34674229100346565,
+      "rewards/accuracy_reward": 0.35416667349636555,
+      "rewards/format_reward": 0.0,
+      "step": 32
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 512.0833435058594,
+      "epoch": 0.352,
+      "grad_norm": 0.0903107300400734,
+      "kl": 0.01560211181640625,
+      "learning_rate": 2.9970636228970757e-06,
+      "loss": 0.0039,
+      "reward": 0.3333333358168602,
+      "reward_std": 0.33841101825237274,
+      "rewards/accuracy_reward": 0.3333333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 33
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 396.4583435058594,
+      "epoch": 0.3626666666666667,
+      "grad_norm": 0.13230890035629272,
+      "kl": 0.0174560546875,
+      "learning_rate": 2.995772224185354e-06,
+      "loss": 0.0103,
+      "reward": 0.41666667722165585,
+      "reward_std": 0.4326419048011303,
+      "rewards/accuracy_reward": 0.41666667722165585,
+      "rewards/format_reward": 0.0,
+      "step": 34
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 536.2083473205566,
+      "epoch": 0.37333333333333335,
+      "grad_norm": 0.1084352657198906,
+      "kl": 0.01836395263671875,
+      "learning_rate": 2.9942465039961373e-06,
+      "loss": 0.0301,
+      "reward": 0.4583333432674408,
+      "reward_std": 0.23899271339178085,
+      "rewards/accuracy_reward": 0.4583333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 35
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 508.3750190734863,
+      "epoch": 0.384,
+      "grad_norm": 0.11348140239715576,
+      "kl": 0.016246795654296875,
+      "learning_rate": 2.992486701342427e-06,
+      "loss": 0.0617,
+      "reward": 0.5416666809469461,
+      "reward_std": 0.4971916303038597,
+      "rewards/accuracy_reward": 0.5416666809469461,
+      "rewards/format_reward": 0.0,
+      "step": 36
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 428.7708492279053,
+      "epoch": 0.39466666666666667,
+      "grad_norm": 0.09864811599254608,
+      "kl": 0.01728057861328125,
+      "learning_rate": 2.9904930919076145e-06,
+      "loss": -0.0272,
+      "reward": 0.6250000149011612,
+      "reward_std": 0.2957112565636635,
+      "rewards/accuracy_reward": 0.6250000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 37
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 519.5833473205566,
+      "epoch": 0.4053333333333333,
+      "grad_norm": 0.08879904448986053,
+      "kl": 0.0256500244140625,
+      "learning_rate": 2.9882659880022905e-06,
+      "loss": 0.0216,
+      "reward": 0.35416667349636555,
+      "reward_std": 0.34674229472875595,
+      "rewards/accuracy_reward": 0.35416667349636555,
+      "rewards/format_reward": 0.0,
+      "step": 38
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 462.75001525878906,
+      "epoch": 0.416,
+      "grad_norm": 0.0630691796541214,
+      "kl": 0.01651763916015625,
+      "learning_rate": 2.985805738515323e-06,
+      "loss": 0.0325,
+      "reward": 0.8125000186264515,
+      "reward_std": 0.28219256550073624,
+      "rewards/accuracy_reward": 0.8125000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 39
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 522.7916831970215,
+      "epoch": 0.4266666666666667,
+      "grad_norm": 0.08862201869487762,
+      "kl": 0.01795196533203125,
+      "learning_rate": 2.9831127288592004e-06,
+      "loss": 0.0388,
+      "reward": 0.5625000074505806,
+      "reward_std": 0.4539917819201946,
+      "rewards/accuracy_reward": 0.5625000074505806,
+      "rewards/format_reward": 0.0,
+      "step": 40
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 448.2916793823242,
+      "epoch": 0.43733333333333335,
+      "grad_norm": 0.09053761512041092,
+      "kl": 0.01364898681640625,
+      "learning_rate": 2.9801873809096543e-06,
+      "loss": -0.0464,
+      "reward": 0.6875000018626451,
+      "reward_std": 0.2525114081799984,
+      "rewards/accuracy_reward": 0.6875000018626451,
+      "rewards/format_reward": 0.0,
+      "step": 41
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 431.5625114440918,
+      "epoch": 0.448,
+      "grad_norm": 0.11968156695365906,
+      "kl": 0.02625274658203125,
+      "learning_rate": 2.97703015293957e-06,
+      "loss": 0.0643,
+      "reward": 0.6875000279396772,
+      "reward_std": 0.4488043636083603,
+      "rewards/accuracy_reward": 0.6875000279396772,
+      "rewards/format_reward": 0.0,
+      "step": 42
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 488.56250762939453,
+      "epoch": 0.45866666666666667,
+      "grad_norm": 0.1140248030424118,
+      "kl": 0.027557373046875,
+      "learning_rate": 2.9736415395471975e-06,
+      "loss": 0.026,
+      "reward": 0.6250000223517418,
+      "reward_std": 0.4152076132595539,
+      "rewards/accuracy_reward": 0.6250000223517418,
+      "rewards/format_reward": 0.0,
+      "step": 43
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 411.1875190734863,
+      "epoch": 0.4693333333333333,
+      "grad_norm": 0.07253677397966385,
+      "kl": 0.0226898193359375,
+      "learning_rate": 2.970022071578666e-06,
+      "loss": 0.076,
+      "reward": 0.8125000149011612,
+      "reward_std": 0.299626849591732,
+      "rewards/accuracy_reward": 0.8125000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 44
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 503.02084732055664,
+      "epoch": 0.48,
+      "grad_norm": 0.08070115745067596,
+      "kl": 0.024505615234375,
+      "learning_rate": 2.9661723160448267e-06,
+      "loss": 0.0922,
+      "reward": 0.5625000111758709,
+      "reward_std": 0.41912321373820305,
+      "rewards/accuracy_reward": 0.5625000111758709,
+      "rewards/format_reward": 0.0,
+      "step": 45
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 453.56251525878906,
+      "epoch": 0.49066666666666664,
+      "grad_norm": 0.06279083341360092,
+      "kl": 0.019269943237304688,
+      "learning_rate": 2.962092876032427e-06,
+      "loss": -0.034,
+      "reward": 0.7708333507180214,
+      "reward_std": 0.21764283254742622,
+      "rewards/accuracy_reward": 0.7708333507180214,
+      "rewards/format_reward": 0.0,
+      "step": 46
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 412.5833511352539,
+      "epoch": 0.5013333333333333,
+      "grad_norm": 0.04880174994468689,
+      "kl": 0.033538818359375,
+      "learning_rate": 2.957784390609632e-06,
+      "loss": -0.0072,
+      "reward": 0.7500000074505806,
+      "reward_std": 0.18404608592391014,
+      "rewards/accuracy_reward": 0.7500000074505806,
+      "rewards/format_reward": 0.0,
+      "step": 47
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 434.12501525878906,
+      "epoch": 0.512,
+      "grad_norm": 0.05984590947628021,
+      "kl": 0.02077484130859375,
+      "learning_rate": 2.9532475347259113e-06,
+      "loss": 0.033,
+      "reward": 0.8125000111758709,
+      "reward_std": 0.2446802221238613,
+      "rewards/accuracy_reward": 0.8125000111758709,
+      "rewards/format_reward": 0.0,
+      "step": 48
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 477.208345413208,
+      "epoch": 0.5226666666666666,
+      "grad_norm": 0.12303212285041809,
+      "kl": 0.02848052978515625,
+      "learning_rate": 2.9484830191063057e-06,
+      "loss": 0.0671,
+      "reward": 0.6250000204890966,
+      "reward_std": 0.3332235962152481,
+      "rewards/accuracy_reward": 0.6250000204890966,
+      "rewards/format_reward": 0.0,
+      "step": 49
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 582.6041831970215,
+      "epoch": 0.5333333333333333,
+      "grad_norm": 0.09904547035694122,
+      "kl": 0.015716552734375,
+      "learning_rate": 2.9434915901400857e-06,
+      "loss": 0.0423,
+      "reward": 0.6458333507180214,
+      "reward_std": 0.41912321001291275,
+      "rewards/accuracy_reward": 0.6458333507180214,
+      "rewards/format_reward": 0.0,
+      "step": 50
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 409.645845413208,
+      "epoch": 0.544,
+      "grad_norm": 0.10279008746147156,
+      "kl": 0.04036712646484375,
+      "learning_rate": 2.938274029763826e-06,
+      "loss": 0.0204,
+      "reward": 0.6250000186264515,
+      "reward_std": 0.3506578877568245,
+      "rewards/accuracy_reward": 0.6250000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 51
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 539.9166831970215,
+      "epoch": 0.5546666666666666,
+      "grad_norm": 0.11620215326547623,
+      "kl": 0.03099822998046875,
+      "learning_rate": 2.9328311553389116e-06,
+      "loss": 0.0362,
+      "reward": 0.7500000223517418,
+      "reward_std": 0.4056045152246952,
+      "rewards/accuracy_reward": 0.7500000223517418,
+      "rewards/format_reward": 0.0,
+      "step": 52
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 539.6875152587891,
+      "epoch": 0.5653333333333334,
+      "grad_norm": 0.0942852720618248,
+      "kl": 0.021152496337890625,
+      "learning_rate": 2.927163819523492e-06,
+      "loss": 0.0986,
+      "reward": 0.5625000149011612,
+      "reward_std": 0.2996268533170223,
+      "rewards/accuracy_reward": 0.5625000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 53
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 448.3125190734863,
+      "epoch": 0.576,
+      "grad_norm": 0.06090138480067253,
+      "kl": 0.026214599609375,
+      "learning_rate": 2.9212729101389084e-06,
+      "loss": 0.041,
+      "reward": 0.7708333395421505,
+      "reward_std": 0.1530931070446968,
+      "rewards/accuracy_reward": 0.7708333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 54
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 634.8750114440918,
+      "epoch": 0.5866666666666667,
+      "grad_norm": 0.10308913886547089,
+      "kl": 0.032123565673828125,
+      "learning_rate": 2.9151593500306103e-06,
+      "loss": 0.021,
+      "reward": 0.4166666753590107,
+      "reward_std": 0.2861081585288048,
+      "rewards/accuracy_reward": 0.4166666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 55
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 456.2708625793457,
+      "epoch": 0.5973333333333334,
+      "grad_norm": 0.09447697550058365,
+      "kl": 0.0720062255859375,
+      "learning_rate": 2.9088240969235863e-06,
+      "loss": 0.0148,
+      "reward": 0.7291666828095913,
+      "reward_std": 0.36417657881975174,
+      "rewards/accuracy_reward": 0.7291666828095913,
+      "rewards/format_reward": 0.0,
+      "step": 56
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 505.68751525878906,
+      "epoch": 0.608,
+      "grad_norm": 0.09954586625099182,
+      "kl": 0.04242706298828125,
+      "learning_rate": 2.902268143272333e-06,
+      "loss": 0.0686,
+      "reward": 0.7708333395421505,
+      "reward_std": 0.2525114119052887,
+      "rewards/accuracy_reward": 0.7708333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 57
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 513.0208511352539,
+      "epoch": 0.6186666666666667,
+      "grad_norm": 0.07663288712501526,
+      "kl": 0.02385711669921875,
+      "learning_rate": 2.895492516105378e-06,
+      "loss": 0.0923,
+      "reward": 0.6250000167638063,
+      "reward_std": 0.4152076132595539,
+      "rewards/accuracy_reward": 0.6250000167638063,
+      "rewards/format_reward": 0.0,
+      "step": 58
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 639.5625076293945,
+      "epoch": 0.6293333333333333,
+      "grad_norm": 0.1266147643327713,
+      "kl": 0.04524993896484375,
+      "learning_rate": 2.888498276864393e-06,
+      "loss": 0.086,
+      "reward": 0.4375000149011612,
+      "reward_std": 0.37377968057990074,
+      "rewards/accuracy_reward": 0.4375000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 59
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 570.666690826416,
+      "epoch": 0.64,
+      "grad_norm": 0.08474475890398026,
+      "kl": 0.02214813232421875,
+      "learning_rate": 2.8812865212379087e-06,
+      "loss": 0.0572,
+      "reward": 0.6666666865348816,
+      "reward_std": 0.3506578877568245,
+      "rewards/accuracy_reward": 0.6666666865348816,
+      "rewards/format_reward": 0.0,
+      "step": 60
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 586.4375152587891,
+      "epoch": 0.6506666666666666,
+      "grad_norm": 0.12999482452869415,
+      "kl": 0.03527069091796875,
+      "learning_rate": 2.8738583789896743e-06,
+      "loss": 0.085,
+      "reward": 0.6666666828095913,
+      "reward_std": 0.43264190107584,
+      "rewards/accuracy_reward": 0.6666666828095913,
+      "rewards/format_reward": 0.0,
+      "step": 61
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 606.8541831970215,
+      "epoch": 0.6613333333333333,
+      "grad_norm": 0.07179641723632812,
+      "kl": 0.03192138671875,
+      "learning_rate": 2.866215013781668e-06,
+      "loss": 0.006,
+      "reward": 0.5000000093132257,
+      "reward_std": 0.4326419048011303,
+      "rewards/accuracy_reward": 0.5000000093132257,
+      "rewards/format_reward": 0.0,
+      "step": 62
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 473.2500190734863,
+      "epoch": 0.672,
+      "grad_norm": 0.1155165582895279,
+      "kl": 0.0264434814453125,
+      "learning_rate": 2.858357622991806e-06,
+      "loss": 0.0919,
+      "reward": 0.7500000260770321,
+      "reward_std": 0.3977733328938484,
+      "rewards/accuracy_reward": 0.7500000260770321,
+      "rewards/format_reward": 0.0,
+      "step": 63
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 568.0833473205566,
+      "epoch": 0.6826666666666666,
+      "grad_norm": 0.08376800268888474,
+      "kl": 0.04048919677734375,
+      "learning_rate": 2.8502874375263646e-06,
+      "loss": 0.0631,
+      "reward": 0.5833333414047956,
+      "reward_std": 0.23116153106093407,
+      "rewards/accuracy_reward": 0.5833333414047956,
+      "rewards/format_reward": 0.0,
+      "step": 64
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 440.8333435058594,
+      "epoch": 0.6933333333333334,
+      "grad_norm": 0.08705708384513855,
+      "kl": 0.04457855224609375,
+      "learning_rate": 2.8420057216271533e-06,
+      "loss": 0.0036,
+      "reward": 0.7083333395421505,
+      "reward_std": 0.32097673043608665,
+      "rewards/accuracy_reward": 0.7083333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 65
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 592.3125228881836,
+      "epoch": 0.704,
+      "grad_norm": 0.1033933237195015,
+      "kl": 0.0416717529296875,
+      "learning_rate": 2.8335137726734607e-06,
+      "loss": 0.1093,
+      "reward": 0.625000013038516,
+      "reward_std": 0.3602609857916832,
+      "rewards/accuracy_reward": 0.625000013038516,
+      "rewards/format_reward": 0.0,
+      "step": 66
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 426.91667556762695,
+      "epoch": 0.7146666666666667,
+      "grad_norm": 0.13923761248588562,
+      "kl": 0.04537200927734375,
+      "learning_rate": 2.8248129209788164e-06,
+      "loss": 0.0978,
+      "reward": 0.7291666865348816,
+      "reward_std": 0.34674229100346565,
+      "rewards/accuracy_reward": 0.7291666865348816,
+      "rewards/format_reward": 0.0,
+      "step": 67
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 488.93752670288086,
+      "epoch": 0.7253333333333334,
+      "grad_norm": 0.08182822912931442,
+      "kl": 0.0292205810546875,
+      "learning_rate": 2.8159045295825857e-06,
+      "loss": -0.0328,
+      "reward": 0.5833333469927311,
+      "reward_std": 0.4152076207101345,
+      "rewards/accuracy_reward": 0.5833333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 68
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 485.5625114440918,
+      "epoch": 0.736,
+      "grad_norm": 0.07835296541452408,
+      "kl": 0.03786468505859375,
+      "learning_rate": 2.806789994036445e-06,
+      "loss": 0.0566,
+      "reward": 0.6875000074505806,
+      "reward_std": 0.3266642391681671,
+      "rewards/accuracy_reward": 0.6875000074505806,
+      "rewards/format_reward": 0.0,
+      "step": 69
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 507.75001525878906,
+      "epoch": 0.7466666666666667,
+      "grad_norm": 0.1188604012131691,
+      "kl": 0.0541534423828125,
+      "learning_rate": 2.7974707421857584e-06,
+      "loss": 0.0698,
+      "reward": 0.5625000055879354,
+      "reward_std": 0.3170611374080181,
+      "rewards/accuracy_reward": 0.5625000055879354,
+      "rewards/format_reward": 0.0,
+      "step": 70
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 556.2708511352539,
+      "epoch": 0.7573333333333333,
+      "grad_norm": 0.07966376841068268,
+      "kl": 0.04036712646484375,
+      "learning_rate": 2.7879482339458973e-06,
+      "loss": 0.0448,
+      "reward": 0.6458333488553762,
+      "reward_std": 0.27258946001529694,
+      "rewards/accuracy_reward": 0.6458333488553762,
+      "rewards/format_reward": 0.0,
+      "step": 71
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 532.9583511352539,
+      "epoch": 0.768,
+      "grad_norm": 0.08758682757616043,
+      "kl": 0.038604736328125,
+      "learning_rate": 2.778223961073537e-06,
+      "loss": 0.0497,
+      "reward": 0.5208333469927311,
+      "reward_std": 0.33713919669389725,
+      "rewards/accuracy_reward": 0.5208333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 72
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 396.0208435058594,
+      "epoch": 0.7786666666666666,
+      "grad_norm": 0.08768614381551743,
+      "kl": 0.0446929931640625,
+      "learning_rate": 2.7682994469329646e-06,
+      "loss": 0.0424,
+      "reward": 0.8333333414047956,
+      "reward_std": 0.16661180183291435,
+      "rewards/accuracy_reward": 0.8333333414047956,
+      "rewards/format_reward": 0.0,
+      "step": 73
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 467.0833549499512,
+      "epoch": 0.7893333333333333,
+      "grad_norm": 0.06757903844118118,
+      "kl": 0.032470703125,
+      "learning_rate": 2.7581762462574334e-06,
+      "loss": 0.0046,
+      "reward": 0.6666666753590107,
+      "reward_std": 0.2861081622540951,
+      "rewards/accuracy_reward": 0.6666666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 74
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 694.7291793823242,
+      "epoch": 0.8,
+      "grad_norm": 0.08728935569524765,
+      "kl": 0.0474700927734375,
+      "learning_rate": 2.747855944905606e-06,
+      "loss": 0.0407,
+      "reward": 0.4375000149011612,
+      "reward_std": 0.36417657881975174,
+      "rewards/accuracy_reward": 0.4375000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 75
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 500.8333549499512,
+      "epoch": 0.8106666666666666,
+      "grad_norm": 0.12662656605243683,
+      "kl": 0.0447845458984375,
+      "learning_rate": 2.73734015961312e-06,
+      "loss": 0.0734,
+      "reward": 0.7291666939854622,
+      "reward_std": 0.4662386514246464,
+      "rewards/accuracy_reward": 0.7291666939854622,
+      "rewards/format_reward": 0.0,
+      "step": 76
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 530.729190826416,
+      "epoch": 0.8213333333333334,
+      "grad_norm": 0.11303848773241043,
+      "kl": 0.0428314208984375,
+      "learning_rate": 2.7266305377393183e-06,
+      "loss": 0.092,
+      "reward": 0.7500000186264515,
+      "reward_std": 0.3506578840315342,
+      "rewards/accuracy_reward": 0.7500000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 77
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 542.1875190734863,
+      "epoch": 0.832,
+      "grad_norm": 0.08298168331384659,
+      "kl": 0.051544189453125,
+      "learning_rate": 2.7157287570091765e-06,
+      "loss": -0.0051,
+      "reward": 0.645833358168602,
+      "reward_std": 0.4758417494595051,
+      "rewards/accuracy_reward": 0.645833358168602,
+      "rewards/format_reward": 0.0,
+      "step": 78
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 609.291675567627,
+      "epoch": 0.8426666666666667,
+      "grad_norm": 0.13366347551345825,
+      "kl": 0.07353973388671875,
+      "learning_rate": 2.7046365252504835e-06,
+      "loss": 0.1406,
+      "reward": 0.6666666865348816,
+      "reward_std": 0.3506578952074051,
+      "rewards/accuracy_reward": 0.6666666865348816,
+      "rewards/format_reward": 0.0,
+      "step": 79
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 436.02084732055664,
+      "epoch": 0.8533333333333334,
+      "grad_norm": 0.07604128122329712,
+      "kl": 0.0461883544921875,
+      "learning_rate": 2.693355580126295e-06,
+      "loss": 0.0742,
+      "reward": 0.8333333432674408,
+      "reward_std": 0.25819891691207886,
+      "rewards/accuracy_reward": 0.8333333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 80
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 544.1458511352539,
+      "epoch": 0.864,
+      "grad_norm": 0.0864926278591156,
+      "kl": 0.08890533447265625,
+      "learning_rate": 2.6818876888627203e-06,
+      "loss": 0.0639,
+      "reward": 0.6875000149011612,
+      "reward_std": 0.290023747831583,
+      "rewards/accuracy_reward": 0.6875000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 81
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 499.02084732055664,
+      "epoch": 0.8746666666666667,
+      "grad_norm": 0.12025976926088333,
+      "kl": 0.0964202880859375,
+      "learning_rate": 2.6702346479720777e-06,
+      "loss": 0.0518,
+      "reward": 0.5833333395421505,
+      "reward_std": 0.22155842930078506,
+      "rewards/accuracy_reward": 0.5833333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 82
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 565.2083549499512,
+      "epoch": 0.8853333333333333,
+      "grad_norm": 0.10178246349096298,
+      "kl": 0.06887054443359375,
+      "learning_rate": 2.6583982829714576e-06,
+      "loss": 0.1361,
+      "reward": 0.6458333488553762,
+      "reward_std": 0.28219256550073624,
+      "rewards/accuracy_reward": 0.6458333488553762,
+      "rewards/format_reward": 0.0,
+      "step": 83
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 551.5833396911621,
+      "epoch": 0.896,
+      "grad_norm": 0.05114445462822914,
+      "kl": 0.0851898193359375,
+      "learning_rate": 2.6463804480967447e-06,
+      "loss": 0.0434,
+      "reward": 0.6875000055879354,
+      "reward_std": 0.2446802221238613,
+      "rewards/accuracy_reward": 0.6875000055879354,
+      "rewards/format_reward": 0.0,
+      "step": 84
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 572.8750190734863,
+      "epoch": 0.9066666666666666,
+      "grad_norm": 0.05700259283185005,
+      "kl": 0.092803955078125,
+      "learning_rate": 2.6341830260121445e-06,
+      "loss": 0.0154,
+      "reward": 0.583333333954215,
+      "reward_std": 0.25642700120806694,
+      "rewards/accuracy_reward": 0.583333333954215,
+      "rewards/format_reward": 0.0,
+      "step": 85
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 571.4791870117188,
+      "epoch": 0.9173333333333333,
+      "grad_norm": 0.10524018853902817,
+      "kl": 0.10492706298828125,
+      "learning_rate": 2.6218079275152485e-06,
+      "loss": 0.0464,
+      "reward": 0.5416666734963655,
+      "reward_std": 0.30354244634509087,
+      "rewards/accuracy_reward": 0.5416666734963655,
+      "rewards/format_reward": 0.0,
+      "step": 86
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 640.2291870117188,
+      "epoch": 0.928,
+      "grad_norm": 0.09200317412614822,
+      "kl": 0.1671142578125,
+      "learning_rate": 2.609257091237699e-06,
+      "loss": 0.0551,
+      "reward": 0.5208333395421505,
+      "reward_std": 0.43655749410390854,
+      "rewards/accuracy_reward": 0.5208333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 87
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 571.8125190734863,
+      "epoch": 0.9386666666666666,
+      "grad_norm": 0.07907970249652863,
+      "kl": 0.14947509765625,
+      "learning_rate": 2.5965324833414926e-06,
+      "loss": 0.0074,
+      "reward": 0.5416666734963655,
+      "reward_std": 0.3410547934472561,
+      "rewards/accuracy_reward": 0.5416666734963655,
+      "rewards/format_reward": 0.0,
+      "step": 88
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 658.6458511352539,
+      "epoch": 0.9493333333333334,
+      "grad_norm": 0.07123763114213943,
+      "kl": 0.1113739013671875,
+      "learning_rate": 2.5836360972109642e-06,
+      "loss": 0.0316,
+      "reward": 0.5000000093132257,
+      "reward_std": 0.23116152733564377,
+      "rewards/accuracy_reward": 0.5000000093132257,
+      "rewards/format_reward": 0.0,
+      "step": 89
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 558.2916831970215,
+      "epoch": 0.96,
+      "grad_norm": 0.09730823338031769,
+      "kl": 0.18788909912109375,
+      "learning_rate": 2.5705699531405174e-06,
+      "loss": 0.0815,
+      "reward": 0.7083333395421505,
+      "reward_std": 0.33057983219623566,
+      "rewards/accuracy_reward": 0.7083333395421505,
+      "rewards/format_reward": 0.0,
+      "step": 90
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 599.0208587646484,
+      "epoch": 0.9706666666666667,
+      "grad_norm": 0.07070768624544144,
+      "kl": 0.12865447998046875,
+      "learning_rate": 2.55733609801813e-06,
+      "loss": 0.0396,
+      "reward": 0.5000000111758709,
+      "reward_std": 0.3410547822713852,
+      "rewards/accuracy_reward": 0.5000000111758709,
+      "rewards/format_reward": 0.0,
+      "step": 91
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 516.7916774749756,
+      "epoch": 0.9813333333333333,
+      "grad_norm": 0.06625442951917648,
+      "kl": 0.0931243896484375,
+      "learning_rate": 2.5439366050046985e-06,
+      "loss": -0.0295,
+      "reward": 0.5000000074505806,
+      "reward_std": 0.2957112602889538,
+      "rewards/accuracy_reward": 0.5000000074505806,
+      "rewards/format_reward": 0.0,
+      "step": 92
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 580.8125076293945,
+      "epoch": 0.992,
+      "grad_norm": 0.15040598809719086,
+      "kl": 0.15094757080078125,
+      "learning_rate": 2.5303735732092655e-06,
+      "loss": 0.1105,
+      "reward": 0.5416666772216558,
+      "reward_std": 0.23116153478622437,
+      "rewards/accuracy_reward": 0.5416666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 93
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 569.8125190734863,
+      "epoch": 1.0106666666666666,
+      "grad_norm": 0.11118441075086594,
+      "kl": 0.1716766357421875,
+      "learning_rate": 2.51664912736018e-06,
+      "loss": 0.0783,
+      "reward": 0.583333345130086,
+      "reward_std": 0.2861081622540951,
+      "rewards/accuracy_reward": 0.583333345130086,
+      "rewards/format_reward": 0.0,
+      "step": 94
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 465.70834732055664,
+      "epoch": 1.0213333333333334,
+      "grad_norm": 0.08838524669408798,
+      "kl": 0.093017578125,
+      "learning_rate": 2.5027654174722495e-06,
+      "loss": 0.0275,
+      "reward": 0.6875000186264515,
+      "reward_std": 0.28219257295131683,
+      "rewards/accuracy_reward": 0.6875000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 95
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 432.81251335144043,
+      "epoch": 1.032,
+      "grad_norm": 0.08520545810461044,
+      "kl": 0.07879638671875,
+      "learning_rate": 2.4887246185099237e-06,
+      "loss": 0.0379,
+      "reward": 0.7083333358168602,
+      "reward_std": 0.26603010296821594,
+      "rewards/accuracy_reward": 0.7083333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 96
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 499.9791793823242,
+      "epoch": 1.0426666666666666,
+      "grad_norm": 0.07848313450813293,
+      "kl": 0.0831756591796875,
+      "learning_rate": 2.474528930046578e-06,
+      "loss": 0.0461,
+      "reward": 0.6666666772216558,
+      "reward_std": 0.22155842557549477,
+      "rewards/accuracy_reward": 0.6666666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 97
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 449.45834732055664,
+      "epoch": 1.0533333333333332,
+      "grad_norm": 0.14378738403320312,
+      "kl": 0.0985260009765625,
+      "learning_rate": 2.4601805759199317e-06,
+      "loss": 0.0735,
+      "reward": 0.7708333507180214,
+      "reward_std": 0.35457348078489304,
+      "rewards/accuracy_reward": 0.7708333507180214,
+      "rewards/format_reward": 0.0,
+      "step": 98
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 703.0625114440918,
+      "epoch": 1.064,
+      "grad_norm": 0.2360164374113083,
+      "kl": 0.1814422607421875,
+      "learning_rate": 2.445681803883678e-06,
+      "loss": 0.0771,
+      "reward": 0.6250000204890966,
+      "reward_std": 0.3332235962152481,
+      "rewards/accuracy_reward": 0.6250000204890966,
+      "rewards/format_reward": 0.0,
+      "step": 99
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 486.9166889190674,
+      "epoch": 1.0746666666666667,
+      "grad_norm": 0.13772936165332794,
+      "kl": 0.12012481689453125,
+      "learning_rate": 2.4310348852553554e-06,
+      "loss": 0.1204,
+      "reward": 0.6666666772216558,
+      "reward_std": 0.3680921718478203,
+      "rewards/accuracy_reward": 0.6666666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 100
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 480.7708396911621,
+      "epoch": 1.0853333333333333,
+      "grad_norm": 0.15142787992954254,
+      "kl": 0.169891357421875,
+      "learning_rate": 2.4162421145605308e-06,
+      "loss": 0.033,
+      "reward": 0.7291666753590107,
+      "reward_std": 0.21764283254742622,
+      "rewards/accuracy_reward": 0.7291666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 101
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 586.4375114440918,
+      "epoch": 1.096,
+      "grad_norm": 0.27943626046180725,
+      "kl": 0.2178497314453125,
+      "learning_rate": 2.4013058091733546e-06,
+      "loss": 0.0708,
+      "reward": 0.541666679084301,
+      "reward_std": 0.23116152733564377,
+      "rewards/accuracy_reward": 0.541666679084301,
+      "rewards/format_reward": 0.0,
+      "step": 102
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 508.4583396911621,
+      "epoch": 1.1066666666666667,
+      "grad_norm": 0.8309665322303772,
+      "kl": 0.48883056640625,
+      "learning_rate": 2.3862283089535265e-06,
+      "loss": 0.1714,
+      "reward": 0.6250000242143869,
+      "reward_std": 0.45271996036171913,
+      "rewards/accuracy_reward": 0.6250000242143869,
+      "rewards/format_reward": 0.0,
+      "step": 103
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 552.9583473205566,
+      "epoch": 1.1173333333333333,
+      "grad_norm": 32.438201904296875,
+      "kl": 7.4153289794921875,
+      "learning_rate": 2.371011975879742e-06,
+      "loss": 0.4224,
+      "reward": 0.7291666753590107,
+      "reward_std": 0.21764283627271652,
+      "rewards/accuracy_reward": 0.7291666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 104
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 584.7708473205566,
+      "epoch": 1.1280000000000001,
+      "grad_norm": 0.9769695401191711,
+      "kl": 1.2292022705078125,
+      "learning_rate": 2.3556591936796777e-06,
+      "loss": 0.2518,
+      "reward": 0.645833345130086,
+      "reward_std": 0.36417657881975174,
+      "rewards/accuracy_reward": 0.645833345130086,
+      "rewards/format_reward": 0.0,
+      "step": 105
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 521.3333511352539,
+      "epoch": 1.1386666666666667,
+      "grad_norm": 0.9792843461036682,
+      "kl": 0.8870162963867188,
+      "learning_rate": 2.340172367456564e-06,
+      "loss": 0.1476,
+      "reward": 0.7083333469927311,
+      "reward_std": 0.3776952736079693,
+      "rewards/accuracy_reward": 0.7083333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 106
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 547.5000190734863,
+      "epoch": 1.1493333333333333,
+      "grad_norm": 168.53076171875,
+      "kl": 22.269760131835938,
+      "learning_rate": 2.3245539233124133e-06,
+      "loss": 1.5006,
+      "reward": 0.6041666846722364,
+      "reward_std": 0.33713918924331665,
+      "rewards/accuracy_reward": 0.6041666846722364,
+      "rewards/format_reward": 0.0,
+      "step": 107
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 455.93750953674316,
+      "epoch": 1.16,
+      "grad_norm": 10.638311386108398,
+      "kl": 8.11297607421875,
+      "learning_rate": 2.308806307967955e-06,
+      "loss": 0.4543,
+      "reward": 0.7291666772216558,
+      "reward_std": 0.2900237515568733,
+      "rewards/accuracy_reward": 0.7291666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 108
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 600.3958549499512,
+      "epoch": 1.1706666666666667,
+      "grad_norm": 11.44261646270752,
+      "kl": 7.8447265625,
+      "learning_rate": 2.2929319883793442e-06,
+      "loss": 0.5332,
+      "reward": 0.666666679084301,
+      "reward_std": 0.23116152733564377,
+      "rewards/accuracy_reward": 0.666666679084301,
+      "rewards/format_reward": 0.0,
+      "step": 109
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 565.8333530426025,
+      "epoch": 1.1813333333333333,
+      "grad_norm": 1.495950698852539,
+      "kl": 2.2734375,
+      "learning_rate": 2.276933451351696e-06,
+      "loss": 0.1701,
+      "reward": 0.6666666809469461,
+      "reward_std": 0.3506578840315342,
+      "rewards/accuracy_reward": 0.6666666809469461,
+      "rewards/format_reward": 0.0,
+      "step": 110
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 545.2083549499512,
+      "epoch": 1.192,
+      "grad_norm": 2.101208448410034,
+      "kl": 2.087677001953125,
+      "learning_rate": 2.2608132031495184e-06,
+      "loss": 0.1528,
+      "reward": 0.5833333488553762,
+      "reward_std": 0.20412414520978928,
+      "rewards/accuracy_reward": 0.5833333488553762,
+      "rewards/format_reward": 0.0,
+      "step": 111
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 518.5000171661377,
+      "epoch": 1.2026666666666666,
+      "grad_norm": 1.4472161531448364,
+      "kl": 0.91241455078125,
+      "learning_rate": 2.244573769104084e-06,
+      "loss": 0.1028,
+      "reward": 0.6875000204890966,
+      "reward_std": 0.40168891847133636,
+      "rewards/accuracy_reward": 0.6875000204890966,
+      "rewards/format_reward": 0.0,
+      "step": 112
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 516.2083435058594,
+      "epoch": 1.2133333333333334,
+      "grad_norm": 0.677949845790863,
+      "kl": 0.6656951904296875,
+      "learning_rate": 2.228217693217826e-06,
+      "loss": 0.0863,
+      "reward": 0.6875000149011612,
+      "reward_std": 0.2996268570423126,
+      "rewards/accuracy_reward": 0.6875000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 113
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 509.60418128967285,
+      "epoch": 1.224,
+      "grad_norm": 0.534664511680603,
+      "kl": 0.5305023193359375,
+      "learning_rate": 2.211747537765808e-06,
+      "loss": 0.0607,
+      "reward": 0.6875000186264515,
+      "reward_std": 0.31970490142703056,
+      "rewards/accuracy_reward": 0.6875000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 114
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 477.6666793823242,
+      "epoch": 1.2346666666666666,
+      "grad_norm": 0.6869253516197205,
+      "kl": 1.214813232421875,
+      "learning_rate": 2.195165882894322e-06,
+      "loss": 0.0821,
+      "reward": 0.8541666865348816,
+      "reward_std": 0.28219256550073624,
+      "rewards/accuracy_reward": 0.8541666865348816,
+      "rewards/format_reward": 0.0,
+      "step": 115
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 644.3958587646484,
+      "epoch": 1.2453333333333334,
+      "grad_norm": 0.5632476210594177,
+      "kl": 1.506103515625,
+      "learning_rate": 2.1784753262166986e-06,
+      "loss": 0.1124,
+      "reward": 0.416666679084301,
+      "reward_std": 0.31314554810523987,
+      "rewards/accuracy_reward": 0.416666679084301,
+      "rewards/format_reward": 0.0,
+      "step": 116
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 519.7500152587891,
+      "epoch": 1.256,
+      "grad_norm": 0.8980614542961121,
+      "kl": 1.355926513671875,
+      "learning_rate": 2.161678482406374e-06,
+      "loss": 0.1538,
+      "reward": 0.6666666753590107,
+      "reward_std": 0.24859581515192986,
+      "rewards/accuracy_reward": 0.6666666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 117
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 610.4375152587891,
+      "epoch": 1.2666666666666666,
+      "grad_norm": 0.5385753512382507,
+      "kl": 0.8503570556640625,
+      "learning_rate": 2.144777982787286e-06,
+      "loss": 0.0987,
+      "reward": 0.5625000223517418,
+      "reward_std": 0.34674229472875595,
+      "rewards/accuracy_reward": 0.5625000223517418,
+      "rewards/format_reward": 0.0,
+      "step": 118
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 460.2708435058594,
+      "epoch": 1.2773333333333334,
+      "grad_norm": 0.40609875321388245,
+      "kl": 0.885009765625,
+      "learning_rate": 2.127776474921661e-06,
+      "loss": 0.1355,
+      "reward": 0.8958333432674408,
+      "reward_std": 0.1801304928958416,
+      "rewards/accuracy_reward": 0.8958333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 119
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 682.7500152587891,
+      "epoch": 1.288,
+      "grad_norm": 1.9950181245803833,
+      "kl": 2.8668365478515625,
+      "learning_rate": 2.1106766221952607e-06,
+      "loss": 0.2645,
+      "reward": 0.5208333507180214,
+      "reward_std": 0.37377968057990074,
+      "rewards/accuracy_reward": 0.5208333507180214,
+      "rewards/format_reward": 0.0,
+      "step": 120
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 571.6875152587891,
+      "epoch": 1.2986666666666666,
+      "grad_norm": 4.098674297332764,
+      "kl": 2.38665771484375,
+      "learning_rate": 2.0934811034001427e-06,
+      "loss": 0.2873,
+      "reward": 0.7083333563059568,
+      "reward_std": 0.3881702311336994,
+      "rewards/accuracy_reward": 0.7083333563059568,
+      "rewards/format_reward": 0.0,
+      "step": 121
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 525.354190826416,
+      "epoch": 1.3093333333333335,
+      "grad_norm": 0.361176073551178,
+      "kl": 1.3050384521484375,
+      "learning_rate": 2.0761926123150166e-06,
+      "loss": 0.0896,
+      "reward": 0.6250000223517418,
+      "reward_std": 0.2686738707125187,
+      "rewards/accuracy_reward": 0.6250000223517418,
+      "rewards/format_reward": 0.0,
+      "step": 122
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 679.2500152587891,
+      "epoch": 1.32,
+      "grad_norm": 0.4487684965133667,
+      "kl": 1.046875,
+      "learning_rate": 2.058813857283244e-06,
+      "loss": 0.159,
+      "reward": 0.5416666809469461,
+      "reward_std": 0.3506578877568245,
+      "rewards/accuracy_reward": 0.5416666809469461,
+      "rewards/format_reward": 0.0,
+      "step": 123
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 574.8958435058594,
+      "epoch": 1.3306666666666667,
+      "grad_norm": 1.1202152967453003,
+      "kl": 1.6433334350585938,
+      "learning_rate": 2.041347560788562e-06,
+      "loss": 0.2923,
+      "reward": 0.6041666902601719,
+      "reward_std": 0.40168892592191696,
+      "rewards/accuracy_reward": 0.6041666902601719,
+      "rewards/format_reward": 0.0,
+      "step": 124
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 678.3333511352539,
+      "epoch": 1.3413333333333333,
+      "grad_norm": 1.3484364748001099,
+      "kl": 3.1891632080078125,
+      "learning_rate": 2.0237964590285924e-06,
+      "loss": 0.1732,
+      "reward": 0.604166679084301,
+      "reward_std": 0.309229951351881,
+      "rewards/accuracy_reward": 0.604166679084301,
+      "rewards/format_reward": 0.0,
+      "step": 125
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 569.7708435058594,
+      "epoch": 1.3519999999999999,
+      "grad_norm": 0.6488833427429199,
+      "kl": 3.36163330078125,
+      "learning_rate": 2.0061633014861975e-06,
+      "loss": 0.3744,
+      "reward": 0.5625000204890966,
+      "reward_std": 0.40168892592191696,
+      "rewards/accuracy_reward": 0.5625000204890966,
+      "rewards/format_reward": 0.0,
+      "step": 126
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 544.5416870117188,
+      "epoch": 1.3626666666666667,
+      "grad_norm": 0.911536455154419,
+      "kl": 2.595245361328125,
+      "learning_rate": 1.988450850498758e-06,
+      "loss": 0.1733,
+      "reward": 0.5208333432674408,
+      "reward_std": 0.1705273948609829,
+      "rewards/accuracy_reward": 0.5208333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 127
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 692.4375152587891,
+      "epoch": 1.3733333333333333,
+      "grad_norm": 0.41733741760253906,
+      "kl": 1.950531005859375,
+      "learning_rate": 1.970661880825437e-06,
+      "loss": 0.138,
+      "reward": 0.29166667349636555,
+      "reward_std": 0.31314554810523987,
+      "rewards/accuracy_reward": 0.29166667349636555,
+      "rewards/format_reward": 0.0,
+      "step": 128
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 605.8541870117188,
+      "epoch": 1.384,
+      "grad_norm": 0.48099005222320557,
+      "kl": 1.07232666015625,
+      "learning_rate": 1.952799179212498e-06,
+      "loss": 0.1656,
+      "reward": 0.5833333469927311,
+      "reward_std": 0.4152076169848442,
+      "rewards/accuracy_reward": 0.5833333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 129
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 450.8333435058594,
+      "epoch": 1.3946666666666667,
+      "grad_norm": 1.2475286722183228,
+      "kl": 0.4179840087890625,
+      "learning_rate": 1.934865543956745e-06,
+      "loss": 0.0396,
+      "reward": 0.6666666772216558,
+      "reward_std": 0.3680921755731106,
+      "rewards/accuracy_reward": 0.6666666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 130
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 664.2708549499512,
+      "epoch": 1.4053333333333333,
+      "grad_norm": 1.1826740503311157,
+      "kl": 1.00048828125,
+      "learning_rate": 1.916863784467152e-06,
+      "loss": 0.1431,
+      "reward": 0.500000013038516,
+      "reward_std": 0.4326419048011303,
+      "rewards/accuracy_reward": 0.500000013038516,
+      "rewards/format_reward": 0.0,
+      "step": 131
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 739.3750228881836,
+      "epoch": 1.416,
+      "grad_norm": 1.281339168548584,
+      "kl": 1.94189453125,
+      "learning_rate": 1.8987967208247542e-06,
+      "loss": 0.2172,
+      "reward": 0.5625000093132257,
+      "reward_std": 0.43655749410390854,
+      "rewards/accuracy_reward": 0.5625000093132257,
+      "rewards/format_reward": 0.0,
+      "step": 132
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 547.166690826416,
+      "epoch": 1.4266666666666667,
+      "grad_norm": 0.7654764652252197,
+      "kl": 2.1027374267578125,
+      "learning_rate": 1.8806671833408633e-06,
+      "loss": 0.2557,
+      "reward": 0.7083333507180214,
+      "reward_std": 0.2861081548035145,
+      "rewards/accuracy_reward": 0.7083333507180214,
+      "rewards/format_reward": 0.0,
+      "step": 133
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 580.0833511352539,
+      "epoch": 1.4373333333333334,
+      "grad_norm": 1.4135154485702515,
+      "kl": 4.193695068359375,
+      "learning_rate": 1.8624780121136837e-06,
+      "loss": 0.3781,
+      "reward": 0.6250000111758709,
+      "reward_std": 0.3506578877568245,
+      "rewards/accuracy_reward": 0.6250000111758709,
+      "rewards/format_reward": 0.0,
+      "step": 134
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 592.7916870117188,
+      "epoch": 1.448,
+      "grad_norm": 3.319873094558716,
+      "kl": 4.4842376708984375,
+      "learning_rate": 1.8442320565833972e-06,
+      "loss": 0.2485,
+      "reward": 0.5833333469927311,
+      "reward_std": 0.2686738781630993,
+      "rewards/accuracy_reward": 0.5833333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 135
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 589.5416831970215,
+      "epoch": 1.4586666666666668,
+      "grad_norm": 1.7027231454849243,
+      "kl": 3.660552978515625,
+      "learning_rate": 1.8259321750857772e-06,
+      "loss": 0.2412,
+      "reward": 0.6041666772216558,
+      "reward_std": 0.2996268533170223,
+      "rewards/accuracy_reward": 0.6041666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 136
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 615.9791831970215,
+      "epoch": 1.4693333333333334,
+      "grad_norm": 0.6865854859352112,
+      "kl": 2.9984588623046875,
+      "learning_rate": 1.8075812344044163e-06,
+      "loss": 0.3306,
+      "reward": 0.5000000167638063,
+      "reward_std": 0.4893604479730129,
+      "rewards/accuracy_reward": 0.5000000167638063,
+      "rewards/format_reward": 0.0,
+      "step": 137
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 529.458345413208,
+      "epoch": 1.48,
+      "grad_norm": 0.3792395293712616,
+      "kl": 1.3320465087890625,
+      "learning_rate": 1.7891821093216275e-06,
+      "loss": 0.1501,
+      "reward": 0.6875000111758709,
+      "reward_std": 0.2350771203637123,
+      "rewards/accuracy_reward": 0.6875000111758709,
+      "rewards/format_reward": 0.0,
+      "step": 138
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 693.7708473205566,
+      "epoch": 1.4906666666666666,
+      "grad_norm": 0.5869024991989136,
+      "kl": 1.708587646484375,
+      "learning_rate": 1.770737682168091e-06,
+      "loss": 0.1275,
+      "reward": 0.5833333414047956,
+      "reward_std": 0.30354244634509087,
+      "rewards/accuracy_reward": 0.5833333414047956,
+      "rewards/format_reward": 0.0,
+      "step": 139
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 493.7291831970215,
+      "epoch": 1.5013333333333332,
+      "grad_norm": 0.7076964378356934,
+      "kl": 0.84979248046875,
+      "learning_rate": 1.7522508423713243e-06,
+      "loss": 0.1598,
+      "reward": 0.7083333432674408,
+      "reward_std": 0.24859581887722015,
+      "rewards/accuracy_reward": 0.7083333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 140
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 595.916690826416,
+      "epoch": 1.512,
+      "grad_norm": 1.1425637006759644,
+      "kl": 0.530303955078125,
+      "learning_rate": 1.7337244860030312e-06,
+      "loss": 0.1132,
+      "reward": 0.5833333432674408,
+      "reward_std": 0.25819891691207886,
+      "rewards/accuracy_reward": 0.5833333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 141
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 609.1041831970215,
+      "epoch": 1.5226666666666666,
+      "grad_norm": 0.8295446634292603,
+      "kl": 1.3100128173828125,
+      "learning_rate": 1.7151615153254192e-06,
+      "loss": 0.193,
+      "reward": 0.645833345130086,
+      "reward_std": 0.37377968430519104,
+      "rewards/accuracy_reward": 0.645833345130086,
+      "rewards/format_reward": 0.0,
+      "step": 142
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 575.9791870117188,
+      "epoch": 1.5333333333333332,
+      "grad_norm": 1.1884087324142456,
+      "kl": 1.55706787109375,
+      "learning_rate": 1.6965648383365416e-06,
+      "loss": 0.1849,
+      "reward": 0.7500000186264515,
+      "reward_std": 0.3506578840315342,
+      "rewards/accuracy_reward": 0.7500000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 143
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 521.9166831970215,
+      "epoch": 1.544,
+      "grad_norm": 0.3929750621318817,
+      "kl": 1.7193450927734375,
+      "learning_rate": 1.6779373683147423e-06,
+      "loss": 0.1654,
+      "reward": 0.6666666734963655,
+      "reward_std": 0.30354244634509087,
+      "rewards/accuracy_reward": 0.6666666734963655,
+      "rewards/format_reward": 0.0,
+      "step": 144
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 621.5625152587891,
+      "epoch": 1.5546666666666666,
+      "grad_norm": 0.9022945165634155,
+      "kl": 3.4391937255859375,
+      "learning_rate": 1.6592820233622743e-06,
+      "loss": 0.269,
+      "reward": 0.5416666883975267,
+      "reward_std": 0.3977733254432678,
+      "rewards/accuracy_reward": 0.5416666883975267,
+      "rewards/format_reward": 0.0,
+      "step": 145
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 637.4583435058594,
+      "epoch": 1.5653333333333332,
+      "grad_norm": 3.3849687576293945,
+      "kl": 6.964111328125,
+      "learning_rate": 1.6406017259481624e-06,
+      "loss": 0.5442,
+      "reward": 0.6041666809469461,
+      "reward_std": 0.41912321373820305,
+      "rewards/accuracy_reward": 0.6041666809469461,
+      "rewards/format_reward": 0.0,
+      "step": 146
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 619.9583587646484,
+      "epoch": 1.576,
+      "grad_norm": 1.0483787059783936,
+      "kl": 2.7703628540039062,
+      "learning_rate": 1.621899402450379e-06,
+      "loss": 0.2715,
+      "reward": 0.708333358168602,
+      "reward_std": 0.3332235962152481,
+      "rewards/accuracy_reward": 0.708333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 147
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 568.8750152587891,
+      "epoch": 1.5866666666666667,
+      "grad_norm": 1.4388060569763184,
+      "kl": 3.212860107421875,
+      "learning_rate": 1.6031779826974138e-06,
+      "loss": 0.2381,
+      "reward": 0.6458333358168602,
+      "reward_std": 0.19756478071212769,
+      "rewards/accuracy_reward": 0.6458333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 148
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 632.5416793823242,
+      "epoch": 1.5973333333333333,
+      "grad_norm": 0.9448708295822144,
+      "kl": 3.631927490234375,
+      "learning_rate": 1.5844403995092944e-06,
+      "loss": 0.3215,
+      "reward": 0.500000013038516,
+      "reward_std": 0.49719163402915,
+      "rewards/accuracy_reward": 0.500000013038516,
+      "rewards/format_reward": 0.0,
+      "step": 149
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 613.2500114440918,
+      "epoch": 1.608,
+      "grad_norm": 0.7443394660949707,
+      "kl": 2.036407470703125,
+      "learning_rate": 1.5656895882381465e-06,
+      "loss": 0.1928,
+      "reward": 0.5000000093132257,
+      "reward_std": 0.23116153106093407,
+      "rewards/accuracy_reward": 0.5000000093132257,
+      "rewards/format_reward": 0.0,
+      "step": 150
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 672.9583511352539,
+      "epoch": 1.6186666666666667,
+      "grad_norm": 0.6227521300315857,
+      "kl": 2.090057373046875,
+      "learning_rate": 1.5469284863083523e-06,
+      "loss": 0.1868,
+      "reward": 0.5625000167638063,
+      "reward_std": 0.34674230217933655,
+      "rewards/accuracy_reward": 0.5625000167638063,
+      "rewards/format_reward": 0.0,
+      "step": 151
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 473.458345413208,
+      "epoch": 1.6293333333333333,
+      "grad_norm": 1.2505531311035156,
+      "kl": 0.9958038330078125,
+      "learning_rate": 1.5281600327563853e-06,
+      "loss": 0.2341,
+      "reward": 0.7500000223517418,
+      "reward_std": 0.2686738707125187,
+      "rewards/accuracy_reward": 0.7500000223517418,
+      "rewards/format_reward": 0.0,
+      "step": 152
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 602.1250228881836,
+      "epoch": 1.6400000000000001,
+      "grad_norm": 0.2619081139564514,
+      "kl": 0.5251922607421875,
+      "learning_rate": 1.5093871677703947e-06,
+      "loss": 0.0727,
+      "reward": 0.6458333469927311,
+      "reward_std": 0.2996268533170223,
+      "rewards/accuracy_reward": 0.6458333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 153
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 698.1250190734863,
+      "epoch": 1.6506666666666665,
+      "grad_norm": 0.3198845386505127,
+      "kl": 1.6552276611328125,
+      "learning_rate": 1.4906128322296056e-06,
+      "loss": 0.1584,
+      "reward": 0.45833333767950535,
+      "reward_std": 0.31314554810523987,
+      "rewards/accuracy_reward": 0.45833333767950535,
+      "rewards/format_reward": 0.0,
+      "step": 154
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 616.2916793823242,
+      "epoch": 1.6613333333333333,
+      "grad_norm": 0.6926825046539307,
+      "kl": 1.0956878662109375,
+      "learning_rate": 1.471839967243615e-06,
+      "loss": 0.1899,
+      "reward": 0.43750000931322575,
+      "reward_std": 0.2996268570423126,
+      "rewards/accuracy_reward": 0.43750000931322575,
+      "rewards/format_reward": 0.0,
+      "step": 155
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 627.5208511352539,
+      "epoch": 1.6720000000000002,
+      "grad_norm": 0.5720909237861633,
+      "kl": 1.5909881591796875,
+      "learning_rate": 1.453071513691648e-06,
+      "loss": 0.115,
+      "reward": 0.4375000074505806,
+      "reward_std": 0.2621145099401474,
+      "rewards/accuracy_reward": 0.4375000074505806,
+      "rewards/format_reward": 0.0,
+      "step": 156
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 598.5625114440918,
+      "epoch": 1.6826666666666665,
+      "grad_norm": 1.2804419994354248,
+      "kl": 2.0683746337890625,
+      "learning_rate": 1.4343104117618536e-06,
+      "loss": 0.1199,
+      "reward": 0.5000000037252903,
+      "reward_std": 0.16661180183291435,
+      "rewards/accuracy_reward": 0.5000000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 157
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 465.5833435058594,
+      "epoch": 1.6933333333333334,
+      "grad_norm": 0.46794331073760986,
+      "kl": 1.4733734130859375,
+      "learning_rate": 1.415559600490706e-06,
+      "loss": 0.103,
+      "reward": 0.750000013038516,
+      "reward_std": 0.2861081548035145,
+      "rewards/accuracy_reward": 0.750000013038516,
+      "rewards/format_reward": 0.0,
+      "step": 158
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 636.3958549499512,
+      "epoch": 1.704,
+      "grad_norm": 0.5448073744773865,
+      "kl": 1.9481201171875,
+      "learning_rate": 1.3968220173025865e-06,
+      "loss": 0.1902,
+      "reward": 0.5000000037252903,
+      "reward_std": 0.19364918768405914,
+      "rewards/accuracy_reward": 0.5000000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 159
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 693.0833473205566,
+      "epoch": 1.7146666666666666,
+      "grad_norm": 0.9531773328781128,
+      "kl": 1.371734619140625,
+      "learning_rate": 1.3781005975496212e-06,
+      "loss": 0.1587,
+      "reward": 0.5208333469927311,
+      "reward_std": 0.3092299550771713,
+      "rewards/accuracy_reward": 0.5208333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 160
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 582.9583549499512,
+      "epoch": 1.7253333333333334,
+      "grad_norm": 0.9278988242149353,
+      "kl": 1.347991943359375,
+      "learning_rate": 1.359398274051838e-06,
+      "loss": 0.2244,
+      "reward": 0.4583333432674408,
+      "reward_std": 0.4248107150197029,
+      "rewards/accuracy_reward": 0.4583333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 161
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 503.5833396911621,
+      "epoch": 1.736,
+      "grad_norm": 0.31701967120170593,
+      "kl": 0.50714111328125,
+      "learning_rate": 1.340717976637726e-06,
+      "loss": 0.0906,
+      "reward": 0.7500000149011612,
+      "reward_std": 0.23116153106093407,
+      "rewards/accuracy_reward": 0.7500000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 162
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 508.7916793823242,
+      "epoch": 1.7466666666666666,
+      "grad_norm": 1.6528522968292236,
+      "kl": 0.932525634765625,
+      "learning_rate": 1.3220626316852584e-06,
+      "loss": 0.2641,
+      "reward": 0.7291666902601719,
+      "reward_std": 0.41129202768206596,
+      "rewards/accuracy_reward": 0.7291666902601719,
+      "rewards/format_reward": 0.0,
+      "step": 163
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 621.0625190734863,
+      "epoch": 1.7573333333333334,
+      "grad_norm": 0.36313435435295105,
+      "kl": 1.24462890625,
+      "learning_rate": 1.303435161663459e-06,
+      "loss": 0.0799,
+      "reward": 0.5416666846722364,
+      "reward_std": 0.2686738669872284,
+      "rewards/accuracy_reward": 0.5416666846722364,
+      "rewards/format_reward": 0.0,
+      "step": 164
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 525.8333435058594,
+      "epoch": 1.768,
+      "grad_norm": 0.7368504405021667,
+      "kl": 0.7718505859375,
+      "learning_rate": 1.2848384846745813e-06,
+      "loss": 0.1392,
+      "reward": 0.7500000204890966,
+      "reward_std": 0.333223607391119,
+      "rewards/accuracy_reward": 0.7500000204890966,
+      "rewards/format_reward": 0.0,
+      "step": 165
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 575.8125190734863,
+      "epoch": 1.7786666666666666,
+      "grad_norm": 0.316484659910202,
+      "kl": 0.6610794067382812,
+      "learning_rate": 1.2662755139969693e-06,
+      "loss": 0.1046,
+      "reward": 0.7291666716337204,
+      "reward_std": 0.3344954252243042,
+      "rewards/accuracy_reward": 0.7291666716337204,
+      "rewards/format_reward": 0.0,
+      "step": 166
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 553.8541793823242,
+      "epoch": 1.7893333333333334,
+      "grad_norm": 0.5553821325302124,
+      "kl": 1.3211212158203125,
+      "learning_rate": 1.2477491576286764e-06,
+      "loss": 0.0911,
+      "reward": 0.6666666772216558,
+      "reward_std": 0.23116153478622437,
+      "rewards/accuracy_reward": 0.6666666772216558,
+      "rewards/format_reward": 0.0,
+      "step": 167
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 554.4166851043701,
+      "epoch": 1.8,
+      "grad_norm": 0.6969295740127563,
+      "kl": 1.73931884765625,
+      "learning_rate": 1.2292623178319094e-06,
+      "loss": 0.1884,
+      "reward": 0.6458333376795053,
+      "reward_std": 0.3170611411333084,
+      "rewards/accuracy_reward": 0.6458333376795053,
+      "rewards/format_reward": 0.0,
+      "step": 168
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 548.1666870117188,
+      "epoch": 1.8106666666666666,
+      "grad_norm": 0.6913095116615295,
+      "kl": 1.616729736328125,
+      "learning_rate": 1.2108178906783732e-06,
+      "loss": 0.2069,
+      "reward": 0.5000000149011612,
+      "reward_std": 0.4701542444527149,
+      "rewards/accuracy_reward": 0.5000000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 169
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 618.916690826416,
+      "epoch": 1.8213333333333335,
+      "grad_norm": 0.9999569058418274,
+      "kl": 3.020416259765625,
+      "learning_rate": 1.1924187655955838e-06,
+      "loss": 0.3262,
+      "reward": 0.5000000149011612,
+      "reward_std": 0.4326419048011303,
+      "rewards/accuracy_reward": 0.5000000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 170
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 509.9791851043701,
+      "epoch": 1.8319999999999999,
+      "grad_norm": 0.4520111382007599,
+      "kl": 0.9397735595703125,
+      "learning_rate": 1.1740678249142229e-06,
+      "loss": 0.1257,
+      "reward": 0.6875000111758709,
+      "reward_std": 0.2350771240890026,
+      "rewards/accuracy_reward": 0.6875000111758709,
+      "rewards/format_reward": 0.0,
+      "step": 171
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 648.6875114440918,
+      "epoch": 1.8426666666666667,
+      "grad_norm": 1.3455004692077637,
+      "kl": 3.2716064453125,
+      "learning_rate": 1.1557679434166027e-06,
+      "loss": 0.2844,
+      "reward": 0.47916667349636555,
+      "reward_std": 0.2350771203637123,
+      "rewards/accuracy_reward": 0.47916667349636555,
+      "rewards/format_reward": 0.0,
+      "step": 172
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 433.7708377838135,
+      "epoch": 1.8533333333333335,
+      "grad_norm": 0.22631153464317322,
+      "kl": 0.7077484130859375,
+      "learning_rate": 1.1375219878863162e-06,
+      "loss": 0.034,
+      "reward": 0.7083333358168602,
+      "reward_std": 0.20148037374019623,
+      "rewards/accuracy_reward": 0.7083333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 173
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 619.8333473205566,
+      "epoch": 1.8639999999999999,
+      "grad_norm": 0.30306950211524963,
+      "kl": 1.13690185546875,
+      "learning_rate": 1.1193328166591372e-06,
+      "loss": 0.1185,
+      "reward": 0.4583333432674408,
+      "reward_std": 0.3776952773332596,
+      "rewards/accuracy_reward": 0.4583333432674408,
+      "rewards/format_reward": 0.0,
+      "step": 174
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 592.5208473205566,
+      "epoch": 1.8746666666666667,
+      "grad_norm": 0.39582908153533936,
+      "kl": 1.418182373046875,
+      "learning_rate": 1.1012032791752457e-06,
+      "loss": 0.1603,
+      "reward": 0.4791666753590107,
+      "reward_std": 0.36417657881975174,
+      "rewards/accuracy_reward": 0.4791666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 175
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 501.5000114440918,
+      "epoch": 1.8853333333333333,
+      "grad_norm": 0.9711954593658447,
+      "kl": 0.571441650390625,
+      "learning_rate": 1.0831362155328477e-06,
+      "loss": 0.166,
+      "reward": 0.833333358168602,
+      "reward_std": 0.3332235962152481,
+      "rewards/accuracy_reward": 0.833333358168602,
+      "rewards/format_reward": 0.0,
+      "step": 176
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 570.8541831970215,
+      "epoch": 1.896,
+      "grad_norm": 0.6740026473999023,
+      "kl": 1.3124847412109375,
+      "learning_rate": 1.0651344560432548e-06,
+      "loss": 0.1433,
+      "reward": 0.645833345130086,
+      "reward_std": 0.36417658254504204,
+      "rewards/accuracy_reward": 0.645833345130086,
+      "rewards/format_reward": 0.0,
+      "step": 177
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 560.0416870117188,
+      "epoch": 1.9066666666666667,
+      "grad_norm": 0.29950958490371704,
+      "kl": 1.2054367065429688,
+      "learning_rate": 1.0472008207875023e-06,
+      "loss": 0.1039,
+      "reward": 0.6041666753590107,
+      "reward_std": 0.1801304966211319,
+      "rewards/accuracy_reward": 0.6041666753590107,
+      "rewards/format_reward": 0.0,
+      "step": 178
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 551.6875228881836,
+      "epoch": 1.9173333333333333,
+      "grad_norm": 0.4891541302204132,
+      "kl": 1.02886962890625,
+      "learning_rate": 1.0293381191745637e-06,
+      "loss": 0.17,
+      "reward": 0.5833333469927311,
+      "reward_std": 0.3332235999405384,
+      "rewards/accuracy_reward": 0.5833333469927311,
+      "rewards/format_reward": 0.0,
+      "step": 179
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 604.8958473205566,
+      "epoch": 1.928,
+      "grad_norm": 0.6707997918128967,
+      "kl": 1.073486328125,
+      "learning_rate": 1.0115491495012424e-06,
+      "loss": 0.1305,
+      "reward": 0.5625000149011612,
+      "reward_std": 0.299626849591732,
+      "rewards/accuracy_reward": 0.5625000149011612,
+      "rewards/format_reward": 0.0,
+      "step": 180
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 538.8750152587891,
+      "epoch": 1.9386666666666668,
+      "grad_norm": 0.7848038673400879,
+      "kl": 1.063873291015625,
+      "learning_rate": 9.938366985138028e-07,
+      "loss": 0.1507,
+      "reward": 0.6250000167638063,
+      "reward_std": 0.2686738707125187,
+      "rewards/accuracy_reward": 0.6250000167638063,
+      "rewards/format_reward": 0.0,
+      "step": 181
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 523.9166812896729,
+      "epoch": 1.9493333333333334,
+      "grad_norm": 0.38517096638679504,
+      "kl": 1.3513946533203125,
+      "learning_rate": 9.762035409714077e-07,
+      "loss": 0.141,
+      "reward": 0.583333345130086,
+      "reward_std": 0.2957112565636635,
+      "rewards/accuracy_reward": 0.583333345130086,
+      "rewards/format_reward": 0.0,
+      "step": 182
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 481.4791717529297,
+      "epoch": 1.96,
+      "grad_norm": 0.7527852058410645,
+      "kl": 1.50286865234375,
+      "learning_rate": 9.586524392114385e-07,
+      "loss": 0.1278,
+      "reward": 0.6250000037252903,
+      "reward_std": 0.26603010296821594,
+      "rewards/accuracy_reward": 0.6250000037252903,
+      "rewards/format_reward": 0.0,
+      "step": 183
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 446.2083435058594,
+      "epoch": 1.9706666666666668,
+      "grad_norm": 0.22020579874515533,
+      "kl": 0.7943267822265625,
+      "learning_rate": 9.411861427167567e-07,
+      "loss": 0.0756,
+      "reward": 0.729166679084301,
+      "reward_std": 0.2350771240890026,
+      "rewards/accuracy_reward": 0.729166679084301,
+      "rewards/format_reward": 0.0,
+      "step": 184
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 653.5625114440918,
+      "epoch": 1.9813333333333332,
+      "grad_norm": 0.4303428828716278,
+      "kl": 2.3649749755859375,
+      "learning_rate": 9.238073876849837e-07,
+      "loss": 0.2236,
+      "reward": 0.5416666883975267,
+      "reward_std": 0.3881702274084091,
+      "rewards/accuracy_reward": 0.5416666883975267,
+      "rewards/format_reward": 0.0,
+      "step": 185
+    },
+    {
+      "clip_ratio": 0.0,
+      "completion_length": 599.3750152587891,
+      "epoch": 1.992,
+      "grad_norm": 2.4530301094055176,
+      "kl": 3.55426025390625,
+      "learning_rate": 9.065188965998574e-07,
+      "loss": 0.3119,
+      "reward": 0.6250000186264515,
+      "reward_std": 0.3602609895169735,
+      "rewards/accuracy_reward": 0.6250000186264515,
+      "rewards/format_reward": 0.0,
+      "step": 186
+    }
+  ],
+  "logging_steps": 1,
+  "max_steps": 279,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 3,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.0,
+  "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:0fdd25831a6977850bf19f892f870a0f9dde2616344ff9cb57c3ec17851adc79
+size 7672

zero_to_fp32.py

@@ -0,0 +1,674 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+#   python zero_to_fp32.py . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import json
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device)
+
+        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, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+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 tqdm(param_shapes.items(), desc='Gathering Sharded Weights'):
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # 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,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None, exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+
+    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, exclude_frozen_parameters)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag, exclude_frozen_parameters)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        state_dict_split = split_torch_state_dict_into_shards(state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard = {tensor: state_dict[tensor].contiguous() for tensor in tensors}
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard, output_path)
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument("output_dir",
+                        type=str,
+                        help="directory to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+                                               args.output_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)