Ensure optimizer in backward works with 2d parallel (#107748)
Summary: Test to ensure optimizer in backward works with 2D parallel.
Test Plan: CI
Differential Revision: D48508057
Pull Request resolved: https://github.com/pytorch/pytorch/pull/107748
Approved by: https://github.com/awgu, https://github.com/fduwjj
