more stringent test for CPUOffloadOptimizer

[ngc92]

Problems with CPUOffloadOptimizer's synchronization (cf. #1649) are not detected by the current test case, as it is very benign, giving ample opportunity for transfers to complete even without explicit sync.

These changes try to make the situation a bit more challenging, decreasing the arithmetic density of the model and ensuring that the critical path is as short as possible.

This PR is deliberately not based on top of #1649, to show that this new test actually catches the synchronization problem. After that PR is merged, tests should pass again.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1650

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 471411d with merge base 3eb18e7 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Code Analysis with Ruff / build (3.9) (gh) (trunk failure)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[gau-nernst]

@ngc92 You can add self.stream.synchronize() at the end of optim step in this PR, then I will convert #1649 to only address the LR scheduler issue.

Thank you again for identifying the issue and coming up with an appropriate test!

[gau-nernst]

Failing ruff check is not caused by this change.