DPOptimizer速度太慢

[xiehuanyi]

我使用DPOptimizer微调GPT，速度太慢，我使用了四条数据，batchsize设置为2，使用正常微调20epochs，花费<4min，但是使用DPOptimizer则无法微调，大概过了40分钟，被迫终止了运行。
我使用的包版本如下：

mindarmour                               1.8.0
mindformers                              0.3.0
mindinsight                              1.8.0
mindspore-ascend                         1.8.1
mindx-elastic                            0.0.1
modelarts-mindspore-model-server         1.0.4

因为每次使用DPOptimier都无法得到运行结果，所以没有具体的时间，我使用的代码如下：
因为微调的代码有数据集，不方便复现。我使用下面的代码也遇到了跑不出结果的问题，请问要如何解决？

from mindformers import GPT2LMHeadModel, GPT2Tokenizer
from mindarmour.privacy.diff_privacy import DPOptimizerClassFactory
import mindspore as ms
model = GPT2LMHeadModel.from_pretrained('gpt2')
model.set_train(False)
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

GaussianSGD = DPOptimizerClassFactory(micro_batches=2)
GaussianSGD.set_mechanisms('Gaussian', norm_bound=1.0, initial_noise_multiplier=1.5)
opt = GaussianSGD.create('Momentum')(params=model.trainable_params(),
                                         learning_rate=0.001,
                                         momentum=0.9)
# opt = ms.nn.Adam(model.trainable_params())
grad_fn = ms.ops.value_and_grad(model, None, opt.parameters, has_aux=False)

model.set_train(True)
inputs = tokenizer(["hello world"],
                   padding='max_length',
                   max_length=model.config.seq_length+1,
                   return_tensors='ms')
# output = model(input_ids=inputs["input_ids"])
# print(output)  # 计算loss
loss, grad = grad_fn(inputs['input_ids'])
res = opt(grad)
print(loss)
print(res)