MindSpore error: wq.weight in the argument 'net' should have the same shape as wq.weight in the argument 'parameter_dict'.

1 System Environment

Hardware Environment (Ascend/GPU/CPU): Ascend

MindSpore version: mindspore=2.3

Execution mode (PyNative/ Graph): Graph

Python version: Python=3.9

Operating system platform: Linux

2. Error Message

2.1 Problem Description

When executing a distributed training task, the following error occurs when loading the ckpt file:

For 'load_param_into_net', the wq.weight in the argument 'net' should have the same shape as wq.weight in the argument 'parameter_dict'. However, the shape of wq.weight in the argument 'net' is (512, 8192), and the shape of wq.weight in the argument 'parameter_dict' is (8192, 8192). Please check whether the loaded checkpoint is correct or whether the batch size and other parameters in 'net' and 'parameter_dict' are the same.

3 Root Cause Analysis

a. The error occurs during the loading of the checkpoint file.

b. The error message is: the shape of the wq.weight parameter in the network is (512, 8192), but the shape of the q.weight parameter in the checkpoint file is (8192, 8192), and the shapes are inconsistent, causing the loading to fail.

c. It was investigated whether the shape of the network or the shape of the checkpoint file did not meet the expected requirements. After investigation, it was found that the shape in the network is split and meets the expected requirements.

d. It was discovered that the parameters saved in the checkpoint have full shapes, not split ones.

e. Further investigation revealed that the save_checkpoint method, by default, merges all parameters and saves them as full parameters, which leads to inconsistency with the shape in the network.

4 Solution

When calling save_checkpoint, manually set the integrated_save parameter to false so that all nodes save the split parameters.

Reload the saved checkpoint file, and the problem is resolved.