Remove unused action expert modules

config.json

@@ -3,7 +3,6 @@
   "action_expert_config": {
     "attn_dropout": 0.0,
     "causal_attn": false,
-    "compile": "blocks",
     "context_layer_norm": true,
     "dropout": 0.0,
     "ffn_multiple_of": 256,
@@ -15,7 +14,6 @@
     "qk_norm": true,
     "qk_norm_eps": 1e-06,
     "rope": true,
-    "rope_on_cross_attention": true,
     "timestep_embed_dim": 256
   },
   "action_expert_depth_gate": false,

model-00004-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:df35c84b32b3460ad7bcca8ecbd3aefbe4d8caa281a030391d3586973965f340
-size 4998106920
+oid sha256:a81faa0f56099dd27590c1088e73b0a84e9fad71a322a90b89eb31dfd283d278
+size 4877619536

model-00005-of-00005.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f89916b6e59a3924f93e0de2fe7d6113eb843f82cf96fb464370864642470c46
-size 2334605176
+oid sha256:6b2eee6db4ad12f8b78fc3b0143aa4bd2510f477cdb2e736c355c41d26850afe
+size 2282630240

model.safetensors.index.json

@@ -1,14 +1,12 @@
 {
   "metadata": {
     "total_parameters": 5485309424,
-    "total_size": 21941237952
+    "total_size": 21768785088
   },
   "weight_map": {
     "lm_head.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.action_embed.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.action_embed.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.0.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.0.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.0.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.0.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.0.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -25,8 +23,6 @@
     "model.action_expert.blocks.0.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.0.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.0.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.1.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.1.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.1.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.1.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.1.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -43,8 +39,6 @@
     "model.action_expert.blocks.1.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.1.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.1.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.10.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.10.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.10.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.10.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.10.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -61,8 +55,6 @@
     "model.action_expert.blocks.10.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.10.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.10.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.11.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.11.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.11.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.11.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.11.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -79,8 +71,6 @@
     "model.action_expert.blocks.11.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.11.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.11.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.12.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.12.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.12.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.12.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.12.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -97,8 +87,6 @@
     "model.action_expert.blocks.12.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.12.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.12.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.13.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.13.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.13.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.13.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.13.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -115,8 +103,6 @@
     "model.action_expert.blocks.13.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.13.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.13.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.14.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.14.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.14.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.14.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.14.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -133,8 +119,6 @@
     "model.action_expert.blocks.14.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.14.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.14.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.15.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.15.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.15.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.15.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.15.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -151,8 +135,6 @@
     "model.action_expert.blocks.15.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.15.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.15.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.16.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.16.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.16.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.16.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.16.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -169,8 +151,6 @@
     "model.action_expert.blocks.16.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.16.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.16.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.17.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.17.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.17.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.17.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.17.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -187,8 +167,6 @@
     "model.action_expert.blocks.17.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.17.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.17.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.18.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.18.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.18.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.18.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.18.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -205,8 +183,6 @@
     "model.action_expert.blocks.18.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.18.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.18.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.19.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.19.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.19.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.19.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.19.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -223,8 +199,6 @@
     "model.action_expert.blocks.19.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.19.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.19.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.2.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.2.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.2.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.2.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.2.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -241,8 +215,6 @@
     "model.action_expert.blocks.2.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.2.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.2.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.20.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.20.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.20.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.20.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.20.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -259,8 +231,6 @@
     "model.action_expert.blocks.20.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.20.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.20.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.21.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.21.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.21.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.21.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.21.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -277,8 +247,6 @@
     "model.action_expert.blocks.21.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.21.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.21.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.22.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.22.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.22.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.22.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.22.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -295,8 +263,6 @@
     "model.action_expert.blocks.22.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.22.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.22.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.23.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.23.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.23.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.23.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.23.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -313,8 +279,6 @@
     "model.action_expert.blocks.23.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.23.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.23.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.24.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.24.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.24.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.24.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.24.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -331,8 +295,6 @@
     "model.action_expert.blocks.24.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.24.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.24.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.25.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.25.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.25.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.25.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.25.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -349,8 +311,6 @@
     "model.action_expert.blocks.25.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.25.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.25.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.26.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.26.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.26.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.26.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.26.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -367,8 +327,6 @@
     "model.action_expert.blocks.26.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.26.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.26.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.27.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.27.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.27.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.27.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.27.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -385,8 +343,6 @@
     "model.action_expert.blocks.27.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.27.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.27.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.28.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.28.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.28.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.28.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.28.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -403,8 +359,6 @@
     "model.action_expert.blocks.28.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.28.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.28.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.29.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.29.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.29.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.29.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.29.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -421,8 +375,6 @@
     "model.action_expert.blocks.29.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.29.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.29.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.3.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.3.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.3.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.3.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.3.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -439,8 +391,6 @@
     "model.action_expert.blocks.3.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.3.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.3.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.30.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.30.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.30.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.30.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.30.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -457,8 +407,6 @@
     "model.action_expert.blocks.30.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.30.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.30.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.31.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.31.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.31.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.31.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.31.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -475,8 +423,6 @@
     "model.action_expert.blocks.31.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.31.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.31.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.32.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.32.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.32.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.32.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.32.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -493,8 +439,6 @@
     "model.action_expert.blocks.32.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.32.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.32.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.33.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.33.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.33.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.33.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.33.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -511,8 +455,6 @@
     "model.action_expert.blocks.33.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.33.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.33.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.34.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.34.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.34.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.34.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.34.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -529,8 +471,6 @@
     "model.action_expert.blocks.34.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.34.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.34.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.35.cross_attn.kv_proj.bias": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.35.cross_attn.kv_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.35.cross_attn.out_proj.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.35.cross_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.35.cross_attn.q_proj.bias": "model-00005-of-00005.safetensors",
@@ -547,8 +487,6 @@
     "model.action_expert.blocks.35.self_attn.out_proj.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.35.self_attn.qkv.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.blocks.35.self_attn.qkv.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.blocks.4.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.4.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.4.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.4.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.4.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -565,8 +503,6 @@
     "model.action_expert.blocks.4.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.4.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.4.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.5.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.5.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.5.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.5.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.5.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -583,8 +519,6 @@
     "model.action_expert.blocks.5.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.5.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.5.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.6.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.6.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.6.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.6.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.6.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -601,8 +535,6 @@
     "model.action_expert.blocks.6.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.6.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.6.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.7.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.7.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.7.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.7.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.7.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -619,8 +551,6 @@
     "model.action_expert.blocks.7.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.7.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.7.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.8.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.8.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.8.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.8.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.8.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -637,8 +567,6 @@
     "model.action_expert.blocks.8.self_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.8.self_attn.qkv.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.8.self_attn.qkv.weight": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.9.cross_attn.kv_proj.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.blocks.9.cross_attn.kv_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.9.cross_attn.out_proj.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.9.cross_attn.out_proj.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.blocks.9.cross_attn.q_proj.bias": "model-00004-of-00005.safetensors",
@@ -661,8 +589,6 @@
     "model.action_expert.final_layer.linear.weight": "model-00005-of-00005.safetensors",
     "model.action_expert.final_layer.modulation.linear.bias": "model-00005-of-00005.safetensors",
     "model.action_expert.final_layer.modulation.linear.weight": "model-00005-of-00005.safetensors",
-    "model.action_expert.state_encoder.bias": "model-00004-of-00005.safetensors",
-    "model.action_expert.state_encoder.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.time_embed.1.bias": "model-00004-of-00005.safetensors",
     "model.action_expert.time_embed.1.weight": "model-00004-of-00005.safetensors",
     "model.action_expert.time_embed.3.bias": "model-00004-of-00005.safetensors",

modeling_molmoact2.py

@@ -315,17 +315,9 @@ class ActionExpertCrossAttention(nn.Module):
             ActionExpertRMSNorm(self.head_dim, eps=qk_norm_eps) if qk_norm else None
         )
         self.q_proj = nn.Linear(hidden_size, hidden_size)
-        self.kv_proj = nn.Linear(hidden_size, hidden_size * 2)
         self.out_proj = nn.Linear(hidden_size, hidden_size)
         self.out_drop = nn.Dropout(proj_dropout)
 
-    def _apply_qk_norm(
-        self, q: torch.Tensor, k: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        if self.q_norm is None or self.k_norm is None:
-            return q, k
-        return self.q_norm(q), self.k_norm(k)
-
     def _as_heads(self, x: torch.Tensor) -> torch.Tensor:
         if x.dim() == 4:
             if x.shape[2] == self.num_heads:
@@ -361,38 +353,18 @@ class ActionExpertCrossAttention(nn.Module):
         self,
         x: torch.Tensor,
         *,
-        kv: Optional[torch.Tensor] = None,
-        kv_k: Optional[torch.Tensor] = None,
-        kv_v: Optional[torch.Tensor] = None,
+        kv_k: torch.Tensor,
+        kv_v: torch.Tensor,
         attn_mask: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        if (kv_k is None) != (kv_v is None):
-            raise ValueError("kv_k and kv_v must both be provided or both be None.")
-        if kv is not None and kv_k is not None:
-            raise ValueError("Provide either kv or kv_k/kv_v, not both.")
         bsz, tgt_len, _ = x.shape
         q = self.q_proj(x).view(bsz, tgt_len, self.num_heads, self.head_dim)
-        if kv_k is not None and kv_v is not None:
-            k = self._as_heads(kv_k)
-            v = self._as_heads(kv_v)
-            k_pre_normed = True
-        else:
-            if kv is None:
-                raise ValueError("cross-attention requires kv or kv_k/kv_v.")
-            src_len = kv.shape[1]
-            kv_proj = self.kv_proj(kv).view(
-                bsz, src_len, 2, self.num_heads, self.head_dim
-            )
-            k = kv_proj[:, :, 0]
-            v = kv_proj[:, :, 1]
-            k_pre_normed = False
+        k = self._as_heads(kv_k)
+        v = self._as_heads(kv_v)
         q = q.transpose(1, 2)
         k = k.transpose(1, 2)
-        if k_pre_normed:
-            if self.q_norm is not None:
-                q = self.q_norm(q)
-        else:
-            q, k = self._apply_qk_norm(q, k)
+        if self.q_norm is not None:
+            q = self.q_norm(q)
         q = q.transpose(1, 2)
         k = k.transpose(1, 2)
         out = self._attention(q, k, v, attn_mask=attn_mask)
@@ -592,12 +564,6 @@ class ActionExpert(nn.Module):
         self.action_embed = nn.Linear(
             config.max_action_dim, config.hidden_size, device=device
         )
-        self.state_encoder = nn.Linear(
-            config.hidden_size, config.hidden_size, device=device
-        )
-        self.state_norm = ActionExpertRMSNorm(
-            config.hidden_size, eps=1e-6, device=device
-        )
         self.context_k_proj = nn.Linear(
             self.llm_kv_dim, config.hidden_size, bias=False, device=device
         )
@@ -629,10 +595,6 @@ class ActionExpert(nn.Module):
                 for _ in range(config.num_layers)
             ]
         )
-        for block in self.blocks:
-            block.cross_attn.kv_proj.weight.requires_grad = False
-            if block.cross_attn.kv_proj.bias is not None:
-                block.cross_attn.kv_proj.bias.requires_grad = False
         self.final_layer = ActionExpertFinalLayer(
             config.hidden_size, config.max_action_dim
         )
@@ -643,8 +605,6 @@ class ActionExpert(nn.Module):
             if isinstance(module, nn.Linear):
                 _init_linear(module)
         _init_linear(self.action_embed)
-        _init_linear(self.state_encoder)
-        self.state_norm.reset_parameters()
         _init_linear(self.context_k_proj)
         _init_linear(self.context_v_proj)
         if isinstance(self.context_norm, ActionExpertRMSNorm):
@@ -654,7 +614,6 @@ class ActionExpert(nn.Module):
             _init_linear(block.self_attn.qkv)
             _init_linear(block.self_attn.out_proj, scale=residual_scale)
             _init_linear(block.cross_attn.q_proj)
-            _init_linear(block.cross_attn.kv_proj)
             _init_linear(block.cross_attn.out_proj, scale=residual_scale)
             _init_linear(block.mlp.up_proj)
             _init_linear(block.mlp.gate_proj)
@@ -680,19 +639,6 @@ class ActionExpert(nn.Module):
             x.shape[0], x.shape[1], self.config.num_heads, self.action_head_dim
         )
 
-    def _encode_states(self, states: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
-        if states is None:
-            return None
-        if states.dim() == 2:
-            states = states.unsqueeze(1)
-        if states.shape[-1] != self.hidden_size:
-            feat_dim = states.shape[-1]
-            if feat_dim < self.hidden_size:
-                states = F.pad(states, (0, self.hidden_size - feat_dim))
-            else:
-                states = states[..., : self.hidden_size]
-        return self.state_norm(self.state_encoder(states))
-
     def _time_conditioning(self, timesteps: torch.Tensor) -> torch.Tensor:
         conditioning = self.time_embed[0](timesteps)
         first_linear = self.time_embed[1]
@@ -709,7 +655,6 @@ class ActionExpert(nn.Module):
     def _prepare_kv_context(
         self,
         encoder_kv_states: Sequence[Tuple[torch.Tensor, torch.Tensor]],
-        encoded_states: Optional[torch.Tensor],
     ) -> Sequence[Tuple[torch.Tensor, torch.Tensor]]:
         if len(encoder_kv_states) != len(self.blocks):
             raise ValueError(
@@ -717,17 +662,9 @@ class ActionExpert(nn.Module):
                 f"got {len(encoder_kv_states)}."
             )
         kv_contexts = []
-        state_heads = (
-            self._reshape_hidden_to_heads(encoded_states)
-            if encoded_states is not None
-            else None
-        )
         for block, (k_in, v_in) in zip(self.blocks, encoder_kv_states):
             k_ctx = self._project_kv_tensor(k_in, self.context_k_proj)
             v_ctx = self._project_kv_tensor(v_in, self.context_v_proj)
-            if state_heads is not None:
-                k_ctx = torch.cat([k_ctx, state_heads], dim=1)
-                v_ctx = torch.cat([v_ctx, state_heads], dim=1)
             k_norm = block.cross_attn.k_norm
             if k_norm is not None:
                 k_ctx = k_norm(k_ctx.transpose(1, 2)).transpose(1, 2)
@@ -737,24 +674,12 @@ class ActionExpert(nn.Module):
     @staticmethod
     def _build_cross_attention_mask(
         encoder_attention_mask: Optional[torch.Tensor],
-        encoded_states: Optional[torch.Tensor],
         batch_size: int,
         dtype: torch.dtype,
     ) -> Optional[torch.Tensor]:
-        state_seq_len = 0 if encoded_states is None else encoded_states.shape[1]
         if encoder_attention_mask is None:
             return None
         mask = encoder_attention_mask[:, None, None, :].to(dtype=dtype)
-        if state_seq_len > 0:
-            ones = torch.ones(
-                batch_size,
-                1,
-                1,
-                state_seq_len,
-                device=mask.device,
-                dtype=mask.dtype,
-            )
-            mask = torch.cat([mask, ones], dim=-1)
         return (1.0 - mask) * torch.finfo(dtype).min
 
     def _build_self_attention_mask(
@@ -792,7 +717,11 @@ class ActionExpert(nn.Module):
         device: torch.device,
         dtype: torch.dtype,
     ) -> ActionExpertContext:
-        encoded_states = self._encode_states(state_embeddings)
+        if state_embeddings is not None:
+            raise ValueError(
+                "MolmoAct2 HF action expert supports only discrete state tokens. "
+                "Continuous state embeddings are not supported."
+            )
         valid_action = None
         if action_attention_mask is not None:
             valid_action = action_attention_mask.to(
@@ -805,10 +734,9 @@ class ActionExpert(nn.Module):
                 device=device,
                 dtype=dtype,
             )
-        kv_contexts = self._prepare_kv_context(encoder_kv_states, encoded_states)
+        kv_contexts = self._prepare_kv_context(encoder_kv_states)
         cross_mask = self._build_cross_attention_mask(
             encoder_attention_mask,
-            encoded_states,
             batch_size,
             dtype,
         )