multihead finetune issue

[GHLau]

when I perform multihead fine-tuning, an error occurs, but when I use naive fine-tuning, no error is reported. How to solve this problem? Thank you

(/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/e3nn/o3/_wigner.py:10: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for weights_only will be flipped to True. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via torch.serialization.add_safe_globals. We recommend you start setting weights_only=True for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
_Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(file), 'constants.pt'))
/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/cli/run_train.py:143: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for weights_only will be flipped to True. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via torch.serialization.add_safe_globals. We recommend you start setting weights_only=True for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
model_foundation = torch.load(
/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/calculators/mace.py:128: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for weights_only will be flipped to True. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via torch.serialization.add_safe_globals. We recommend you start setting weights_only=True for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
torch.load(f=model_path, map_location=device)
Traceback (most recent call last):
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/io/extxyz.py", line 419, in _read_xyz_frame
line = next(lines)
^^^^^^^^^^^
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/tools/multihead_tools.py", line 163, in assemble_mp_data
select_samples(dict_to_namespace(args_samples))
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/cli/fine_tuning_select.py", line 248, in select_samples
atoms_list_pt = ase.io.read(args.configs_pt, index=":")
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/io/formats.py", line 797, in read
return list(_iread(filename, index, format, io, parallel=parallel,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/parallel.py", line 302, in new_generator
for result in generator(*args, **kwargs):
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/io/formats.py", line 866, in _iread
for dct in io.read(fd, *args, **kwargs):
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/io/formats.py", line 624, in wrap_read_function
yield from read(filename, index, **kwargs)
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/io/extxyz.py", line 726, in read_xyz
yield _read_xyz_frame(fileobj, natoms, properties_parser, nvec)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/ase/io/extxyz.py", line 421, in _read_xyz_frame
raise XYZError('ase.io.extxyz: Frame has {} atoms, expected {}'
ase.io.extxyz.XYZError: ase.io.extxyz: Frame has 41 atoms, expected 64

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/bin/mace_run_train", line 8, in
sys.exit(main())
^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/cli/run_train.py", line 62, in main
run(args)
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/cli/run_train.py", line 261, in run
collections = assemble_mp_data(args, tag, head_configs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc2hdd/home/guanghualiu/anaconda3/envs/mace_multihead_env/lib/python3.12/site-packages/mace/tools/multihead_tools.py", line 183, in assemble_mp_data
raise RuntimeError(
RuntimeError: Model or descriptors download failed and no local model found
)

[ilyes319]

Hello,
Can you share your input script and pytorch environement please.

[GHLau]

mace 0.3.7 main branch

pytorch 2.4
（ conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia）
（conda install
e3nn
numpy
opt_einsum
ase
torch-ema
prettytable
matscipy
h5py
torchmetrics
python-hostlist
configargparse
GitPython
pyYAML
tqdm
matplotlib
pandas）

mace_run_train
--name="MACE"
--foundation_model="mace_agnesi_small.model"
--multiheads_finetuning=True
--train_file="train.xyz"
--valid_fraction=0.05
--test_file="test.xyz"
--energy_weight=1.0
--forces_weight=1.0
--E0s="foundation"
--lr=0.01
--scaling="rms_forces_scaling"
--batch_size=2
--max_num_epochs=6
--ema
--ema_decay=0.99
--amsgrad
--default_dtype="float32"
--device=cpu
--seed=3

[ilyes319]

Do you use a node that has internet access? If not you should run a test run on the login node to download everything and then execute on the GPU node. Also make sure your node has enough memory.

[GHLau]

yes! when i download everything, it works successfully.