Fix memory regression when using FSDP (#1788)
Fix memory regression when using FSDP (#1788)

Summary:
FSDP overloads nn.Module.apply and inserts a step where it all-gathers params before calling apply (added in https://github.com/facebookresearch/fairscale/commit/fa1b85fbbe75f058b39f1bcf027de42e6ddbd487). This is important for the typical use case of nn.Module.apply -- weight initialization. But here we were using apply totally unnecessarily. It's easier to just loop over all the modules and call set_num_updates directly

Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/1788

Reviewed By: sshleifer

Differential Revision: D27622168

Pulled By: myleott

fbshipit-source-id: f5462107ad251cf7834b20a0eaccbe2f685da8f8

pytorch/fairseq,pytorch/fairseq,pytorch/fairseq