Python--解决占用重复内存问题

报错信息：RuntimeError: Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: [{'base_model.encoder.embed_tokens.weight', 'base_model.shared.weight', 'base_model.decoder.embed_tokens.weight'}]. A potential way to correctly save your model is to use save_model. More information at https://huggingface.co/docs/safetensors/torch_shared_tensors

根据报错信息，问题出在使用 safetensors 库保存模型时，遇到了一些共享权重的张量，导致无法正常保存模型。具体来说，报错指出了多个共享权重：'base_model.encoder.embed_tokens.weight'、'base_model.shared.weight' 和 'base_model.decoder.embed_tokens.weight'。这类共享权重在保存时可能会占用重复的内存，从而引发报错。

可以通过以下几种方式来解决这个问题：

1. 使用 Trainer 提供的 save_model 方法

根据错误信息的提示，推荐使用 Trainer 类的 save_model 方法来保存模型，而不是直接使用 safetensors。Trainer 已经内置了一些机制来处理共享权重的情况，确保模型可以被正确保存。

在代码的保存模型部分，调用了 self.save_model(staging_output_dir)，这是正确的方式，但需要进一步改进：

# 改用 Trainer 类的 `save_model` 函数
trainer.save_model(output_dir)

确保 output_dir 路径是有效的，并且在调用 trainer.train() 时，检查是否有手动操作干扰了模型的保存机制。

2. 修改 safetensors 的保存策略

如果必须使用 safetensors，则需要确保保存时不保存共享的张量。可以使用 PyTorch 自带的保存功能，而不是 safetensors，因为 PyTorch 会自动处理共享权重。替换相关代码如下：

# 使用 PyTorch 原生保存方式，而不是 safetensors
torch.save(model.state_dict(), os.path.join(output_dir, "pytorch_model.bin"))

这将避免 safetensors 导致的共享内存问题。

3. 使用 Hugging Face 的标准保存方法

如果 safetensors 的问题依然存在，可以完全切换回 Hugging Face 官方推荐的标准保存方法。这通常不需要 safetensors，并且可以正确处理共享权重。尝试在 Trainer 中指定 save_total_limit 和 save_steps 参数，以确保定期保存模型。

training_args = TrainingArguments(
    output_dir=output_dir,
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=1,
    weight_decay=0.01,
    logging_dir=f'{output_dir}/logs',
    logging_steps=300,
    save_total_limit=2,  # 限制保存的总数量
    save_steps=500,      # 每500步保存一次模型
)

4. 禁用 safetensors

可以直接让 Trainer 在保存模型时使用传统的 .bin 格式，而不是 safetensors。在 Hugging Face transformers 中，默认情况下如果检测到 safetensors 可用，Trainer 会尝试使用它进行保存。可以通过禁用 safetensors 来规避这个问题。

解决方法： 在 TrainingArguments 中添加 save_safetensors=False 选项。

修改 training_args 部分的代码：

training_args = TrainingArguments(
    output_dir=output_dir,
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=1,
    weight_decay=0.01,
    logging_dir=f'{output_dir}/logs',
    logging_steps=300,
    save_total_limit=2,
    save_steps=500,
    save_safetensors=False  # 禁用 safetensors 保存
)

这样，Trainer 会使用标准的 .bin 文件来保存模型，避免 safetensors 引起的共享内存问题。

5. 手动保存模型

如果仍然希望使用 safetensors，可以手动保存模型，绕过 Trainer 的自动保存机制。

可以在训练结束后，调用 model.save_pretrained() 来手动保存模型。例如：

trainer.train()

# 手动保存模型
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

这种方法避免了 Trainer 自动保存时使用 safetensors 可能带来的问题。

6. 检查并移除共享权重

另一种方法是手动解除共享权重关系，但这可能会带来更多的问题，因为 BART 模型设计时共享权重是有其目的的。如果有特定的需求，想要确保不同模块的权重不共享，可以自行修改模型，但这涉及到深层的模型结构调整，不推荐作为优先选择。

结论

最简单直接的解决方案是通过 TrainingArguments 设置 save_safetensors=False 来禁用 safetensors，这样可以避免共享张量保存时引发的错误。如果希望手动控制模型保存，也可以调用 model.save_pretrained() 来保存模型。