Upload folder using huggingface_hub

checkpoint-2000/added_tokens.json

@@ -0,0 +1,11 @@
+{
+  "</box>": 92552,
+  "</img>": 92545,
+  "</quad>": 92548,
+  "</ref>": 92550,
+  "<IMG_CONTEXT>": 92546,
+  "<box>": 92551,
+  "<img>": 92544,
+  "<quad>": 92547,
+  "<ref>": 92549
+}

checkpoint-2000/chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{'<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>' + '
+'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

checkpoint-2000/config.json

@@ -0,0 +1,188 @@
+{
+  "action_chunk_size": 60,
+  "action_config": {
+    "_attn_implementation_autoset": true,
+    "action_chunk_size": 60,
+    "action_dim": 26,
+    "architectures": [
+      "ActionExpertModel"
+    ],
+    "attn_implementation": "eager",
+    "auto_map": {
+      "AutoConfig": "configuration_action_expert.ActionExpertConfig",
+      "AutoModel": "modeling_action_expert.ActionExpertModel"
+    },
+    "bias": false,
+    "dtype": "bfloat16",
+    "head_dim": 64,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "input_hidden_size": 2048,
+    "intermediate_size": 4096,
+    "max_position_embeddings": 32768,
+    "model_type": "action_expert",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 8,
+    "pad_token_id": 2,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": {
+      "factor": 2.0,
+      "type": "dynamic"
+    },
+    "rope_theta": 1000000,
+    "state_dim": 26,
+    "state_token_num": 3,
+    "use_bfloat16": true,
+    "use_cache": true,
+    "use_flash_attn": false
+  },
+  "architectures": [
+    "GO1Model"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_go1.GO1ModelConfig",
+    "AutoModel": "modeling_go1.GO1Model",
+    "AutoModelForCausalLM": "modeling_internlm2_go1.py.InternLM2ForCausalLMGO1"
+  },
+  "bos_token_id": 1,
+  "downsample_ratio": 0.5,
+  "dtype": "bfloat16",
+  "dynamic_image_size": false,
+  "eos_token_id": 2,
+  "flow_matching": {
+    "flow_matching_final_weight": 10.0,
+    "flow_matching_weight": 1.0,
+    "num_steps": 10,
+    "rng": 42
+  },
+  "force_image_size": 448,
+  "img_context_token_id": 92546,
+  "information_fusion_config": {
+    "action_chunk_size": 60,
+    "action_dim": 26,
+    "attn_implementation": "eager",
+    "bias": true,
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "input_hidden_size": 2048,
+    "intermediate_size": 11008,
+    "max_position_embeddings": 2048,
+    "model_type": "information_fusion",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 12,
+    "num_key_value_heads": 16,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 10000,
+    "state_dim": 26,
+    "use_cache": false
+  },
+  "initializer_range": 0.02,
+  "latent_planner_config": {
+    "action_dim": 1,
+    "attn_implementation": "eager",
+    "bias": false,
+    "head_dim": 64,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "input_hidden_size": 2048,
+    "intermediate_size": 2048,
+    "max_position_embeddings": 2048,
+    "model_type": "intermidiate_action_expert",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 10000,
+    "state_token_num": 0,
+    "use_cache": true,
+    "vocab_size": 32
+  },
+  "latent_planning": false,
+  "llm_config": {
+    "_attn_implementation_autoset": true,
+    "architectures": [
+      "InternLM2ForCausalLMGO1"
+    ],
+    "attn_implementation": "flash_attention_2",
+    "auto_map": {
+      "AutoConfig": "configuration_internlm2.InternLM2Config",
+      "AutoModel": "modeling_internlm2_go1.InternLM2ForCausalLMGO1",
+      "AutoModelForCausalLM": "modeling_internlm2_go1.py.InternLM2ForCausalLMGO1"
+    },
+    "bias": false,
+    "dtype": "bfloat16",
+    "hidden_act": "silu",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 8192,
+    "max_position_embeddings": 32768,
+    "model_type": "internlm2",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 8,
+    "pad_token_id": 2,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": {
+      "factor": 2.0,
+      "type": "dynamic"
+    },
+    "rope_theta": 1000000,
+    "use_bfloat16": true,
+    "use_cache": true,
+    "vocab_size": 92553
+  },
+  "max_dynamic_patch": 6,
+  "min_dynamic_patch": 1,
+  "model_type": "go1",
+  "noise_scheduler_config": {
+    "beta_schedule": "squaredcos_cap_v2",
+    "clip_sample": false,
+    "num_inference_timesteps": 5,
+    "num_train_timesteps": 1000,
+    "prediction_type": "sample"
+  },
+  "norm": true,
+  "output_attentions": false,
+  "pad2square": false,
+  "pad_token_id": 2,
+  "ps_version": "v2",
+  "select_layer": -1,
+  "template": "internlm2-chat",
+  "transformers_version": null,
+  "use_backbone_lora": 0,
+  "use_llm_lora": 0,
+  "use_thumbnail": false,
+  "vision_config": {
+    "_attn_implementation_autoset": true,
+    "architectures": [
+      "InternVisionModel"
+    ],
+    "attention_dropout": 0.0,
+    "drop_path_rate": 0.1,
+    "dropout": 0.0,
+    "dtype": "bfloat16",
+    "hidden_act": "gelu",
+    "hidden_size": 1024,
+    "image_size": 448,
+    "initializer_factor": 1.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "layer_norm_eps": 1e-06,
+    "model_type": "intern_vit_6b",
+    "norm_type": "layer_norm",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "patch_size": 14,
+    "qk_normalization": false,
+    "qkv_bias": true,
+    "use_bfloat16": true,
+    "use_flash_attn": true
+  }
+}

checkpoint-2000/global_step2000/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt

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

checkpoint-2000/global_step2000/mp_rank_00_model_states.pt

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

checkpoint-2000/latest

@@ -0,0 +1 @@
+global_step2000

checkpoint-2000/model-00001-of-00002.safetensors

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

checkpoint-2000/model-00002-of-00002.safetensors

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

checkpoint-2000/model.safetensors.index.json

@@ -0,0 +1,815 @@
+{
+  "metadata": {
+    "total_parameters": 2591644698,
+    "total_size": 5183289396
+  },
+  "weight_map": {
+    "action_adaptor.0.bias": "model-00002-of-00002.safetensors",
+    "action_adaptor.0.weight": "model-00002-of-00002.safetensors",
+    "action_adaptor.2.bias": "model-00002-of-00002.safetensors",
+    "action_adaptor.2.weight": "model-00002-of-00002.safetensors",
+    "action_adaptor.4.bias": "model-00002-of-00002.safetensors",
+    "action_adaptor.4.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.0.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.0.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.0.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.0.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.0.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.0.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.0.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.1.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.10.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.11.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.12.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.13.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.14.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.15.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.16.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.17.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.18.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.18.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.18.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.18.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.18.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.18.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.18.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.attention.wqkv.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.19.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.2.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.2.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.2.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.2.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.2.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.2.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.2.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.20.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.20.attention.wqkv.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.20.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.20.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.20.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.20.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.20.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.attention.wqkv.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.21.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.attention.wqkv.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.22.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.attention.wo.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.attention.wqkv.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.attention_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.feed_forward.w1.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.feed_forward.w2.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.feed_forward.w3.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.23.ffn_norm.weight": "model-00002-of-00002.safetensors",
+    "action_model.layers.3.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.3.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.3.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.3.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.3.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.3.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.3.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.4.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.5.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.6.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.7.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.8.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "action_model.layers.9.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "action_model.norm.weight": "model-00002-of-00002.safetensors",
+    "final_layer.ffn_final.fc1.bias": "model-00002-of-00002.safetensors",
+    "final_layer.ffn_final.fc1.weight": "model-00002-of-00002.safetensors",
+    "final_layer.ffn_final.fc2.bias": "model-00002-of-00002.safetensors",
+    "final_layer.ffn_final.fc2.weight": "model-00002-of-00002.safetensors",
+    "final_layer.norm_final.weight": "model-00002-of-00002.safetensors",
+    "freq_embedder.mlp.0.bias": "model-00002-of-00002.safetensors",
+    "freq_embedder.mlp.0.weight": "model-00002-of-00002.safetensors",
+    "freq_embedder.mlp.2.bias": "model-00002-of-00002.safetensors",
+    "freq_embedder.mlp.2.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.0.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.0.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.1.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.1.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.10.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.10.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.11.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.11.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.12.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.12.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.13.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.13.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.14.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.14.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.15.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.15.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.16.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.16.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.17.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.17.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.18.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.18.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.19.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.19.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.2.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.2.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.20.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.20.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.21.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.21.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.22.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.22.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.23.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.23.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.3.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.3.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.4.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.4.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.5.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.5.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.6.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.6.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.7.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.7.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.8.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.8.weight": "model-00002-of-00002.safetensors",
+    "k_proj_layers.9.bias": "model-00002-of-00002.safetensors",
+    "k_proj_layers.9.weight": "model-00002-of-00002.safetensors",
+    "language_model.model.layers.0.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.0.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.0.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.0.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.0.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.0.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.0.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.1.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.10.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.11.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.12.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.13.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.14.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.15.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.16.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.17.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.18.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.19.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.2.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.20.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.21.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.22.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.23.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.3.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.4.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.5.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.6.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.7.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.8.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.attention.wo.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.attention.wqkv.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.attention_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.feed_forward.w1.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.feed_forward.w2.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.feed_forward.w3.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.layers.9.ffn_norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.norm.weight": "model-00001-of-00002.safetensors",
+    "language_model.model.tok_embeddings.weight": "model-00001-of-00002.safetensors",
+    "language_model.output.weight": "model-00001-of-00002.safetensors",
+    "mlp1.0.bias": "model-00001-of-00002.safetensors",
+    "mlp1.0.weight": "model-00001-of-00002.safetensors",
+    "mlp1.1.bias": "model-00001-of-00002.safetensors",
+    "mlp1.1.weight": "model-00001-of-00002.safetensors",
+    "mlp1.3.bias": "model-00001-of-00002.safetensors",
+    "mlp1.3.weight": "model-00001-of-00002.safetensors",
+    "state_adaptor.0.bias": "model-00002-of-00002.safetensors",
+    "state_adaptor.0.weight": "model-00002-of-00002.safetensors",
+    "state_adaptor.2.bias": "model-00002-of-00002.safetensors",
+    "state_adaptor.2.weight": "model-00002-of-00002.safetensors",
+    "state_adaptor.4.bias": "model-00002-of-00002.safetensors",
+    "state_adaptor.4.weight": "model-00002-of-00002.safetensors",
+    "time_embedder.mlp.0.bias": "model-00002-of-00002.safetensors",
+    "time_embedder.mlp.0.weight": "model-00002-of-00002.safetensors",
+    "time_embedder.mlp.2.bias": "model-00002-of-00002.safetensors",
+    "time_embedder.mlp.2.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.0.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.0.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.1.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.1.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.10.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.10.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.11.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.11.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.12.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.12.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.13.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.13.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.14.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.14.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.15.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.15.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.16.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.16.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.17.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.17.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.18.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.18.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.19.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.19.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.2.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.2.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.20.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.20.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.21.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.21.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.22.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.22.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.23.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.23.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.3.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.3.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.4.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.4.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.5.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.5.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.6.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.6.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.7.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.7.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.8.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.8.weight": "model-00002-of-00002.safetensors",
+    "v_proj_layers.9.bias": "model-00002-of-00002.safetensors",
+    "v_proj_layers.9.weight": "model-00002-of-00002.safetensors",
+    "vision_model.embeddings.class_embedding": "model-00001-of-00002.safetensors",
+    "vision_model.embeddings.patch_embedding.bias": "model-00001-of-00002.safetensors",
+    "vision_model.embeddings.patch_embedding.weight": "model-00001-of-00002.safetensors",
+    "vision_model.embeddings.position_embedding": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.0.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.1.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.10.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.11.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.12.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.13.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.14.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.15.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.16.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.17.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.18.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.19.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.2.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.20.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.21.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.22.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.23.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.3.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.4.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.5.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.6.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.7.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.8.norm2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.attn.proj.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.attn.proj.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.attn.qkv.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.attn.qkv.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.ls1": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.ls2": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc2.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.norm1.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.norm1.weight": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.norm2.bias": "model-00001-of-00002.safetensors",
+    "vision_model.encoder.layers.9.norm2.weight": "model-00001-of-00002.safetensors"
+  }
+}

checkpoint-2000/rng_state.pth

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

checkpoint-2000/special_tokens_map.json

@@ -0,0 +1,47 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|action_start|>",
+    "<|action_end|>",
+    "<|interpreter|>",
+    "<|plugin|>",
+    "<img>",
+    "</img>",
+    "<IMG_CONTEXT>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>"
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint-2000/tokenization_internlm2.py

@@ -0,0 +1,235 @@
+# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on transformers/src/transformers/models/llama/tokenization_llama.py
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tokenization classes for InternLM."""
+import os
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {'vocab_file': './tokenizer.model'}
+
+PRETRAINED_VOCAB_FILES_MAP = {}
+
+
+# Modified from transformers.model.llama.tokenization_llama.LlamaTokenizer
+class InternLM2Tokenizer(PreTrainedTokenizer):
+    """
+    Construct a InternLM2 tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    model_input_names = ['input_ids', 'attention_mask']
+    _auto_class = 'AutoTokenizer'
+
+    def __init__(
+        self,
+        vocab_file,
+        unk_token='<unk>',
+        bos_token='<s>',
+        eos_token='</s>',
+        pad_token='</s>',
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        add_bos_token=True,
+        add_eos_token=False,
+        decode_with_prefix_space=False,
+        clean_up_tokenization_spaces=False,
+        **kwargs,
+    ):
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+        self.vocab_file = vocab_file
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+        self.decode_with_prefix_space = decode_with_prefix_space
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+        self._no_prefix_space_tokens = None
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    @property
+    def no_prefix_space_tokens(self):
+        if self._no_prefix_space_tokens is None:
+            vocab = self.convert_ids_to_tokens(list(range(self.vocab_size)))
+            self._no_prefix_space_tokens = {i for i, tok in enumerate(vocab) if not tok.startswith('▁')}
+        return self._no_prefix_space_tokens
+
+    @property
+    def vocab_size(self):
+        """Returns vocab size"""
+        return self.sp_model.get_piece_size()
+
+    @property
+    def bos_token_id(self) -> Optional[int]:
+        return self.sp_model.bos_id()
+
+    @property
+    def eos_token_id(self) -> Optional[int]:
+        return self.sp_model.eos_id()
+
+    def get_vocab(self):
+        """Returns vocab as a dict"""
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text):
+        """Returns a tokenized string."""
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        token = self.sp_model.IdToPiece(index)
+        return token
+
+    def _maybe_add_prefix_space(self, tokens, decoded):
+        if tokens and tokens[0] not in self.no_prefix_space_tokens:
+            return ' ' + decoded
+        else:
+            return decoded
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        current_sub_tokens = []
+        out_string = ''
+        prev_is_special = False
+        for token in tokens:
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                if not prev_is_special:
+                    out_string += ' '
+                out_string += self.sp_model.decode(current_sub_tokens) + token
+                prev_is_special = True
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+                prev_is_special = False
+        out_string += self.sp_model.decode(current_sub_tokens)
+        out_string = self.clean_up_tokenization(out_string)
+        out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string)
+        return out_string[1:]
+
+    def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if not os.path.isdir(save_directory):
+            logger.error(f'Vocabulary path ({save_directory}) should be a directory')
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, 'wb') as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        if self.add_bos_token:
+            bos_token_ids = [self.bos_token_id]
+        else:
+            bos_token_ids = []
+
+        output = bos_token_ids + token_ids_0
+
+        if token_ids_1 is not None:
+            output = output + token_ids_1
+
+        if self.add_eos_token:
+            output = output + [self.eos_token_id]
+
+        return output
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        if token_ids_1 is None:
+            return [1] + ([0] * len(token_ids_0)) + [1]
+        return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
+        use of token type ids, therefore a list of zeros is returned.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of zeros.
+        """
+        eos = [self.eos_token_id]
+
+        if token_ids_1 is None:
+            return len(token_ids_0 + eos) * [0]
+        return len(token_ids_0 + eos + token_ids_1 + eos) * [0]

checkpoint-2000/tokenizer.model

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

checkpoint-2000/tokenizer_config.json

@@ -0,0 +1,179 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92538": {
+      "content": "<|plugin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92539": {
+      "content": "<|interpreter|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92540": {
+      "content": "<|action_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92541": {
+      "content": "<|action_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92542": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92543": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92544": {
+      "content": "<img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92545": {
+      "content": "</img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92546": {
+      "content": "<IMG_CONTEXT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92547": {
+      "content": "<quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92548": {
+      "content": "</quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92549": {
+      "content": "<ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92550": {
+      "content": "</ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92551": {
+      "content": "<box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "92552": {
+      "content": "</box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|action_start|>",
+    "<|action_end|>",
+    "<|interpreter|>",
+    "<|plugin|>",
+    "<img>",
+    "</img>",
+    "<IMG_CONTEXT>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_internlm2.InternLM2Tokenizer",
+      null
+    ]
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "model_max_length": 4096,
+  "pad_token": "</s>",
+  "tokenizer_class": "InternLM2Tokenizer",
+  "unk_token": "<unk>"
+}

checkpoint-2000/trainer_state.json

@@ -0,0 +1,1434 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 2.7181719260065287,
+  "eval_steps": 500,
+  "global_step": 2000,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.013601741022850925,
+      "grad_norm": 2.812274217605591,
+      "learning_rate": 2.0000000000000003e-06,
+      "loss": 5.2086,
+      "step": 10
+    },
+    {
+      "epoch": 0.02720348204570185,
+      "grad_norm": 2.162468433380127,
+      "learning_rate": 4.000000000000001e-06,
+      "loss": 5.2051,
+      "step": 20
+    },
+    {
+      "epoch": 0.040805223068552776,
+      "grad_norm": 1.7208396196365356,
+      "learning_rate": 6e-06,
+      "loss": 5.2105,
+      "step": 30
+    },
+    {
+      "epoch": 0.0544069640914037,
+      "grad_norm": 2.556692361831665,
+      "learning_rate": 8.000000000000001e-06,
+      "loss": 5.1625,
+      "step": 40
+    },
+    {
+      "epoch": 0.06800870511425462,
+      "grad_norm": 8.235758781433105,
+      "learning_rate": 1e-05,
+      "loss": 4.6379,
+      "step": 50
+    },
+    {
+      "epoch": 0.08161044613710555,
+      "grad_norm": 10.163893699645996,
+      "learning_rate": 1.2e-05,
+      "loss": 3.8391,
+      "step": 60
+    },
+    {
+      "epoch": 0.09521218715995647,
+      "grad_norm": 11.116437911987305,
+      "learning_rate": 1.4000000000000001e-05,
+      "loss": 3.0016,
+      "step": 70
+    },
+    {
+      "epoch": 0.1088139281828074,
+      "grad_norm": 10.696512222290039,
+      "learning_rate": 1.6000000000000003e-05,
+      "loss": 2.4004,
+      "step": 80
+    },
+    {
+      "epoch": 0.12241566920565833,
+      "grad_norm": 12.367400169372559,
+      "learning_rate": 1.8e-05,
+      "loss": 1.8969,
+      "step": 90
+    },
+    {
+      "epoch": 0.13601741022850924,
+      "grad_norm": 13.723505973815918,
+      "learning_rate": 2e-05,
+      "loss": 1.4616,
+      "step": 100
+    },
+    {
+      "epoch": 0.14961915125136016,
+      "grad_norm": 22.034290313720703,
+      "learning_rate": 2.2000000000000003e-05,
+      "loss": 1.2379,
+      "step": 110
+    },
+    {
+      "epoch": 0.1632208922742111,
+      "grad_norm": 8.55447769165039,
+      "learning_rate": 2.4e-05,
+      "loss": 1.0951,
+      "step": 120
+    },
+    {
+      "epoch": 0.17682263329706202,
+      "grad_norm": 11.120705604553223,
+      "learning_rate": 2.6000000000000002e-05,
+      "loss": 0.9654,
+      "step": 130
+    },
+    {
+      "epoch": 0.19042437431991294,
+      "grad_norm": 13.087006568908691,
+      "learning_rate": 2.8000000000000003e-05,
+      "loss": 0.8391,
+      "step": 140
+    },
+    {
+      "epoch": 0.20402611534276388,
+      "grad_norm": 12.556931495666504,
+      "learning_rate": 3e-05,
+      "loss": 0.7075,
+      "step": 150
+    },
+    {
+      "epoch": 0.2176278563656148,
+      "grad_norm": 14.245359420776367,
+      "learning_rate": 3.2000000000000005e-05,
+      "loss": 0.6853,
+      "step": 160
+    },
+    {
+      "epoch": 0.2312295973884657,
+      "grad_norm": 12.10800838470459,
+      "learning_rate": 3.4000000000000007e-05,
+      "loss": 0.606,
+      "step": 170
+    },
+    {
+      "epoch": 0.24483133841131666,
+      "grad_norm": 11.180713653564453,
+      "learning_rate": 3.6e-05,
+      "loss": 0.5734,
+      "step": 180
+    },
+    {
+      "epoch": 0.2584330794341676,
+      "grad_norm": 12.638975143432617,
+      "learning_rate": 3.8e-05,
+      "loss": 0.545,
+      "step": 190
+    },
+    {
+      "epoch": 0.2720348204570185,
+      "grad_norm": 11.650677680969238,
+      "learning_rate": 4e-05,
+      "loss": 0.5162,
+      "step": 200
+    },
+    {
+      "epoch": 0.28563656147986943,
+      "grad_norm": 11.61837387084961,
+      "learning_rate": 4.2e-05,
+      "loss": 0.491,
+      "step": 210
+    },
+    {
+      "epoch": 0.2992383025027203,
+      "grad_norm": 5.453939914703369,
+      "learning_rate": 4.4000000000000006e-05,
+      "loss": 0.4378,
+      "step": 220
+    },
+    {
+      "epoch": 0.31284004352557127,
+      "grad_norm": 7.908788204193115,
+      "learning_rate": 4.600000000000001e-05,
+      "loss": 0.4299,
+      "step": 230
+    },
+    {
+      "epoch": 0.3264417845484222,
+      "grad_norm": 13.816142082214355,
+      "learning_rate": 4.8e-05,
+      "loss": 0.4113,
+      "step": 240
+    },
+    {
+      "epoch": 0.3400435255712731,
+      "grad_norm": 11.430989265441895,
+      "learning_rate": 5e-05,
+      "loss": 0.4371,
+      "step": 250
+    },
+    {
+      "epoch": 0.35364526659412404,
+      "grad_norm": 6.766596794128418,
+      "learning_rate": 5.2000000000000004e-05,
+      "loss": 0.387,
+      "step": 260
+    },
+    {
+      "epoch": 0.367247007616975,
+      "grad_norm": 10.817462921142578,
+      "learning_rate": 5.4000000000000005e-05,
+      "loss": 0.3643,
+      "step": 270
+    },
+    {
+      "epoch": 0.3808487486398259,
+      "grad_norm": 6.278716564178467,
+      "learning_rate": 5.6000000000000006e-05,
+      "loss": 0.3477,
+      "step": 280
+    },
+    {
+      "epoch": 0.3944504896626768,
+      "grad_norm": 13.244380950927734,
+      "learning_rate": 5.8e-05,
+      "loss": 0.3763,
+      "step": 290
+    },
+    {
+      "epoch": 0.40805223068552776,
+      "grad_norm": 4.513805866241455,
+      "learning_rate": 6e-05,
+      "loss": 0.3093,
+      "step": 300
+    },
+    {
+      "epoch": 0.42165397170837865,
+      "grad_norm": 4.978238582611084,
+      "learning_rate": 6.2e-05,
+      "loss": 0.306,
+      "step": 310
+    },
+    {
+      "epoch": 0.4352557127312296,
+      "grad_norm": 9.654162406921387,
+      "learning_rate": 6.400000000000001e-05,
+      "loss": 0.3157,
+      "step": 320
+    },
+    {
+      "epoch": 0.44885745375408054,
+      "grad_norm": 5.806159496307373,
+      "learning_rate": 6.6e-05,
+      "loss": 0.302,
+      "step": 330
+    },
+    {
+      "epoch": 0.4624591947769314,
+      "grad_norm": 6.283710479736328,
+      "learning_rate": 6.800000000000001e-05,
+      "loss": 0.2885,
+      "step": 340
+    },
+    {
+      "epoch": 0.47606093579978237,
+      "grad_norm": 6.1870951652526855,
+      "learning_rate": 7e-05,
+      "loss": 0.289,
+      "step": 350
+    },
+    {
+      "epoch": 0.4896626768226333,
+      "grad_norm": 3.082080364227295,
+      "learning_rate": 7.2e-05,
+      "loss": 0.2584,
+      "step": 360
+    },
+    {
+      "epoch": 0.5032644178454843,
+      "grad_norm": 5.254792213439941,
+      "learning_rate": 7.4e-05,
+      "loss": 0.2752,
+      "step": 370
+    },
+    {
+      "epoch": 0.5168661588683352,
+      "grad_norm": 6.462503433227539,
+      "learning_rate": 7.6e-05,
+      "loss": 0.2639,
+      "step": 380
+    },
+    {
+      "epoch": 0.530467899891186,
+      "grad_norm": 10.89343547821045,
+      "learning_rate": 7.800000000000001e-05,
+      "loss": 0.2893,
+      "step": 390
+    },
+    {
+      "epoch": 0.544069640914037,
+      "grad_norm": 5.2178192138671875,
+      "learning_rate": 8e-05,
+      "loss": 0.2415,
+      "step": 400
+    },
+    {
+      "epoch": 0.5576713819368879,
+      "grad_norm": 8.687518119812012,
+      "learning_rate": 8.2e-05,
+      "loss": 0.2385,
+      "step": 410
+    },
+    {
+      "epoch": 0.5712731229597389,
+      "grad_norm": 4.670180320739746,
+      "learning_rate": 8.4e-05,
+      "loss": 0.2297,
+      "step": 420
+    },
+    {
+      "epoch": 0.5848748639825898,
+      "grad_norm": 3.699598550796509,
+      "learning_rate": 8.6e-05,
+      "loss": 0.2241,
+      "step": 430
+    },
+    {
+      "epoch": 0.5984766050054406,
+      "grad_norm": 3.260232925415039,
+      "learning_rate": 8.800000000000001e-05,
+      "loss": 0.216,
+      "step": 440
+    },
+    {
+      "epoch": 0.6120783460282916,
+      "grad_norm": 3.4559569358825684,
+      "learning_rate": 9e-05,
+      "loss": 0.2275,
+      "step": 450
+    },
+    {
+      "epoch": 0.6256800870511425,
+      "grad_norm": 3.7167000770568848,
+      "learning_rate": 9.200000000000001e-05,
+      "loss": 0.2133,
+      "step": 460
+    },
+    {
+      "epoch": 0.6392818280739935,
+      "grad_norm": 4.1776628494262695,
+      "learning_rate": 9.4e-05,
+      "loss": 0.2194,
+      "step": 470
+    },
+    {
+      "epoch": 0.6528835690968444,
+      "grad_norm": 4.609129428863525,
+      "learning_rate": 9.6e-05,
+      "loss": 0.2088,
+      "step": 480
+    },
+    {
+      "epoch": 0.6664853101196954,
+      "grad_norm": 5.062413215637207,
+      "learning_rate": 9.8e-05,
+      "loss": 0.2185,
+      "step": 490
+    },
+    {
+      "epoch": 0.6800870511425462,
+      "grad_norm": 3.3144760131835938,
+      "learning_rate": 0.0001,
+      "loss": 0.2176,
+      "step": 500
+    },
+    {
+      "epoch": 0.6936887921653971,
+      "grad_norm": 2.5662150382995605,
+      "learning_rate": 0.00010200000000000001,
+      "loss": 0.2108,
+      "step": 510
+    },
+    {
+      "epoch": 0.7072905331882481,
+      "grad_norm": 2.718777656555176,
+      "learning_rate": 0.00010400000000000001,
+      "loss": 0.1734,
+      "step": 520
+    },
+    {
+      "epoch": 0.720892274211099,
+      "grad_norm": 2.934107780456543,
+      "learning_rate": 0.00010600000000000002,
+      "loss": 0.209,
+      "step": 530
+    },
+    {
+      "epoch": 0.73449401523395,
+      "grad_norm": 2.3642656803131104,
+      "learning_rate": 0.00010800000000000001,
+      "loss": 0.1958,
+      "step": 540
+    },
+    {
+      "epoch": 0.7480957562568009,
+      "grad_norm": 2.765012502670288,
+      "learning_rate": 0.00011000000000000002,
+      "loss": 0.1857,
+      "step": 550
+    },
+    {
+      "epoch": 0.7616974972796517,
+      "grad_norm": 2.482921600341797,
+      "learning_rate": 0.00011200000000000001,
+      "loss": 0.186,
+      "step": 560
+    },
+    {
+      "epoch": 0.7752992383025027,
+      "grad_norm": 2.26837420463562,
+      "learning_rate": 0.00011399999999999999,
+      "loss": 0.1783,
+      "step": 570
+    },
+    {
+      "epoch": 0.7889009793253536,
+      "grad_norm": 2.1319069862365723,
+      "learning_rate": 0.000116,
+      "loss": 0.1731,
+      "step": 580
+    },
+    {
+      "epoch": 0.8025027203482046,
+      "grad_norm": 1.4668622016906738,
+      "learning_rate": 0.000118,
+      "loss": 0.1854,
+      "step": 590
+    },
+    {
+      "epoch": 0.8161044613710555,
+      "grad_norm": 2.2079055309295654,
+      "learning_rate": 0.00012,
+      "loss": 0.1806,
+      "step": 600
+    },
+    {
+      "epoch": 0.8297062023939065,
+      "grad_norm": 1.9991375207901,
+      "learning_rate": 0.000122,
+      "loss": 0.1682,
+      "step": 610
+    },
+    {
+      "epoch": 0.8433079434167573,
+      "grad_norm": 1.889404296875,
+      "learning_rate": 0.000124,
+      "loss": 0.162,
+      "step": 620
+    },
+    {
+      "epoch": 0.8569096844396082,
+      "grad_norm": 1.7021093368530273,
+      "learning_rate": 0.000126,
+      "loss": 0.1521,
+      "step": 630
+    },
+    {
+      "epoch": 0.8705114254624592,
+      "grad_norm": 2.153775215148926,
+      "learning_rate": 0.00012800000000000002,
+      "loss": 0.1656,
+      "step": 640
+    },
+    {
+      "epoch": 0.8841131664853101,
+      "grad_norm": 2.285336494445801,
+      "learning_rate": 0.00013000000000000002,
+      "loss": 0.1609,
+      "step": 650
+    },
+    {
+      "epoch": 0.8977149075081611,
+      "grad_norm": 1.8330312967300415,
+      "learning_rate": 0.000132,
+      "loss": 0.151,
+      "step": 660
+    },
+    {
+      "epoch": 0.911316648531012,
+      "grad_norm": 1.712404489517212,
+      "learning_rate": 0.000134,
+      "loss": 0.1577,
+      "step": 670
+    },
+    {
+      "epoch": 0.9249183895538629,
+      "grad_norm": 1.6210945844650269,
+      "learning_rate": 0.00013600000000000003,
+      "loss": 0.1641,
+      "step": 680
+    },
+    {
+      "epoch": 0.9385201305767138,
+      "grad_norm": 1.661000370979309,
+      "learning_rate": 0.000138,
+      "loss": 0.1543,
+      "step": 690
+    },
+    {
+      "epoch": 0.9521218715995647,
+      "grad_norm": 1.754515290260315,
+      "learning_rate": 0.00014,
+      "loss": 0.1621,
+      "step": 700
+    },
+    {
+      "epoch": 0.9657236126224157,
+      "grad_norm": 1.6415514945983887,
+      "learning_rate": 0.000142,
+      "loss": 0.144,
+      "step": 710
+    },
+    {
+      "epoch": 0.9793253536452666,
+      "grad_norm": 2.2498250007629395,
+      "learning_rate": 0.000144,
+      "loss": 0.1502,
+      "step": 720
+    },
+    {
+      "epoch": 0.9929270946681176,
+      "grad_norm": 1.6925517320632935,
+      "learning_rate": 0.000146,
+      "loss": 0.1546,
+      "step": 730
+    },
+    {
+      "epoch": 1.0054406964091405,
+      "grad_norm": 1.8807954788208008,
+      "learning_rate": 0.000148,
+      "loss": 0.1409,
+      "step": 740
+    },
+    {
+      "epoch": 1.0190424374319913,
+      "grad_norm": 1.611089825630188,
+      "learning_rate": 0.00015000000000000001,
+      "loss": 0.1578,
+      "step": 750
+    },
+    {
+      "epoch": 1.0326441784548421,
+      "grad_norm": 1.1505298614501953,
+      "learning_rate": 0.000152,
+      "loss": 0.1371,
+      "step": 760
+    },
+    {
+      "epoch": 1.0462459194776932,
+      "grad_norm": 1.5354456901550293,
+      "learning_rate": 0.000154,
+      "loss": 0.1407,
+      "step": 770
+    },
+    {
+      "epoch": 1.059847660500544,
+      "grad_norm": 1.178806185722351,
+      "learning_rate": 0.00015600000000000002,
+      "loss": 0.141,
+      "step": 780
+    },
+    {
+      "epoch": 1.073449401523395,
+      "grad_norm": 1.3755444288253784,
+      "learning_rate": 0.00015800000000000002,
+      "loss": 0.1287,
+      "step": 790
+    },
+    {
+      "epoch": 1.087051142546246,
+      "grad_norm": 1.5800870656967163,
+      "learning_rate": 0.00016,
+      "loss": 0.1382,
+      "step": 800
+    },
+    {
+      "epoch": 1.1006528835690967,
+      "grad_norm": 1.1685993671417236,
+      "learning_rate": 0.000162,
+      "loss": 0.1363,
+      "step": 810
+    },
+    {
+      "epoch": 1.1142546245919478,
+      "grad_norm": 1.1938755512237549,
+      "learning_rate": 0.000164,
+      "loss": 0.1316,
+      "step": 820
+    },
+    {
+      "epoch": 1.1278563656147986,
+      "grad_norm": 1.2022426128387451,
+      "learning_rate": 0.000166,
+      "loss": 0.1313,
+      "step": 830
+    },
+    {
+      "epoch": 1.1414581066376497,
+      "grad_norm": 1.1900382041931152,
+      "learning_rate": 0.000168,
+      "loss": 0.1256,
+      "step": 840
+    },
+    {
+      "epoch": 1.1550598476605005,
+      "grad_norm": 1.067172884941101,
+      "learning_rate": 0.00017,
+      "loss": 0.1385,
+      "step": 850
+    },
+    {
+      "epoch": 1.1686615886833516,
+      "grad_norm": 1.4434224367141724,
+      "learning_rate": 0.000172,
+      "loss": 0.1382,
+      "step": 860
+    },
+    {
+      "epoch": 1.1822633297062024,
+      "grad_norm": 1.0884168148040771,
+      "learning_rate": 0.000174,
+      "loss": 0.1266,
+      "step": 870
+    },
+    {
+      "epoch": 1.1958650707290532,
+      "grad_norm": 1.3909893035888672,
+      "learning_rate": 0.00017600000000000002,
+      "loss": 0.1215,
+      "step": 880
+    },
+    {
+      "epoch": 1.2094668117519043,
+      "grad_norm": 0.9344027042388916,
+      "learning_rate": 0.00017800000000000002,
+      "loss": 0.119,
+      "step": 890
+    },
+    {
+      "epoch": 1.2230685527747551,
+      "grad_norm": 1.1694083213806152,
+      "learning_rate": 0.00018,
+      "loss": 0.1192,
+      "step": 900
+    },
+    {
+      "epoch": 1.2366702937976062,
+      "grad_norm": 0.9874443411827087,
+      "learning_rate": 0.000182,
+      "loss": 0.1278,
+      "step": 910
+    },
+    {
+      "epoch": 1.250272034820457,
+      "grad_norm": 1.2893680334091187,
+      "learning_rate": 0.00018400000000000003,
+      "loss": 0.1318,
+      "step": 920
+    },
+    {
+      "epoch": 1.263873775843308,
+      "grad_norm": 1.3346811532974243,
+      "learning_rate": 0.00018600000000000002,
+      "loss": 0.128,
+      "step": 930
+    },
+    {
+      "epoch": 1.277475516866159,
+      "grad_norm": 0.9889335632324219,
+      "learning_rate": 0.000188,
+      "loss": 0.1372,
+      "step": 940
+    },
+    {
+      "epoch": 1.2910772578890097,
+      "grad_norm": 1.2218222618103027,
+      "learning_rate": 0.00019,
+      "loss": 0.1307,
+      "step": 950
+    },
+    {
+      "epoch": 1.3046789989118608,
+      "grad_norm": 0.897546112537384,
+      "learning_rate": 0.000192,
+      "loss": 0.1217,
+      "step": 960
+    },
+    {
+      "epoch": 1.3182807399347116,
+      "grad_norm": 0.9673519730567932,
+      "learning_rate": 0.000194,
+      "loss": 0.1094,
+      "step": 970
+    },
+    {
+      "epoch": 1.3318824809575625,
+      "grad_norm": 0.9263612031936646,
+      "learning_rate": 0.000196,
+      "loss": 0.1111,
+      "step": 980
+    },
+    {
+      "epoch": 1.3454842219804135,
+      "grad_norm": 1.0785088539123535,
+      "learning_rate": 0.00019800000000000002,
+      "loss": 0.12,
+      "step": 990
+    },
+    {
+      "epoch": 1.3590859630032643,
+      "grad_norm": 0.8844039440155029,
+      "learning_rate": 0.0002,
+      "loss": 0.1211,
+      "step": 1000
+    },
+    {
+      "epoch": 1.3726877040261154,
+      "grad_norm": 0.8651229739189148,
+      "learning_rate": 0.0001999999906373993,
+      "loss": 0.1291,
+      "step": 1010
+    },
+    {
+      "epoch": 1.3862894450489662,
+      "grad_norm": 1.0854979753494263,
+      "learning_rate": 0.000199999962549599,
+      "loss": 0.125,
+      "step": 1020
+    },
+    {
+      "epoch": 1.3998911860718173,
+      "grad_norm": 1.303252935409546,
+      "learning_rate": 0.00019999991573660427,
+      "loss": 0.1295,
+      "step": 1030
+    },
+    {
+      "epoch": 1.4134929270946681,
+      "grad_norm": 1.0601307153701782,
+      "learning_rate": 0.00019999985019842397,
+      "loss": 0.1363,
+      "step": 1040
+    },
+    {
+      "epoch": 1.427094668117519,
+      "grad_norm": 0.8127331733703613,
+      "learning_rate": 0.0001999997659350703,
+      "loss": 0.1124,
+      "step": 1050
+    },
+    {
+      "epoch": 1.44069640914037,
+      "grad_norm": 0.867559015750885,
+      "learning_rate": 0.0001999996629465591,
+      "loss": 0.1157,
+      "step": 1060
+    },
+    {
+      "epoch": 1.4542981501632208,
+      "grad_norm": 0.8963221907615662,
+      "learning_rate": 0.0001999995412329096,
+      "loss": 0.1197,
+      "step": 1070
+    },
+    {
+      "epoch": 1.467899891186072,
+      "grad_norm": 0.7211653590202332,
+      "learning_rate": 0.00019999940079414464,
+      "loss": 0.11,
+      "step": 1080
+    },
+    {
+      "epoch": 1.4815016322089227,
+      "grad_norm": 0.8746846914291382,
+      "learning_rate": 0.00019999924163029048,
+      "loss": 0.104,
+      "step": 1090
+    },
+    {
+      "epoch": 1.4951033732317738,
+      "grad_norm": 0.6718381643295288,
+      "learning_rate": 0.00019999906374137694,
+      "loss": 0.1041,
+      "step": 1100
+    },
+    {
+      "epoch": 1.5087051142546246,
+      "grad_norm": 0.7146100997924805,
+      "learning_rate": 0.00019999886712743732,
+      "loss": 0.0894,
+      "step": 1110
+    },
+    {
+      "epoch": 1.5223068552774754,
+      "grad_norm": 0.784561276435852,
+      "learning_rate": 0.00019999865178850845,
+      "loss": 0.0885,
+      "step": 1120
+    },
+    {
+      "epoch": 1.5359085963003265,
+      "grad_norm": 0.610625147819519,
+      "learning_rate": 0.00019999841772463066,
+      "loss": 0.0853,
+      "step": 1130
+    },
+    {
+      "epoch": 1.5495103373231773,
+      "grad_norm": 0.5619096159934998,
+      "learning_rate": 0.00019999816493584775,
+      "loss": 0.0788,
+      "step": 1140
+    },
+    {
+      "epoch": 1.5631120783460282,
+      "grad_norm": 0.5502200126647949,
+      "learning_rate": 0.00019999789342220708,
+      "loss": 0.0802,
+      "step": 1150
+    },
+    {
+      "epoch": 1.5767138193688792,
+      "grad_norm": 0.6030136942863464,
+      "learning_rate": 0.0001999976031837595,
+      "loss": 0.083,
+      "step": 1160
+    },
+    {
+      "epoch": 1.5903155603917303,
+      "grad_norm": 0.647160530090332,
+      "learning_rate": 0.00019999729422055928,
+      "loss": 0.0897,
+      "step": 1170
+    },
+    {
+      "epoch": 1.603917301414581,
+      "grad_norm": 0.7512218952178955,
+      "learning_rate": 0.00019999696653266437,
+      "loss": 0.08,
+      "step": 1180
+    },
+    {
+      "epoch": 1.617519042437432,
+      "grad_norm": 0.5823985934257507,
+      "learning_rate": 0.00019999662012013612,
+      "loss": 0.0772,
+      "step": 1190
+    },
+    {
+      "epoch": 1.631120783460283,
+      "grad_norm": 0.6494550108909607,
+      "learning_rate": 0.00019999625498303932,
+      "loss": 0.0772,
+      "step": 1200
+    },
+    {
+      "epoch": 1.6447225244831338,
+      "grad_norm": 0.5662053823471069,
+      "learning_rate": 0.00019999587112144244,
+      "loss": 0.0695,
+      "step": 1210
+    },
+    {
+      "epoch": 1.6583242655059847,
+      "grad_norm": 0.4981078803539276,
+      "learning_rate": 0.0001999954685354173,
+      "loss": 0.0742,
+      "step": 1220
+    },
+    {
+      "epoch": 1.6719260065288357,
+      "grad_norm": 0.5595643520355225,
+      "learning_rate": 0.00019999504722503927,
+      "loss": 0.0693,
+      "step": 1230
+    },
+    {
+      "epoch": 1.6855277475516868,
+      "grad_norm": 0.49906110763549805,
+      "learning_rate": 0.00019999460719038732,
+      "loss": 0.0692,
+      "step": 1240
+    },
+    {
+      "epoch": 1.6991294885745374,
+      "grad_norm": 0.5288122296333313,
+      "learning_rate": 0.00019999414843154375,
+      "loss": 0.0689,
+      "step": 1250
+    },
+    {
+      "epoch": 1.7127312295973884,
+      "grad_norm": 0.48555830121040344,
+      "learning_rate": 0.00019999367094859452,
+      "loss": 0.0655,
+      "step": 1260
+    },
+    {
+      "epoch": 1.7263329706202395,
+      "grad_norm": 0.5002060532569885,
+      "learning_rate": 0.00019999317474162905,
+      "loss": 0.0707,
+      "step": 1270
+    },
+    {
+      "epoch": 1.7399347116430903,
+      "grad_norm": 0.4510345458984375,
+      "learning_rate": 0.0001999926598107402,
+      "loss": 0.06,
+      "step": 1280
+    },
+    {
+      "epoch": 1.7535364526659412,
+      "grad_norm": 0.5075559020042419,
+      "learning_rate": 0.00019999212615602445,
+      "loss": 0.0675,
+      "step": 1290
+    },
+    {
+      "epoch": 1.7671381936887922,
+      "grad_norm": 0.5471305251121521,
+      "learning_rate": 0.0001999915737775817,
+      "loss": 0.0661,
+      "step": 1300
+    },
+    {
+      "epoch": 1.780739934711643,
+      "grad_norm": 0.6649473905563354,
+      "learning_rate": 0.00019999100267551538,
+      "loss": 0.0746,
+      "step": 1310
+    },
+    {
+      "epoch": 1.7943416757344939,
+      "grad_norm": 0.6705607175827026,
+      "learning_rate": 0.00019999041284993245,
+      "loss": 0.075,
+      "step": 1320
+    },
+    {
+      "epoch": 1.807943416757345,
+      "grad_norm": 0.645964503288269,
+      "learning_rate": 0.00019998980430094334,
+      "loss": 0.0825,
+      "step": 1330
+    },
+    {
+      "epoch": 1.821545157780196,
+      "grad_norm": 0.48304426670074463,
+      "learning_rate": 0.00019998917702866202,
+      "loss": 0.0726,
+      "step": 1340
+    },
+    {
+      "epoch": 1.8351468988030468,
+      "grad_norm": 0.5829260349273682,
+      "learning_rate": 0.00019998853103320592,
+      "loss": 0.0674,
+      "step": 1350
+    },
+    {
+      "epoch": 1.8487486398258977,
+      "grad_norm": 0.6563496589660645,
+      "learning_rate": 0.00019998786631469603,
+      "loss": 0.0666,
+      "step": 1360
+    },
+    {
+      "epoch": 1.8623503808487487,
+      "grad_norm": 0.5735076069831848,
+      "learning_rate": 0.00019998718287325676,
+      "loss": 0.0686,
+      "step": 1370
+    },
+    {
+      "epoch": 1.8759521218715995,
+      "grad_norm": 0.5888078212738037,
+      "learning_rate": 0.0001999864807090162,
+      "loss": 0.0728,
+      "step": 1380
+    },
+    {
+      "epoch": 1.8895538628944504,
+      "grad_norm": 0.5383855700492859,
+      "learning_rate": 0.00019998575982210572,
+      "loss": 0.0691,
+      "step": 1390
+    },
+    {
+      "epoch": 1.9031556039173014,
+      "grad_norm": 0.4812714159488678,
+      "learning_rate": 0.0001999850202126604,
+      "loss": 0.0591,
+      "step": 1400
+    },
+    {
+      "epoch": 1.9167573449401525,
+      "grad_norm": 0.5054184794425964,
+      "learning_rate": 0.00019998426188081865,
+      "loss": 0.0657,
+      "step": 1410
+    },
+    {
+      "epoch": 1.9303590859630033,
+      "grad_norm": 0.4280984103679657,
+      "learning_rate": 0.0001999834848267225,
+      "loss": 0.0619,
+      "step": 1420
+    },
+    {
+      "epoch": 1.9439608269858542,
+      "grad_norm": 0.7245299816131592,
+      "learning_rate": 0.0001999826890505175,
+      "loss": 0.064,
+      "step": 1430
+    },
+    {
+      "epoch": 1.9575625680087052,
+      "grad_norm": 0.5748353004455566,
+      "learning_rate": 0.0001999818745523526,
+      "loss": 0.0646,
+      "step": 1440
+    },
+    {
+      "epoch": 1.971164309031556,
+      "grad_norm": 0.5662197470664978,
+      "learning_rate": 0.00019998104133238034,
+      "loss": 0.0646,
+      "step": 1450
+    },
+    {
+      "epoch": 1.9847660500544069,
+      "grad_norm": 0.47684717178344727,
+      "learning_rate": 0.00019998018939075673,
+      "loss": 0.0626,
+      "step": 1460
+    },
+    {
+      "epoch": 1.998367791077258,
+      "grad_norm": 0.4988132417201996,
+      "learning_rate": 0.00019997931872764132,
+      "loss": 0.0596,
+      "step": 1470
+    },
+    {
+      "epoch": 2.010881392818281,
+      "grad_norm": 0.48226118087768555,
+      "learning_rate": 0.0001999784293431971,
+      "loss": 0.0579,
+      "step": 1480
+    },
+    {
+      "epoch": 2.0244831338411315,
+      "grad_norm": 0.44725948572158813,
+      "learning_rate": 0.0001999775212375907,
+      "loss": 0.0575,
+      "step": 1490
+    },
+    {
+      "epoch": 2.0380848748639826,
+      "grad_norm": 0.5599634051322937,
+      "learning_rate": 0.00019997659441099206,
+      "loss": 0.0594,
+      "step": 1500
+    },
+    {
+      "epoch": 2.0516866158868337,
+      "grad_norm": 0.47687241435050964,
+      "learning_rate": 0.00019997564886357476,
+      "loss": 0.0601,
+      "step": 1510
+    },
+    {
+      "epoch": 2.0652883569096843,
+      "grad_norm": 0.38273605704307556,
+      "learning_rate": 0.0001999746845955159,
+      "loss": 0.0551,
+      "step": 1520
+    },
+    {
+      "epoch": 2.0788900979325353,
+      "grad_norm": 0.5022798180580139,
+      "learning_rate": 0.00019997370160699602,
+      "loss": 0.0628,
+      "step": 1530
+    },
+    {
+      "epoch": 2.0924918389553864,
+      "grad_norm": 0.4850836992263794,
+      "learning_rate": 0.00019997269989819916,
+      "loss": 0.0562,
+      "step": 1540
+    },
+    {
+      "epoch": 2.1060935799782374,
+      "grad_norm": 0.45668402314186096,
+      "learning_rate": 0.0001999716794693129,
+      "loss": 0.0542,
+      "step": 1550
+    },
+    {
+      "epoch": 2.119695321001088,
+      "grad_norm": 0.4036906361579895,
+      "learning_rate": 0.00019997064032052837,
+      "loss": 0.0556,
+      "step": 1560
+    },
+    {
+      "epoch": 2.133297062023939,
+      "grad_norm": 0.46593400835990906,
+      "learning_rate": 0.00019996958245204009,
+      "loss": 0.0563,
+      "step": 1570
+    },
+    {
+      "epoch": 2.14689880304679,
+      "grad_norm": 0.4298928380012512,
+      "learning_rate": 0.00019996850586404615,
+      "loss": 0.0531,
+      "step": 1580
+    },
+    {
+      "epoch": 2.1605005440696408,
+      "grad_norm": 0.40037456154823303,
+      "learning_rate": 0.00019996741055674816,
+      "loss": 0.0519,
+      "step": 1590
+    },
+    {
+      "epoch": 2.174102285092492,
+      "grad_norm": 0.4456872344017029,
+      "learning_rate": 0.00019996629653035126,
+      "loss": 0.0477,
+      "step": 1600
+    },
+    {
+      "epoch": 2.187704026115343,
+      "grad_norm": 0.5066975355148315,
+      "learning_rate": 0.000199965163785064,
+      "loss": 0.049,
+      "step": 1610
+    },
+    {
+      "epoch": 2.2013057671381935,
+      "grad_norm": 0.39652055501937866,
+      "learning_rate": 0.0001999640123210985,
+      "loss": 0.0485,
+      "step": 1620
+    },
+    {
+      "epoch": 2.2149075081610445,
+      "grad_norm": 0.41883769631385803,
+      "learning_rate": 0.00019996284213867033,
+      "loss": 0.051,
+      "step": 1630
+    },
+    {
+      "epoch": 2.2285092491838956,
+      "grad_norm": 0.4462110996246338,
+      "learning_rate": 0.0001999616532379987,
+      "loss": 0.0515,
+      "step": 1640
+    },
+    {
+      "epoch": 2.2421109902067466,
+      "grad_norm": 0.5779175758361816,
+      "learning_rate": 0.00019996044561930622,
+      "loss": 0.0569,
+      "step": 1650
+    },
+    {
+      "epoch": 2.2557127312295973,
+      "grad_norm": 0.42402154207229614,
+      "learning_rate": 0.00019995921928281894,
+      "loss": 0.0519,
+      "step": 1660
+    },
+    {
+      "epoch": 2.2693144722524483,
+      "grad_norm": 0.642371654510498,
+      "learning_rate": 0.00019995797422876654,
+      "loss": 0.0612,
+      "step": 1670
+    },
+    {
+      "epoch": 2.2829162132752994,
+      "grad_norm": 0.44712305068969727,
+      "learning_rate": 0.0001999567104573822,
+      "loss": 0.0555,
+      "step": 1680
+    },
+    {
+      "epoch": 2.29651795429815,
+      "grad_norm": 0.5107985138893127,
+      "learning_rate": 0.0001999554279689025,
+      "loss": 0.0629,
+      "step": 1690
+    },
+    {
+      "epoch": 2.310119695321001,
+      "grad_norm": 0.4252520203590393,
+      "learning_rate": 0.00019995412676356762,
+      "loss": 0.0511,
+      "step": 1700
+    },
+    {
+      "epoch": 2.323721436343852,
+      "grad_norm": 0.413331001996994,
+      "learning_rate": 0.0001999528068416212,
+      "loss": 0.0504,
+      "step": 1710
+    },
+    {
+      "epoch": 2.337323177366703,
+      "grad_norm": 0.4216318726539612,
+      "learning_rate": 0.0001999514682033104,
+      "loss": 0.0539,
+      "step": 1720
+    },
+    {
+      "epoch": 2.3509249183895538,
+      "grad_norm": 0.5241215825080872,
+      "learning_rate": 0.0001999501108488859,
+      "loss": 0.0551,
+      "step": 1730
+    },
+    {
+      "epoch": 2.364526659412405,
+      "grad_norm": 0.48192480206489563,
+      "learning_rate": 0.00019994873477860185,
+      "loss": 0.0591,
+      "step": 1740
+    },
+    {
+      "epoch": 2.378128400435256,
+      "grad_norm": 0.5567557215690613,
+      "learning_rate": 0.00019994733999271596,
+      "loss": 0.0522,
+      "step": 1750
+    },
+    {
+      "epoch": 2.3917301414581065,
+      "grad_norm": 0.4027807414531708,
+      "learning_rate": 0.00019994592649148933,
+      "loss": 0.0524,
+      "step": 1760
+    },
+    {
+      "epoch": 2.4053318824809575,
+      "grad_norm": 0.40923604369163513,
+      "learning_rate": 0.0001999444942751867,
+      "loss": 0.0533,
+      "step": 1770
+    },
+    {
+      "epoch": 2.4189336235038086,
+      "grad_norm": 0.37654557824134827,
+      "learning_rate": 0.00019994304334407622,
+      "loss": 0.0482,
+      "step": 1780
+    },
+    {
+      "epoch": 2.432535364526659,
+      "grad_norm": 0.37279561161994934,
+      "learning_rate": 0.00019994157369842964,
+      "loss": 0.0448,
+      "step": 1790
+    },
+    {
+      "epoch": 2.4461371055495102,
+      "grad_norm": 0.34426528215408325,
+      "learning_rate": 0.0001999400853385221,
+      "loss": 0.0467,
+      "step": 1800
+    },
+    {
+      "epoch": 2.4597388465723613,
+      "grad_norm": 0.4583146870136261,
+      "learning_rate": 0.00019993857826463231,
+      "loss": 0.0501,
+      "step": 1810
+    },
+    {
+      "epoch": 2.4733405875952124,
+      "grad_norm": 0.4296802878379822,
+      "learning_rate": 0.00019993705247704245,
+      "loss": 0.0475,
+      "step": 1820
+    },
+    {
+      "epoch": 2.486942328618063,
+      "grad_norm": 0.3417333960533142,
+      "learning_rate": 0.00019993550797603828,
+      "loss": 0.0475,
+      "step": 1830
+    },
+    {
+      "epoch": 2.500544069640914,
+      "grad_norm": 0.3391024172306061,
+      "learning_rate": 0.000199933944761909,
+      "loss": 0.0428,
+      "step": 1840
+    },
+    {
+      "epoch": 2.514145810663765,
+      "grad_norm": 0.4173300862312317,
+      "learning_rate": 0.00019993236283494728,
+      "loss": 0.0487,
+      "step": 1850
+    },
+    {
+      "epoch": 2.527747551686616,
+      "grad_norm": 0.41140830516815186,
+      "learning_rate": 0.00019993076219544938,
+      "loss": 0.0499,
+      "step": 1860
+    },
+    {
+      "epoch": 2.5413492927094667,
+      "grad_norm": 0.4615647494792938,
+      "learning_rate": 0.00019992914284371497,
+      "loss": 0.0485,
+      "step": 1870
+    },
+    {
+      "epoch": 2.554951033732318,
+      "grad_norm": 0.4390008747577667,
+      "learning_rate": 0.00019992750478004738,
+      "loss": 0.0476,
+      "step": 1880
+    },
+    {
+      "epoch": 2.568552774755169,
+      "grad_norm": 0.36641502380371094,
+      "learning_rate": 0.00019992584800475322,
+      "loss": 0.0421,
+      "step": 1890
+    },
+    {
+      "epoch": 2.5821545157780195,
+      "grad_norm": 0.43274542689323425,
+      "learning_rate": 0.00019992417251814282,
+      "loss": 0.048,
+      "step": 1900
+    },
+    {
+      "epoch": 2.5957562568008705,
+      "grad_norm": 0.39262470602989197,
+      "learning_rate": 0.0001999224783205299,
+      "loss": 0.0482,
+      "step": 1910
+    },
+    {
+      "epoch": 2.6093579978237216,
+      "grad_norm": 0.3618634343147278,
+      "learning_rate": 0.0001999207654122316,
+      "loss": 0.0464,
+      "step": 1920
+    },
+    {
+      "epoch": 2.622959738846572,
+      "grad_norm": 0.3481171727180481,
+      "learning_rate": 0.0001999190337935688,
+      "loss": 0.0478,
+      "step": 1930
+    },
+    {
+      "epoch": 2.6365614798694232,
+      "grad_norm": 0.33801841735839844,
+      "learning_rate": 0.0001999172834648657,
+      "loss": 0.0444,
+      "step": 1940
+    },
+    {
+      "epoch": 2.6501632208922743,
+      "grad_norm": 0.337028443813324,
+      "learning_rate": 0.00019991551442645006,
+      "loss": 0.0403,
+      "step": 1950
+    },
+    {
+      "epoch": 2.663764961915125,
+      "grad_norm": 0.3380604684352875,
+      "learning_rate": 0.0001999137266786531,
+      "loss": 0.0483,
+      "step": 1960
+    },
+    {
+      "epoch": 2.677366702937976,
+      "grad_norm": 0.44155481457710266,
+      "learning_rate": 0.0001999119202218096,
+      "loss": 0.045,
+      "step": 1970
+    },
+    {
+      "epoch": 2.690968443960827,
+      "grad_norm": 0.43442797660827637,
+      "learning_rate": 0.00019991009505625784,
+      "loss": 0.0442,
+      "step": 1980
+    },
+    {
+      "epoch": 2.704570184983678,
+      "grad_norm": 0.4352855980396271,
+      "learning_rate": 0.00019990825118233957,
+      "loss": 0.0463,
+      "step": 1990
+    },
+    {
+      "epoch": 2.7181719260065287,
+      "grad_norm": 0.5073570609092712,
+      "learning_rate": 0.00019990638860040006,
+      "loss": 0.0492,
+      "step": 2000
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 73600,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 100,
+  "save_steps": 2000,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 5.54846097526789e+21,
+  "train_batch_size": 32,
+  "trial_name": null,
+  "trial_params": null
+}

checkpoint-2000/training_args.bin

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

checkpoint-2000/zero_to_fp32.py

@@ -0,0 +1,760 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+#   python zero_to_fp32.py . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device, weights_only=False)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+    total_files = len(files)
+    state_dicts = []
+    for f in tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        # an memory-efficient approach for sharding
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+        shard_state_dict = to_torch_tensor(shard_state_dict)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info("Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info("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)

dataset_stats.json

@@ -0,0 +1 @@
+{"observation.depth_images.camera_top": {"min": [[[0.0]], [[0.0]], [[0.0]]], "max": [[[1.0]], [[1.0]], [[1.0]]], "mean": [[[0.8737428225379095]], [[0.8737428225379095]], [[0.8737428225379095]]], "std": [[[0.01473990820408429]], [[0.01473990820408429]], [[0.01473990820408429]]], "count": [42299], "q01": [[[-3.921568627450984e-13]], [[-3.921568627450984e-13]], [[-3.921568627450984e-13]]], "q10": [[[-3.921568627450984e-13]], [[-3.921568627450984e-13]], [[-3.921568627450984e-13]]], "q50": [[[0.9998853739951559]], [[0.9998853739951559]], [[0.9998853739951559]]], "q90": [[[0.9999770747993448]], [[0.9999770747993448]], [[0.9999770747993448]]], "q99": [[[0.9999977074802886]], [[0.9999977074802886]], [[0.9999977074802886]]]}, "index": {"min": [0], "max": [117602], "mean": [58801.0], "std": [33949.061852526436], "count": [117603], "q01": [58616.72104526005], "q10": [58650.536767449725], "q50": [58800.778728192265], "q90": [58951.391855374524], "q99": [58985.269199886774]}, "frame_index": {"min": [0], "max": [894], "mean": [187.54551329472886], "std": [165.61093180395812], "count": [117603], "q01": [3.2665585547004277], "q10": [37.08228074445854], "q50": [187.30599470761794], "q90": [337.9373686692468], "q99": [371.8147131815649]}, "timestamp": {"min": [0.0], "max": [29.8], "mean": [6.251517109824292], "std": [5.520364393465271], "count": [117603], "q01": [0.10888528506161543], "q10": [1.2360760247376563], "q50": [6.244090840652991], "q90": [11.264578955719138], "q99": [12.393823772813889]}, "episode_index": {"min": [0], "max": [406], "mean": [225.93453398297657], "std": [130.58113875389165], "count": [117603], "q01": [225.93453398297657], "q10": [225.93453398297657], "q50": [225.93453398297657], "q90": [225.9345339829766], "q99": [225.9345339829766]}, "task_index": {"min": [0], "max": [0], "mean": [0.0], "std": [0.0], "count": [117603], "q01": [3.9999999999994196e-16], "q10": [3.999999999999417e-15], "q50": [1.9999999999997088e-14], "q90": [3.5999999999994735e-14], "q99": [3.959999999999425e-14]}, "observation.depth_images.camera_middle": {"min": [[[0.0]], [[0.0]], [[0.0]]], "max": [[[1.0]], [[1.0]], [[1.0]]], "mean": [[[0.7371832889075387]], [[0.7371832889075387]], [[0.7371832889075387]]], "std": [[[0.028255914071435326]], [[0.028255914071435326]], [[0.028255914071435326]]], "count": [42299], "q01": [[[-3.921568627450984e-13]], [[-3.921568627450984e-13]], [[-3.921568627450984e-13]]], "q10": [[[-3.921568627450984e-13]], [[-3.921568627450984e-13]], [[-3.921568627450984e-13]]], "q50": [[[0.9998641467443063]], [[0.9998641467443063]], [[0.9998641467443063]]], "q90": [[[0.9999728293491758]], [[0.9999728293491758]], [[0.9999728293491758]]], "q99": [[[0.9999972829352701]], [[0.9999972829352701]], [[0.9999972829352701]]]}, "state": {"min": [-0.3410964012145996, -0.017065448686480522, -0.8808728456497192, -0.3584683835506439, -0.47099998593330383, -0.25070762634277344, -0.3522544205188751, 0.20474791526794434, -0.951011598110199, -1.1395604610443115, 0.004790774080902338, -0.47099998593330383, -0.42500001192092896, -0.4710000455379486, 0.5233548879623413, -1.7453292608261108, 0.013509199023246765, 0.16424892842769623, 0.07441703230142593, 0.0, -1.7453292608261108, -1.7453292608261108, 0.0, 0.0, 0.0, 0.0], "max": [0.3343784511089325, 0.4551280438899994, 0.4898811876773834, 0.0, 0.17262005805969238, 0.4259999990463257, 0.4710000455379486, 1.42960786819458, 0.20000000298023224, 0.42028796672821045, 2.2654612064361572, 0.47099998593330383, 0.42500001192092896, 0.4710000455379486, 1.7453292608261108, -0.2723813056945801, 1.0848406553268433, 1.1712498664855957, 1.1572402715682983, 1.5509999990463257, 1.2000000476837158, -9.685753866506275e-06, 1.5509999990463257, 1.6920000314712524, 1.6920000314712524, 1.6920000314712524], "mean": [-0.07206582229242957, 0.17966891811792696, -0.1668464674341303, -0.0041360377980851635, -0.02855430856304381, 0.08563092750740967, 0.2893196088073914, 0.8363146289336536, -0.05387761385531189, -0.15336964771372463, 1.091904723826614, -0.28583224883559294, -0.2571242879282505, -0.04026229581134297, 1.631143633205173, -1.4191315095174262, 0.36479706485884494, 0.5091594637107267, 0.5016290432196715, 0.47609916631462623, -1.1053045955704097, -0.7251840643011697, 0.4714432698585686, 0.6842672242540222, 0.7672084591632831, 0.7534114637983107], "std": [0.13414144999809616, 0.06972650106920732, 0.25217907699852415, 0.02540099917055559, 0.11312150010916003, 0.10591266561362785, 0.19800768856144005, 0.1916404538214021, 0.1745209627611984, 0.2986915744066566, 0.37307377725837665, 0.27900556679351535, 0.14265063229848485, 0.2649863375472597, 0.203832594109652, 0.4301649488141138, 0.16008396160430588, 0.15833114838469356, 0.16450132461071795, 0.16326665338442858, 0.8143790653334856, 0.6069117444348734, 0.5677597965745057, 0.7476876651920975, 0.7628839306288268, 0.7625733012414243], "count": [117603], "q01": [-0.0900946626436895, 0.16442706509038613, -0.20515712685889975, -0.004887999590566002, -0.03502268398111734, 0.07959235241462335, 0.27133121740580957, 0.6528161505957908, -0.2659094747240229, -0.361794528686148, 0.7081018407532926, -0.46517954403538603, -0.35239264602504955, -0.2537239629698275, 1.6100977368982725, -1.4325886990180976, 0.3588537656561231, 0.504233242489377, 0.49273811896613856, 0.4674098045882135, -1.6605862154200803, -1.6429466853531651, 0.006851068033765835, 0.005850401283922236, 0.034195805934406874, 0.020448910315615902], "q10": [-0.08278976715559705, 0.16850630087211108, -0.19412164436363818, -0.004544541940541722, -0.03320866031030702, 0.0811968865892971, 0.2789572943566766, 0.6937583445861262, -0.20366920404476815, -0.32381025855426354, 0.7443205089939386, -0.4562365811872135, -0.3381223613521823, -0.21151295480617066, 1.620000935922196, -1.4291823533635943, 0.36067983515185703, 0.5057284518066023, 0.49521285631026923, 0.46997778254506795, -1.6132910523021013, -1.5679720148272271, 0.009773494327837565, 0.008090842227112085, 0.04456819834005742, 0.03460155597809609], "q50": [-0.0705348028684113, 0.18082898738071787, -0.1715710318250208, -0.004124543876975338, -0.028435645887978263, 0.08592637465300615, 0.29015337989340223, 0.841900340888932, -0.02172713823777844, -0.16216342465147707, 1.0210222394371753, -0.2633752477380261, -0.2654837467769859, -0.07935002363205927, 1.6318213337985956, -1.4217994092806205, 0.3648724970624905, 0.5095826853827928, 0.5021259825897231, 0.47654999878707843, -1.083572161473688, -0.5472054899472031, 0.2721429125038979, 0.6072634940305482, 0.6835550812469716, 0.6589050052086065], "q90": [-0.06286363941678236, 0.18891258937965152, -0.1354511074955689, -0.0037521858137056314, -0.024000949785724444, 0.08969127916363645, 0.2990914372977373, 0.9873103317771706, 0.03926613510762872, 0.015348704249211642, 1.5924075453087019, -0.11458629958531305, -0.1622620541588408, 0.15796104530149713, 1.6410317231895248, -1.4046785675889646, 0.3688311613039933, 0.5119426943225508, 0.5072208331283555, 0.48188938842079043, -0.62753262009664, -0.16851437724747884, 1.1652704487342436, 1.4333100584877039, 1.5411606659190882, 1.5361229393636513], "q99": [-0.059267575979736066, 0.19206527590878586, -0.12157279662742974, -0.0032695989562949138, -0.022354343712563778, 0.09092141466970259, 0.3036732241330353, 1.010810550033975, 0.05152326953739953, 0.11427531251100229, 1.7256451499216334, -0.07658159229284045, -0.12156362557106531, 0.23086158800094628, 1.6456560759243761, -1.400590835353288, 0.37035896563608134, 0.5138242272546271, 0.5097698584210922, 0.48475492430734096, -0.5283688574885755, -0.08377068065516621, 1.1850499603560902, 1.4442007498241958, 1.5479421826214554, 1.5470789057322587]}, "action": {"min": [-0.3389834761619568, -0.016907263547182083, -0.8794836401939392, -0.35402196645736694, -0.4709988832473755, -0.25057452917099, -0.3520881235599518, 0.2073279768228531, -0.9472358822822571, -1.1349681615829468, 0.009744583629071712, -0.4709988832473755, -0.42499902844429016, -1.4476113319396973, 0.5789587497711182, -1.7453292608261108, 0.013713635504245758, 0.16434887051582336, 0.07455907762050629, 4.245934346904196e-43, -1.7453292608261108, -1.7453292608261108, 0.0, 0.0, 0.0, 0.0], "max": [0.3341060280799866, 0.4546440839767456, 0.4896775186061859, 0.0, 0.17095127701759338, 0.4259990155696869, 0.9498989582061768, 1.4282805919647217, 0.19999954104423523, 0.43650567531585693, 2.253171682357788, 0.4709988832473755, 0.42499902844429016, 1.4465306997299194, 1.7453292608261108, -0.2728305459022522, 1.0846863985061646, 1.170958399772644, 1.156921148300171, 1.5509964227676392, 1.1999973058700562, -0.00046248978469520807, 1.5509964227676392, 1.6919960975646973, 1.6919960975646973, 1.6919960975646973], "mean": [-0.072062257609025, 0.1796734644739688, -0.16686331682866412, -0.0041351625959732225, -0.02852587046182658, 0.08565139129763345, 0.32956871001500315, 0.8355832817843305, -0.053134900457304526, -0.15265108764233445, 1.0894641430920489, -0.28540526575117153, -0.2575528473375333, -0.1262773200934414, 1.6311239782964502, -1.4192232057339058, 0.3648181842468053, 0.509163687405163, 0.5016597717877881, 0.4761237225457119, -1.1083735234476446, -0.7194708674206187, 0.46565009546621755, 0.6773420993678213, 0.7596605267104654, 0.7458658080296252], "std": [0.1341501065346926, 0.06972543554831165, 0.25211092742325175, 0.025382737904455964, 0.1130934105372598, 0.10589149914855811, 0.24532509622423124, 0.19137450795467298, 0.17346203701594373, 0.2989028806396463, 0.3708233988929887, 0.27894503456919906, 0.14230884887582138, 0.4576078292580101, 0.2038335208976552, 0.430100287209941, 0.1600713615414153, 0.15831408714073195, 0.16449934022152263, 0.16327920467584384, 0.8134898763692819, 0.603978072248423, 0.5649302168019992, 0.745754495145648, 0.7609507401745298, 0.7605825637964025], "count": [117603], "q01": [-0.08973729920989976, 0.1645976790662157, -0.20458951804463482, -0.004874202549583956, -0.03482574262124167, 0.07971509062882592, 0.30420597828507784, 0.6527587994501476, -0.2646251819438719, -0.36011265521491115, 0.7091029334884766, -0.4649887144841815, -0.35195990401961436, -0.4576455162499903, 1.6106382374401185, -1.432363790534601, 0.3590471331147991, 0.5042930752638022, 0.49288848604852414, 0.46757302854113886, -1.659711526428826, -1.6372496984979608, 0.006930907002576065, 0.005880558584038221, 0.034295387056921346, 0.02066212792632346], "q10": [-0.08273926163608117, 0.16854173190536725, -0.19408818756871984, -0.004508137714585739, -0.0331904104886033, 0.08121539567932311, 0.3136233512582266, 0.693144717812317, -0.20257035082552827, -0.32309771050310204, 0.7444061289768307, -0.4560105913966955, -0.3380824074876659, -0.40098712019743266, 1.6200639702629633, -1.4291839753538067, 0.3607006011976306, 0.5057366310468573, 0.49524575506096724, 0.46999732793000715, -1.6137498901483112, -1.5669113454421353, 0.009537449144312819, 0.0077853381815340674, 0.04414572443335411, 0.03418651956656013], "q50": [-0.07053305820219094, 0.18082641711411718, -0.17161606493473158, -0.004124987320009037, -0.028412243059807035, 0.08594706311879181, 0.3304782613207927, 0.8416115215533719, -0.021393517658311895, -0.16172603149867956, 1.0200302076258285, -0.26317536400011404, -0.26601014348674606, -0.17193387216764938, 1.6318250131570682, -1.4219394018659648, 0.36489079155891385, 0.5095974080373759, 0.502150824075827, 0.47656647103680383, -1.084467645695311, -0.541072469758272, 0.26532405899650024, 0.5980663829096573, 0.6736640718928992, 0.6484285009233538], "q90": [-0.0629178950179798, 0.18889062538451337, -0.1355250712340314, -0.003753144772577669, -0.023989247093333905, 0.08968235107935792, 0.3450012624506413, 0.9862027842628612, 0.039176925736125745, 0.016397138002638475, 1.5838448781169798, -0.11480706927083592, -0.16267037683146732, 0.16910018803031487, 1.6409919418960317, -1.4047737059652932, 0.3688356387792457, 0.5119281054285479, 0.5072521002643746, 0.4819028246975299, -0.6281812157310809, -0.16911760944302, 1.1644663662945496, 1.4327910248120364, 1.5408417009754056, 1.5354879258082261], "q99": [-0.05948614329976619, 0.1918574831135316, -0.12224194084860017, -0.0032829581361166196, -0.022470318704383582, 0.09085078092607207, 0.35246323822909376, 1.009487347622399, 0.05102437258844651, 0.11558852257291917, 1.722479834389032, -0.07742505615537215, -0.12283282492243937, 0.2537340214318874, 1.6454435218049546, -1.4009079608592592, 0.37029122424714694, 0.5136805481001522, 0.5096214037537491, 0.4846387095877261, -0.5373701207187795, -0.09244570347598105, 1.1841573043495919, 1.4436111705631345, 1.547584790705695, 1.5463872847427946]}, "cam_head_color": {"min": [[[0.0]], [[0.0]], [[0.0]]], "max": [[[1.0]], [[1.0]], [[1.0]]], "mean": [[[0.4158378660980179]], [[0.4476603398442878]], [[0.4604056766521079]]], "std": [[[0.04568349135173331]], [[0.0398879075799359]], [[0.03867090045986252]]], "count": [42299], "q01": [[[0.010803154657167389]], [[0.028855566163912962]], [[0.02618866991628496]]], "q10": [[[0.06614998816630643]], [[0.0965055780024583]], [[0.10173520661601497]]], "q50": [[[0.4025348264841603]], [[0.4495228986290825]], [[0.47094461193441506]]], "q90": [[[0.8042724440128624]], [[0.8027035566852121]], [[0.8116003231512168]]], "q99": [[[0.8926717697935003]], [[0.8771265633571808]], [[0.8777958268663902]]]}, "cam_hand_left_color": {"min": [[[0.0]], [[0.0]], [[0.0]]], "max": [[[1.0]], [[1.0]], [[1.0]]], "mean": [[[0.4628092887389989]], [[0.48527135800812327]], [[0.48852257640981633]]], "std": [[[0.03426552366211674]], [[0.031842301244512906]], [[0.03227401838744624]]], "count": [42299], "q01": [[[0.0016811978932110387]], [[0.03955406968803756]], [[0.03149705127898501]]], "q10": [[[0.06575519334489975]], [[0.11975457219415195]], [[0.11542931169145909]]], "q50": [[[0.44927931153963657]], [[0.46918153366830506]], [[0.48076844767956345]]], "q90": [[[0.8848905032993664]], [[0.8815046131329304]], [[0.8804975627721096]]], "q99": [[[0.9837387985467572]], [[0.9863311417944179]], [[0.9864727468590169]]]}}

go1_air_sft_libero.py

@@ -0,0 +1,107 @@
+import os
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from transformers import TrainingArguments
+
+from go1.configs.go1_base_cfg import BaseDatasetArguments, BaseModelArguments, BaseSpaceArguments
+from go1.tools.env_parse import get_bool_env
+
+# 获取运行名称和调试模式环境变量
+RUNNAME = os.environ.get("RUNNAME")
+DEBUG_MODE = get_bool_env("DEBUG_MODE")
+
+
+@dataclass
+class DatasetArguments(BaseDatasetArguments):
+    """数据集相关配置参数"""
+    # 数据集类型，默认为lerobot格式
+    dataset_type: Optional[str] = field(default="lerobot")
+    # 数据集根目录路径列表
+    data_root_dir: Optional[List[str]] = field(
+        default_factory=lambda: [
+            "/home/public/lerobot_datasets/w_bot_lerobot_3.0/Step1_pick_up_large_workpieces/"
+        ],
+    )
+    # 数据预处理变换操作列表
+    transforms: Optional[List[str]] = field(default_factory=lambda: [dict(type="Normalize")])
+
+
+@dataclass
+class GOModelArguments(BaseModelArguments):
+    """模型相关配置参数"""
+    # 预训练模型路径或名称
+    model_name_or_path: str = field(default="/home/hanqingqi/checkpoint/go-1-Air/")
+    # 是否冻结大语言模型参数
+    freeze_llm: bool = field(default=True)#field(default=False if not DEBUG_MODE else True)
+    # 是否冻结视觉主干网络参数
+    freeze_backbone: bool = field(default=True)# field(default=False if not DEBUG_MODE else True)
+    # 是否冻结MLP参数
+    freeze_mlp: bool = field(default=False) #field(default=False if not DEBUG_MODE else True)
+    # 动作序列块大小
+    action_chunk_size: int = field(default=60)
+    # 是否启用潜在空间规划
+    latent_planning: bool = field(default=False)
+
+
+@dataclass
+class GOTrainingArguments(TrainingArguments):
+    """训练相关配置参数"""
+    # 模型输出目录
+    output_dir: str = field(default=f"experiment/{RUNNAME}")
+    # 是否覆盖输出目录
+    overwrite_output_dir: bool = field(default=True)
+    # 数据加载器使用的进程数
+    dataloader_num_workers: int = field(default=32 if not DEBUG_MODE else 0)
+    # 是否使用bf16混合精度训练
+    bf16: bool = field(default=True)
+    # 训练轮数
+    num_train_epochs: float = field(default=100.0)
+    # 每个设备的训练批次大小
+    per_device_train_batch_size: int = field(default=32 if not DEBUG_MODE else 2)
+    # 梯度累积步数
+    gradient_accumulation_steps: int = field(default=5)
+    # 学习率
+    learning_rate: float = field(default=2e-4)
+    # 权重衰减系数
+    weight_decay: float = field(default=0.01)
+    # 学习率调度器类型
+    lr_scheduler_type: str = field(default="cosine")
+    # 学习率预热步数
+    warmup_steps: int = field(default=1000)
+    # 是否进行训练
+    do_train: bool = field(default=True)
+    # DeepSpeed配置文件路径
+    deepspeed: str = field(default="go1/zero_stage1_config.json")
+
+    # 模型保存策略
+    save_strategy: str = field(default="steps")
+    # 每多少步保存一次模型
+    save_steps: int = field(default=2000)
+    # 最多保存模型的数量
+    save_total_limit: int = field(default=100)
+    # 每多少步记录一次日志
+    logging_steps: int = field(default=10)
+    # 日志报告目标
+    report_to: str = field(default="tensorboard")
+
+
+@dataclass
+class SpaceArguments(BaseSpaceArguments):
+    """状态空间和动作空间配置参数"""
+    # 状态维度
+    state_dim: int = field(default=26)
+    # 动作维度
+    action_dim: int = field(default=26)
+    # 空间数据重映射配置
+    space_repack: dict = field(
+        default_factory=lambda: {
+            "state": "observation.state",    # 键为模型中的名称，值为数据在数据字典中的键名
+            "action": "action",
+            "cam_head_color": "observation.images.camera_top",
+            "cam_hand_left_color": "observation.images.camera_middle",
+            "final_prompt": "task",
+        }
+    )
+    # 控制频率
+    ctrl_freq: int = field(default=30)

runs/Dec06_16-18-34_user-SYS-821GE-TNHR/events.out.tfevents.1765009124.user-SYS-821GE-TNHR.2163381.0

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