PyTorch Lightning中如何避免保存预训练子模块的检查点

在PyTorch Lightning项目中，当模型包含大型预训练子模块(如LLM、VAE等)时，默认的检查点机制会保存所有子模块的状态，这会导致存储空间和时间的浪费。本文将介绍如何优雅地解决这一问题。

问题背景

在构建复杂模型时，我们经常会使用预训练好的子模块作为模型的组成部分。这些子模块通常在整个训练过程中保持冻结状态，不会被更新。然而，PyTorch Lightning默认的检查点机制会保存所有子模块的参数，这带来了两个问题：

1. 存储空间浪费：大型预训练模型(如LLM)的参数可能占用数GB空间
2. 时间浪费：每次保存检查点时都需要序列化这些不变的数据

解决方案

PyTorch Lightning提供了两种主要方式来解决这个问题：

1. 自定义state_dict方法

通过重写模型的state_dict方法，可以精确控制哪些参数需要保存。例如：

def state_dict(self, *args, **kwargs):
    # 获取完整的状态字典
    state_dict = super().state_dict(*args, **kwargs)
    
    # 移除不需要保存的子模块
    for key in list(state_dict.keys()):
        if key.startswith("vae."):  # 假设vae是预训练的子模块
            del state_dict[key]
    
    return state_dict

这种方法提供了最大的灵活性，可以精确控制保存哪些参数。

2. 使用strict_loading=False选项

从PyTorch Lightning 2.2版本开始，可以设置self.strict_loading = False来允许加载部分检查点。这样即使检查点中不包含某些子模块的参数，模型也能正常加载。

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.strict_loading = False
        self.vae = AutoencoderKL.from_pretrained(...)  # 预训练的子模块

最佳实践

结合上述两种方法，可以构建既节省存储空间又便于部署的模型：

1. 对于完全冻结的预训练子模块，从state_dict中排除
2. 设置strict_loading=False以确保模型能加载部分检查点
3. 在文档中明确说明哪些子模块需要单独加载

注意事项

1. 确保排除的子模块确实不需要训练
2. 在部署时需要单独提供预训练子模块的权重
3. 测试模型加载逻辑，确保排除子模块后不影响功能

通过合理使用这些技术，可以显著减少模型检查点的大小，加快保存和加载速度，同时保持模型的完整功能。