[Easy][FSDP] Change `assert` -> `p_assert` (#85052)
This changes a few `assert`s to `p_assert()`s because they can run in the backward (some are in the forward, but AC can make them run in the backward).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85052
Approved by: https://github.com/zhaojuanmao
