Error by CUequivariance execution

[Seunghyo-Noh]

Hi, I want to use Cu-equvarience for accelerated training, but I am encountering the following error. How can I fix it?
I attached captued error image.

• The command I ran is sevenn "input_test.yaml -s --enable_cueq"
• I am using sevennet version 0.11.0 multi-fidelity
• I modified some presets provided by sevennet and attached the file.
  input_test.yaml.txt

Traceback (most recent call last):
File "/home/test/miniconda3/envs/sevenn/bin/sevenn", line 8, in
sys.exit(main())
File "/home/test/miniconda3/envs/sevenn/lib/python3.10/site-packages/sevenn/main/sevenn.py", line 190, in main
run(args)
File "/home/test/miniconda3/envs/sevenn/lib/python3.10/site-packages/sevenn/main/sevenn.py", line 129, in run
train_v2(global_config, working_dir)
File "/home/test/miniconda3/envs/sevenn/lib/python3.10/site-packages/sevenn/scripts/train.py", line 84, in train_v2
trainer.load_state_dicts(*state_dicts, strict=False)
File "/home/test/miniconda3/envs/sevenn/lib/python3.10/site-packages/sevenn/train/trainer.py", line 234, in load_state_dicts
self.model.load_state_dict(model_state_dict, strict=strict)
File "/home/test/miniconda3/envs/sevenn/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2215, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for AtomGraphSequential:
size mismatch for onehot_to_feature_x.linear.weight: copying a param with shape torch.Size([11392]) from checkpoint, the shape in current model is torch.Size([1, 11392]).
size mismatch for 0_self_connection_intro.linear.weight: copying a param with shape torch.Size([28672]) from checkpoint, the shape in current model is torch.Size([1, 28672]).
size mismatch for 0_self_interaction_1.linear.weight: copying a param with shape torch.Size([16640]) from checkpoint, the shape in current model is torch.Size([1, 16640]).
size mismatch for 0_self_interaction_2.linear.weight: copying a param with shape torch.Size([41408]) from checkpoint, the shape in current model is torch.Size([1, 41408]).
size mismatch for 1_self_connection_intro.linear.weight: copying a param with shape torch.Size([33792]) from checkpoint, the shape in current model is torch.Size([1, 33792]).
size mismatch for 1_self_interaction_1.linear.weight: copying a param with shape torch.Size([21760]) from checkpoint, the shape in current model is torch.Size([1, 21760]).
size mismatch for 1_self_interaction_2.linear.weight: copying a param with shape torch.Size([86464]) from checkpoint, the shape in current model is torch.Size([1, 86464]).
size mismatch for 2_self_connection_intro.linear.weight: copying a param with shape torch.Size([33792]) from checkpoint, the shape in current model is torch.Size([1, 33792]).
size mismatch for 2_self_interaction_1.linear.weight: copying a param with shape torch.Size([21760]) from checkpoint, the shape in current model is torch.Size([1, 21760]).
size mismatch for 2_self_interaction_2.linear.weight: copying a param with shape torch.Size([86464]) from checkpoint, the shape in current model is torch.Size([1, 86464]).
size mismatch for 3_self_connection_intro.linear.weight: copying a param with shape torch.Size([33792]) from checkpoint, the shape in current model is torch.Size([1, 33792]).
size mismatch for 3_self_interaction_1.linear.weight: copying a param with shape torch.Size([21760]) from checkpoint, the shape in current model is torch.Size([1, 21760]).
size mismatch for 3_self_interaction_2.linear.weight: copying a param with shape torch.Size([86464]) from checkpoint, the shape in current model is torch.Size([1, 86464]).
size mismatch for 4_self_connection_intro.linear.weight: copying a param with shape torch.Size([16384]) from checkpoint, the shape in current model is torch.Size([1, 16384]).
size mismatch for 4_self_interaction_1.linear.weight: copying a param with shape torch.Size([21760]) from checkpoint, the shape in current model is torch.Size([1, 21760]).
size mismatch for 4_self_interaction_2.linear.weight: copying a param with shape torch.Size([28928]) from checkpoint, the shape in current model is torch.Size([1, 28928]).
size mismatch for reduce_input_to_hidden.linear.weight: copying a param with shape torch.Size([8320]) from checkpoint, the shape in current model is torch.Size([1, 8320]).

[YutackPark]

Thanks for the bug report. This is apparently a bug and should be fixed.
Maybe one of these is the cause: multi-fidel + cuEquivariance or checkpoint_continue + cuEquivariance
Both path is not tested. I'll let you know after the fix.

[Seunghyo-Noh]

Hi @YutackPark

1. I have two more questions and I would like to inquire about the Cu_equivariant execution.
   In the case of SevenNet-0, I noticed that when I ran "sevenn --enable_cueq", the calculations were accelerated very quickly.
   However, when I ran cueq with the l3i5 setting, I found it difficult to experience an improvement in calculation speed, and it seemed to take longer based on my experience.
   I am curious if cueq is not applicable in the l3i5 setting, and if it will be supported in the future, when the update might take place.

2. The current calculations are based on float32 bits. Is there a way to convert them to float16 or mixed precision to save memory? Or I am curious if Automatic Mixed Precision (AMP) can be applied in SevenNet.

Best regards.

[YutackPark]

Hi, @Seunghyo-Noh , sorry for I have no update for the original issue. I'm currently a bit busy for various reasons..

1. cuEquivariance itself is still in alpha and performance improvement is not consistent. Actually, it is improved result you experienced from 7net-0, compared to cueq verision 0.1.0. When it first came out, I wasn't able to run 7net-0 (it stucked [Bug] torch.autograd.grad never ends, if ChannelWiseTensorProduct is used with specific configuration NVIDIA/cuEquivariance#33).

I have no idea which tensor product config is optimized or not, as it is closed source. At leaset, there is no filtering routine that preventing SevenNet to leverage cuEq by config.

2. Nope. It is not implemented yet (+ float64). But I suspect that AMP or float16 will not help as model is not bottlenecked by matrix multiplication.

We're also working on optimizing the model. I'll update when it became available.

[Seunghyo-Noh]

Hi, @Seunghyo-Noh , sorry for I have no update for the original issue. I'm currently a bit busy for various reasons..

1. cuEquivariance itself is still in alpha and performance improvement is not consistent. Actually, it is improved result you experienced from 7net-0, compared to cueq verision 0.1.0. When it first came out, I wasn't able to run 7net-0 (it stucked [Bug] torch.autograd.grad never ends, if ChannelWiseTensorProduct is used with specific configuration NVIDIA/cuEquivariance#33).

I have no idea which tensor product config is optimized or not, as it is closed source. At leaset, there is no filtering routine that preventing SevenNet to leverage cuEq by config.

2. Nope. It is not implemented yet (+ float64). But I suspect that AMP or float16 will not help as model is not bottlenecked by matrix multiplication.

We're also working on optimizing the model. I'll update when it became available.

@YutackPark
Thank you for your kind comments.
For original issue, I found the following phenomena.

1. training (w/o -cueq) -> continue (w/ -cueq) # not work
2. training (w/ -cueq) -> continue (w/ -cueq) # work
3. training (w/ -cueq) -> continue (w/o -cueq) # not work
4. training (w/o -cueq) -> continue (w/o -cueq) # work

Fortunately, if I turn on the option (-cueq) from scratch, it can be continued by keeping (-cueq) option.
Probably, it is not critical issue if I turn on --cueq from scratch training.

I hope this information help your bug fix.
Thanks again.