Move `fairscale::with_wrapped_modules` logic to a separate argument in the trainer

Rather than wrapping the model before passing it to the trainer. We want to do this because the `checkpoint_wrapper` in `with_wrapped_modules` creates a `weakref` for the wrapped modules, rendering the model unpicklable.