I think algorithms of these two papers may help you.

https://arxiv.org/abs/2110.02711 https://openreview.net/forum?id=pd1P2eUBVfq (I just read it, and so impressed)

Actually there is a gap because of one-timestep. But if betas are sufficiently low, we can ignore the gap.