[FSDP] Add/refactor unit test for wrap (#71803)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71803
1. Extra check for wrapping with override args,
2. Enhance UT to make sure
`wrap` doesn't wrap outside of ctx.
ghstack-source-id: 147753225
Test Plan: CI
Reviewed By: zhaojuanmao
Differential Revision: D33774512
fbshipit-source-id: 1f8d60bdf9b3ba257fee465064a0e25235b3622b
(cherry picked from commit 9ab775b29eddcd193c11398184bee8beffed0327)
