Sequential offloading doesn't work when using pytest, but does seem to work outside of tests.
This is an issue, because we can't properly test sequential offloading on Stable Audio.
Additional information: I've managed to pinpoint where the first NaN appears. It's in the autoencoder level when passing through the decoder's conv1. It might has to do with the fact that we're using weight_norm, but at this point, I really don't know since it works outside of the testing file.
Reproduction
Offloading with NaN
test_sequential_cpu_offload_forward_pass and test_sequential_offload_forward_pass_twice
Describe the bug
Sequential offloading doesn't work when using
pytest
, but does seem to work outside of tests. This is an issue, because we can't properly test sequential offloading on Stable Audio.Additional information: I've managed to pinpoint where the first
NaN
appears. It's in theautoencoder
level when passing through the decoder'sconv1
. It might has to do with the fact that we're usingweight_norm
, but at this point, I really don't know since it works outside of the testing file.Reproduction
Offloading with NaN
test_sequential_cpu_offload_forward_pass
andtest_sequential_offload_forward_pass_twice
Working offloading
Logs
No response
System Info
Who can help?
cc @yiyixuxu @sayakpaul