search

OpenSenseNova / piccolo-embedding

code for piccolo embedding model from SenseTime
93 stars 4 forks source link

微调显存占用异常 #15

Open STHSF opened 2 weeks ago

STHSF commented 2 weeks ago

你好，我在使用piccolo-embedding代码分别加载stella-mrl-large-zh-v3.5-1792d和tao-8k两种模型分别在3090和A800上对比显存占用存在一些差异，请问这个怎么解释？

分别使用3090和A800加载stella-mrl-large-zh-v3.5-1792d（MAX_LENGTH=512）时显存占用基本符合预期， 3090单卡最大可以支撑112个batch， A800单卡最大可以支撑240个batch。但是加载tao-8k（MAX_LENGTH=8192）出现了不一样的表现， 3090单卡最大可以支撑12个batch，A800单卡最大可以支撑6个batch。

stella-mrl-large-zh-v3.5-1792d的模型大小为1248M和tao-8k的模型大小为1278M，两者的差异只有embedding层的30M差异。

训练过程中的显存占用细节如下：

image

image

hjq133 commented 2 weeks ago

我之前没有用过tao-8k这个模型，不知道这个模型会不会有坑。

另外： “3090单卡最大可以支撑12个batch，A800单卡最大可以支撑6个batch” 这个结论是你完全过完所有数据得到的吗，还是只跑了几个iteration？

因为sample batch的时候会有随机性，有时候一开始就会sample到特别长的batch导致out of memory。 手动构造几百条MAX_LENGTH > 8192的数据，再跑下实验试试呢

STHSF commented 2 weeks ago

数据量比较大没有完全跑完。但是也跑了四五百个iteration，并且尝试了很多次，感觉随机性的可能性比较小。 另外，我这边手动构造了几条MAX_LENGTH > 8192的数据进行了复现。以下是分别在3090和a800上的训练日志（附件）以及显存占用截图。训练日志中我打印了训练过程中不同类型数据的shape和模型的结构信息，没看出来啥问题，请大佬@hjq133 帮忙看看，感谢

[default0]:cls_contrast
[default0]:text_ids torch.Size([12, 8172])
[default0]:text_pos_ids torch.Size([12, 6])
[default0]:text_neg_ids torch.Size([120, 6])
[default1]:cosent
[default1]:text_ids torch.Size([12, 64])
[default1]:text_pair_ids torch.Size([12, 8192])
[default0]:STEmbedder(
[default0]:  (embedder): LastMeanEmbedder(
[default0]:    (encoder): BertModel(
[default0]:      (embeddings): BertEmbeddings(
[default0]:        (word_embeddings): Embedding(21128, 1024, padding_idx=0)
[default0]:        (position_embeddings): Embedding(8192, 1024)
[default0]:        (token_type_embeddings): Embedding(2, 1024)
[default0]:        (LayerNorm): LayerNorm((1024,), eps=1e-12, elementwise_affine=True)
[default0]:        (dropout): Dropout(p=0.1, inplace=False)
[default0]:      )
[default0]:      (encoder): BertEncoder(
[default0]:        (layer): ModuleList(
[default0]:          (0-23): 24 x BertLayer(
[default0]:            (attention): BertAttention(
[default0]:              (self): BertSdpaSelfAttention(
[default0]:                (query): Linear(in_features=1024, out_features=1024, bias=True)
[default0]:                (key): Linear(in_features=1024, out_features=1024, bias=True)
[default0]:                (value): Linear(in_features=1024, out_features=1024, bias=True)
[default0]:                (dropout): Dropout(p=0.1, inplace=False)
[default0]:              )
[default0]:              (output): BertSelfOutput(
[default0]:                (dense): Linear(in_features=1024, out_features=1024, bias=True)
[default0]:                (LayerNorm): LayerNorm((1024,), eps=1e-12, elementwise_affine=True)
[default0]:                (dropout): Dropout(p=0.1, inplace=False)
[default0]:              )
[default0]:            )
[default0]:            (intermediate): BertIntermediate(
[default0]:              (dense): Linear(in_features=1024, out_features=4096, bias=True)
[default0]:              (intermediate_act_fn): GELUActivation()
[default0]:            )
[default0]:            (output): BertOutput(
[default0]:              (dense): Linear(in_features=4096, out_features=1024, bias=True)
[default0]:              (LayerNorm): LayerNorm((1024,), eps=1e-12, elementwise_affine=True)
[default0]:              (dropout): Dropout(p=0.1, inplace=False)
[default0]:            )
[default0]:          )
[default0]:        )
[default0]:      )
[default0]:      (pooler): BertPooler(
[default0]:        (dense): Linear(in_features=1024, out_features=1024, bias=True)
[default0]:        (activation): Tanh()
[default0]:      )
[default0]:    )
[default0]:  )
[default0]:  (retri_contrst_loss): RetriContrastLoss(
[default0]:    (_cross_entropy_loss): CrossEntropyLoss()
[default0]:  )
[default0]:  (cosent_loss): CoSentLoss()
[default0]:  (cls_contrast_loss): ClsContrastLoss(
[default0]:    (_cross_entropy_loss): CrossEntropyLoss()
[default0]:  )
[default0]:  (scaling_layer): ScalingLayer(
[default0]:    (linear): Linear(in_features=1024, out_features=1792, bias=True)
[default0]:  )
[default0]:)

tao_8k_3090*2_test.log

image

tao_8k_a800*2_test.log image

hjq133 commented 2 weeks ago

emmm 这个确实有点难发现原因。会不会是neg数量不一致的原因呢，你把neg全部设置为0，保证数据只会来自于text和text_pos，显存会正常吗。

以及推理过程中，显存占用是正常的吗，是不是目前只会在训练侧出现问题

STHSF commented 2 weeks ago

我构造的数据格式跟data_example下的数据是完全一样的，只是把文本的长度复制扩充到8192，text_neg的list里的数据也只有一条，然后我的NEG_NUM参数取的也是1。

推理我还没有试过。不过我重新找了另外一个8k的模型（USER-bge-m3）来测试是否也有同样的问题，使用前面相同的人造数据进行测试发现bge-m3在3090跑完完整的流程最大只能支持bs=2，在A800上最大bs=10，这个跟原来stella-v3.5是一致的。

对比tao-8k和bge-m3的模型config.json文件，发现tao-8k多了一些pooler操作，但不知道是不是这个影响。

 "pooler_fc_size": 768,
 "pooler_num_attention_heads": 12,
 "pooler_num_fc_layers": 3,
 "pooler_size_per_head": 128,
 "pooler_type": "first_token_transform",

以下是bge-m3的模型结构：

STEmbedder(
  (embedder): LastMeanEmbedder(
    (encoder): XLMRobertaModel(
      (embeddings): XLMRobertaEmbeddings(
        (word_embeddings): Embedding(46166, 1024, padding_idx=1)
        (position_embeddings): Embedding(8194, 1024, padding_idx=1)
        (token_type_embeddings): Embedding(1, 1024)
        (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (dropout): Dropout(p=0.1, inplace=False)
      )
      (encoder): XLMRobertaEncoder(
        (layer): ModuleList(
          (0-23): 24 x XLMRobertaLayer(
            (attention): XLMRobertaAttention(
              (self): XLMRobertaSelfAttention(
                (query): Linear(in_features=1024, out_features=1024, bias=True)
                (key): Linear(in_features=1024, out_features=1024, bias=True)
                (value): Linear(in_features=1024, out_features=1024, bias=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
              (output): XLMRobertaSelfOutput(
                (dense): Linear(in_features=1024, out_features=1024, bias=True)
                (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
                (dropout): Dropout(p=0.1, inplace=False)
              )
            )
            (intermediate): XLMRobertaIntermediate(
              (dense): Linear(in_features=1024, out_features=4096, bias=True)
              (intermediate_act_fn): GELUActivation()
            )
            (output): XLMRobertaOutput(
              (dense): Linear(in_features=4096, out_features=1024, bias=True)
              (LayerNorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
              (dropout): Dropout(p=0.1, inplace=False)
            )
          )
        )
      )
      (pooler): XLMRobertaPooler(
        (dense): Linear(in_features=1024, out_features=1024, bias=True)
        (activation): Tanh()
      )
    )
  )
  (retri_contrst_loss): RetriContrastLoss(
    (_cross_entropy_loss): CrossEntropyLoss()
  )
  (cosent_loss): CoSentLoss()
  (cls_contrast_loss): ClsContrastLoss(
    (_cross_entropy_loss): CrossEntropyLoss()
  )
  (scaling_layer): ScalingLayer(
    (linear): Linear(in_features=1024, out_features=1792, bias=True)
  )
)
hjq133 commented 2 weeks ago

应该不是那个pooler的原因，那个应该只是写在config里，但实际embedding forward的时候应该不会过那些层的。我记得我之前的config里也有这个pooler，https://huggingface.co/sensenova/piccolo-large-zh-v2/blob/main/config.json

STHSF commented 1 week ago

我分别在3090和A800上部署了tao-8k的推理服务，显存占用没有什么特别的，很是奇怪。显存占用分别如下：

image

image