Regarding the third point, maybe my description above is not clear enough.
I compared the GPU memory usage at each step of the code that does not use the CUDA extension and the code that uses the CUDA extension. Both use the exact same optimizer. My cuda extension only replaces the layer modules in the model.
The former only adds about 5M of GPU memory usage during the first execution of optimizer.step(), while the latter adds about 1G of GPU memory usage. There is a huge difference between the two.
I don’t know why the CUDA extension causes such a difference? Can you provide some possible clues?