Pytorch 保留的数据比需要的多

Question

我正在尝试微调句子转换器。 问题是我遇到了 OOM 错误（我正在使用 google-cloud 来训练模型）。

我不断发现 pytorch 预留了 ~13GB（有 ~14GB）可用，因此没有空间容纳任何批次。

如果我尝试计算实际使用的内存，大约是 1.3GB

from sentence_transformers import models
model_name = "alexandrainst/scandi-nli-large"
word_embedding_model = models.Transformer(model_name, max_seq_length=512)
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
model.to("cuda")

param_size = 0
for param in model.parameters():
    param_size += param.nelement() * param.element_size()
buffer_size = 0
for buffer in model.buffers():
     buffer_size += buffer.nelement() * buffer.element_size()
size_all_mb = (param_size + buffer_size) / 1024 ** 2
print('model size: {:.3f}MB'.format(size_all_mb)) # ~1300
torch.cuda.memory_reserved()/(1024**2) # ~1300

我尝试调用

torch.cuda.empty_cache()

并设置

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

但出现同样的错误。

有没有办法不让 pytorch 保留内存（或至少减少它）而只使用所需的内存？

Answer 1

您正在计算模型权重所需的内存 - 这是训练所需总内存的一小部分。

训练模型时，您还可以为模型的激活、梯度和优化器状态分配内存。

Pytorch 不会“保留比所需更多的内存”——你只需要那么多内存来完成你想要做的事情。

要减少微调所需的内存，您可以查看以下内容：

• 使用混合精度训练
• 减少批量大小并使用梯度累积
• 梯度检查点
• 仅微调模型的最后一层
• 使用LORA等高效的微调方法

既然你正在使用Huggingface，这个将会很有用