使用Opacus PrivacyEngine与PyTorch Lightning时遇RecursionError求助

问题现象 #

运行基于Opacus PrivacyEngine封装基础推荐模型的DPModel时，触发递归错误：

RecursionError: maximum recursion depth exceeded while calling a Python object

代码结构 #

基础模型（BaseModel.py） #

class BaseModel(pl.LightningModule):
    # 带用户/物品嵌入的基础推荐系统
    def forward(self, user_ids, item_ids):
        user_embed = self.user_embedding(user_ids)
        item_embed = self.item_embedding(item_ids)
        x = torch.cat([user_embed, item_embed], dim=1)
        return self.fc(x).squeeze()

DP模型（DPModel.py） #

class DPModel(BaseModel):
    def on_train_start(self):
        if self.enable_dp:
            self.dp_model, dp_optimizer, _ = self.privacy_engine.make_private(
                module=self,
                optimizer=optimizer,
                data_loader=train_loader,
                noise_multiplier=self.hparams.noise_multiplier,
                max_grad_norm=self.hparams.max_grad_norm
            )
    
    def forward(self, *args, **kwargs):
        if hasattr(self, 'dp_model'):
            return self.dp_model(*args, **kwargs)
        return super().forward(*args, **kwargs)

数据模块 #

def train_dataloader(self):
    return DPDataLoader(  # 使用Opacus的DPDataLoader
        self.train_dataset,
        sample_rate=self.batch_size/len(self.train_dataset),
        num_workers=self.num_workers
    )

已尝试方案 #

• 在__init__中初始化self.dp_model = None
• 调用make_private前保存原始模型引用
• 更换Opacus（1.5.3）和PyTorch（2.6.0）的不同版本

问题提问 #

如何正确集成Opacus与PyTorch Lightning以避免递归问题？具体来说：

1. 正确使用make_private进行模型封装的方式是什么？
2. 是否需要修改Lightning访问模型参数的方式？

解决方案 #

递归错误根源 #

当前实现触发无限递归的核心原因：DPModel继承自BaseModel，在on_train_start中用make_private封装当前实例self得到dp_model，随后重写的forward又调用self.dp_model(*args, **kwargs)。但dp_model本质是封装后的DPModel实例，它的forward最终会调用原模型的forward，从而形成循环调用链，触发RecursionError。

1. 正确使用make_private的方式 #

不需要额外保存dp_model，直接让make_private修改当前LightningModule实例的行为即可，具体调整如下：

class DPModel(BaseModel):
    def __init__(self, enable_dp, hparams, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.enable_dp = enable_dp
        self.hparams = hparams
        self.privacy_engine = PrivacyEngine()  # 提前初始化PrivacyEngine

    def on_train_start(self):
        if self.enable_dp:
            # 直接修改当前模型实例，无需保存dp_model
            self.privacy_engine.make_private(
                module=self,
                optimizer=self.trainer.optimizers[0],  # 从Lightning Trainer获取已初始化的优化器
                data_loader=self.trainer.train_dataloader,
                noise_multiplier=self.hparams.noise_multiplier,
                max_grad_norm=self.hparams.max_grad_norm
            )

    # 无需重写forward方法，make_private已自动替换模型的forward逻辑

核心要点：

• 不要封装后再手动调用dp_model，make_private会直接替换当前模型的参数、梯度计算逻辑，原有的forward调用会自动走DP流程。
• 必须在on_train_start钩子中执行封装：此时Lightning已完成模型设备迁移、优化器初始化，符合Opacus的要求。

2. 是否需要修改Lightning访问模型参数的方式？ #

不需要修改Lightning默认的参数访问逻辑，但要注意两个细节：

• 确保make_private在模型参数被移动到设备后执行（on_train_start是合适的时机），避免设备不匹配问题。
• 不要手动修改模型参数的引用，make_private会自动处理参数的隐私化包装。

内容的提问来源于stack exchange，提问作者drey_1