Support storage of large parameters

[hbwx24]

PR types

Bug fixes

PR changes

APIs

Describe

paddle.save 和 paddle.static.save 存储模型，单个参数超过4 GiB时报错：

一、PR主要修改
1.将元素个数超过2^22的参数拆分，拆成元素个数不超过2^22的子向量。

2.在用pickle.dump保存的字典中会增加被拆分的子向量和拆分后的信息：

{'UnpackBigParamInfor@@':  {'_l.weight': {'OriginShape': (10, 67108864), 'slices': ['_l.weight@@.0', '_l.weight@@.1'...] }, 
 '_l.weight@@.0'：numpy data, 
 '_l.weight@@.1'：numpy data,...}

'UnpackBigParamInfor@@':拆分信息的名字；
'_l.weight'：被拆分的参数的名；
OriginShape：'_l.weight'的shape;
slices：被拆分后的的子向量名，例如：'_l.weight'被拆分成了'_l.weight@@.0'、'_l.weight@@.1'......
'_l.weight@@.0'：子向量的名字。

3.load的时候利用拆分后的信息还原被拆分的向量。

二、Example：
1.paddle.static.save/load:

import paddle
import paddle.static as static
import numpy as np

paddle.enable_static()
W = 2**15
x = static.data(name="x", shape=[None, W], dtype='float32')
y = static.nn.fc(x, W)
z = static.nn.fc(y, 10)

place = paddle.CPUPlace()
exe = static.Executor(place)
exe.run(static.default_startup_program())
prog = static.default_main_program()

inputs = np.random.randn(2, W).astype("float32")
result_z = exe.run(program=prog, feed={"x": inputs}, fetch_list=[z.name])

static.save(prog, "./big/temp")
static.load(prog, "./big/temp")
result_load = exe.run(program=prog, feed={"x": inputs}, fetch_list=[z.name])

print(np.sum(np.abs(result_z[0] - result_load[0])))

2.paddle.save/load:

import paddle
import numpy as np

W = 2**15
N = 2**1


class Layer(paddle.nn.Layer):
    def __init__(self):
        super(Layer, self).__init__()
        for i in range(N):
            setattr(self, "l_" + str(i), paddle.nn.Linear(W, W))

    def forward(self, x):
        for i in range(N):
            x = getattr(self, "l_" + str(i))(x)
        return x


layer = Layer()

save_dict = layer.state_dict()
print(len(save_dict))
path = "big/linear" + ".pdparams"
paddle.save(layer.state_dict(), path)
dict_load = paddle.load(path)

for key, value in save_dict.items():
    print("{} dev:{}".format(key,
                             np.sum(np.abs(dict_load[key] - value.numpy()))))

[paddle-bot-old]

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[paddle-bot-old]

✅ This PR's description meets the template requirements!
Please wait for other CI results.

[chenwhql]

LGTM

[chenwhql]

LGTM