使用 HuggingFace transformers 进行模型分片存储

发表于 阅读次数：

在HuggingFace transformers 的模型库中常常会见到这种pytorch_model.bin.index.json 参数名 与 模型bin文件的映射表，在如今的大模型中更为常见。本文主要介绍如何利用HuggingFace transformers进模型保存分片。

pytorch_model.bin.index.json 

1
2
3
4
5
6
7
8
9
10
11
{
  "metadata": {
    "total_size": 8396800
  },
  "weight_map": {
    "liner1.bias": "pytorch_model-00001-of-00002.bin",
    "liner1.weight": "pytorch_model-00001-of-00002.bin",
    "liner2.bias": "pytorch_model-00002-of-00002.bin",
    "liner2.weight": "pytorch_model-00002-of-00002.bin"
  }
}

权重结构：

1
2
3
4
5
saved_pt
├── config.json
├── pytorch_model-00001-of-00002.bin
├── pytorch_model-00002-of-00002.bin
└── pytorch_model.bin.index.json

PyTorch

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
# PyTorch 使用HuggingFace transformers保存模型方式

class ModelConfig(PretrainedConfig):
    model_type = "ptmodel"

    def __init__(self, size=1024, **kwargs):
        super().__init__(**kwargs)
        self.size = size


class Model(PreTrainedModel):
    def __init__(self, config: ModelConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.config = config
        self.liner1 = torch.nn.Linear(self.config.size, self.config.size)
        self.liner2 = torch.nn.Linear(self.config.size, self.config.size)

    def forward(self, input):
        out = self.liner1(input)
        out = self.liner2(out)
        return out
      
if __name__ == '__main__':
    saved_path = "saved_pt"
    os.makedirs(saved_path, exist_ok=True)

    config = ModelConfig(size=1024)

    model = Model(config=config)
    data = torch.ones((1024, 1024))
    out = model(data)
    print(f"model output:", out)
    model.save_pretrained(saved_path, max_shard_size="6MB")
    # 打印model参数
    print(model.state_dict())

    model1 = Model(config=config).from_pretrained(saved_path, from_pt=True, config=config)
    out1 = model1(data)
    print(f"model1 output:", out)
    # 打印model1 参数
    print(model1.state_dict())

权重结构：

1
2
3
4
5
6
7
8
9
'''
>>> tree saved_pt 
   
saved_pt
├── config.json
├── pytorch_model-00001-of-00002.bin
├── pytorch_model-00002-of-00002.bin
└── pytorch_model.bin.index.json
'''

TF2.x

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
# TF2.x 使用HuggingFace transformers保存模型方式

class ModelConfig(PretrainedConfig):
    model_type = "tfmodel"

    def __init__(self, size=1024, **kwargs):
        super().__init__(**kwargs)
        self.size = size


config = ModelConfig(size=1024)


@dataclass
class ModelOut(ModelOutput):
    logits: tf.Tensor = None


class Model(TFPreTrainedModel):

    @property
    def dummy_inputs(self):
        dummy = tf.constant(tf.ones((self.config.size, self.config.size), dtype=tf.float32))
        return dummy

    def __init__(self, config: ModelConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.config = config

        # TF2.x 一定要记住给name赋值，否则会根据整个类图进行增量自动命名
        self.dense1 = tf.keras.layers.Dense(self.config.size, name="dense1")
        self.dense2 = tf.keras.layers.Dense(self.config.size, name="dense2")

    def call(self, inputs, training=None, mask=None):
        out = self.dense1(inputs)
        out = self.dense2(out)
        return out

    @tf.function(
        input_signature=[
            (tf.TensorSpec(shape=(None, config.size), dtype=tf.float32, name="inputs"))
        ]
    )
    def serving(self, inputs):
        output = self.call(inputs)
        return self.serving_output(output)

    def serving_output(self, output):
        return ModelOut(logits=output)


if __name__ == '__main__':
    saved_path = "saved_tf"
    os.makedirs(saved_path, exist_ok=True)

    config = ModelConfig(size=1024)
    data = tf.ones((1024, 1024))
    model = Model(config=config)
    out = model(data)
    print(f"model output:", out)
    model.save_pretrained(saved_path, max_shard_size="6MB")
    # 打印model参数
    print(model.trainable_variables)

    model1 = Model(config=config).from_pretrained(saved_path, config=config)
    out1 = model1(data)
    print(f"model1 output:", out)
    # 打印model1参数
    print(model1.trainable_variables)

权重结构：

1
2
3
4
5
6
7
8
9
10
'''
>>> tree saved_tf

saved_tf
├── config.json
├── tf_model-00001-of-00002.h5
├── tf_model-00002-of-00002.h5
└── tf_model.h5.index.json

'''