diff --git "a/step_75000/_METADATA" "b/step_75000/_METADATA"
deleted file mode 100644--- "a/step_75000/_METADATA"
+++ /dev/null
@@ -1 +0,0 @@
-{"tree_metadata": {"('train_rng',)": {"key_metadata": [{"key": "train_rng", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'step')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "step", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_0', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_1', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_10', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_11', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_2', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_3', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_4', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_5', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_6', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_7', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_8', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'Block_9', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'lm_head', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "lm_head", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'pos_emb', 'embedding')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "pos_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'params', 'tok_emb', 'embedding')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "params", "key_type": 2}, {"key": "tok_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '0')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "0", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('train_state', 'opt_state', '1', '0', 'count')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_0', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_1', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_10', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_11', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_2', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_3', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_4', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_5', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_6', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_7', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_8', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'Block_9', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'lm_head', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "lm_head", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'pos_emb', 'embedding')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "pos_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'mu', 'tok_emb', 'embedding')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "mu", "key_type": 2}, {"key": "tok_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_0', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_0", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_1', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_1", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_10', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_10", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_11', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_11", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_2', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_2", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_3', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_3", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_4', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_4", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_5', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_5", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_6', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_6", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_7', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_7", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_8', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_8", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'ExpertMLP_0', 'b1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'ExpertMLP_0', 'b2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "b2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'ExpertMLP_0', 'w1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w1", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'ExpertMLP_0', 'w2')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "ExpertMLP_0", "key_type": 2}, {"key": "w2", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'MultiHeadAttention_0', 'k_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "k_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'MultiHeadAttention_0', 'out_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'MultiHeadAttention_0', 'q_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "q_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'MultiHeadAttention_0', 'v_proj', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "MultiHeadAttention_0", "key_type": 2}, {"key": "v_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'Block_9', 'Router_0', 'gate', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "Block_9", "key_type": 2}, {"key": "Router_0", "key_type": 2}, {"key": "gate", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'lm_head', 'kernel')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "lm_head", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'pos_emb', 'embedding')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "pos_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '0', 'nu', 'tok_emb', 'embedding')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "0", "key_type": 1}, {"key": "nu", "key_type": 2}, {"key": "tok_emb", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}, "('train_state', 'opt_state', '1', '1')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "1", "key_type": 1}], "value_metadata": {"value_type": "None", "skip_deserialize": true}}, "('train_state', 'opt_state', '1', '2', 'count')": {"key_metadata": [{"key": "train_state", "key_type": 2}, {"key": "opt_state", "key_type": 2}, {"key": "1", "key_type": 1}, {"key": "2", "key_type": 1}, {"key": "count", "key_type": 2}], "value_metadata": {"value_type": "np.ndarray", "skip_deserialize": false}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}
\ No newline at end of file