Has anyone managed to run evaluate.py with the provided ARC-AGI-2 checkpoint?

[alexander-rakhlin]

ARC-AGI-2
I'm gettng size mismatch error during train_state initialization in evaluate.py:

train_state.model.load_state_dict(torch.load(eval_cfg.checkpoint, map_location="cuda"), assign=True)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[81], line 1
----> 1 train_state.model.load_state_dict(torch.load(eval_cfg.checkpoint, map_location="cuda"), assign=True)

File /opt/conda/lib/python3.11/site-packages/torch/nn/modules/module.py:2624, in Module.load_state_dict(self, state_dict, strict, assign)
   2616         error_msgs.insert(
   2617             0,
   2618             "Missing key(s) in state_dict: {}. ".format(
   2619                 ", ".join(f'"{k}"' for k in missing_keys)
   2620             ),
   2621         )
   2623 if len(error_msgs) > 0:
-> 2624     raise RuntimeError(
   2625         "Error(s) in loading state_dict for {}:\n\t{}".format(
   2626             self.__class__.__name__, "\n\t".join(error_msgs)
   2627         )
   2628     )
   2629 return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for OptimizedModule:
	size mismatch for _orig_mod.model.inner.puzzle_emb.weights: copying a param with shape torch.Size([1045829, 512]) from checkpoint, the shape in current model is torch.Size([1921251, 512]).

[alexander-rakhlin]

all_config.yaml is just as trained checkpoint

{
    "arch": {
        "H_cycles": 2,
        "H_layers": 4,
        "L_cycles": 2,
        "L_layers": 4,
        "expansion": 4,
        "halt_exploration_prob": 0.1,
        "halt_max_steps": 16,
        "hidden_size": 512,
        "loss": {
            "loss_type": "stablemax_cross_entropy",
            "name": "losses@ACTLossHead"
        },
        "name": "hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1",
        "num_heads": 8,
        "pos_encodings": "rope",
        "puzzle_emb_ndim": 512
    },
    "beta1": 0.9,
    "beta2": 0.95,
    "checkpoint_every_eval": true,
    "checkpoint_path": "checkpoints/Arc-2-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 bright-mustang",
    "data_path": "data/arc-2-aug-1000",
    "epochs": 100000,
    "eval_interval": 10000,
    "eval_save_outputs": [],
    "global_batch_size": 768,
    "lr": 0.0001,
    "lr_min_ratio": 1.0,
    "lr_warmup_steps": 2000,
    "project_name": "Arc-2-aug-1000 ACT-torch",
    "puzzle_emb_lr": 0.01,
    "puzzle_emb_weight_decay": 0.1,
    "run_name": "HierarchicalReasoningModel_ACTV1 bright-mustang",
    "seed": 0,
    "weight_decay": 0.1
}

[alexander-rakhlin]

While the developers are looking for an answer to why their checkpoint isn't loading, I'll answer for them. The problem is caused by the fact that the model architecture is parameterized by the dataset size. Not the dimensionality, but the size.

This error:

RuntimeError: Error(s) in loading state_dict for OptimizedModule:
	size mismatch for _orig_mod.model.inner.puzzle_emb.weights: copying a param with shape torch.Size([1045829, 512]) from checkpoint, the shape in current model is torch.Size([1921251, 512]).

is caused by the fact that _orig_mod.model.inner.puzzle_emb has a dimension of 1921251x512 but is initialized with a size of 1045829x512. 1921251 and 1045829 are the num_puzzle_identifiers that were involved in the training. For some reason, they are different in the developers' dataset and in our dataset. And since they parameterize the model (although they ideally shouldn't), the model breaks. To load the developers' checkpoint, you can replace "num_puzzle_identifiers": 1921251 with 1045829 in dataset.json.

[jin-qq]

True. In the file pretrain.py, change line 116 to num_puzzle_identifiers=1045829

[Bejro]

I'm also looking forward to having that issue resolved. It's possible to load the model but the distortion caused by the ids count shift makes the model accuracy drop to 0.