Non shared params is empty when large translation model is traind by FSDP.

I'm trying to train translation model which is using transformer encoder-decoder model with FSDP.
But the same problem as https://github.com/facebookresearch/fairscale/issues/883 occures even if fairscale 0.4.5 fairseq 1.0.0a0+4a7835b.
When I run the following command, I got the same error.

``` bash
fairseq-train data-bin/iwslt14.tokenized.de-en --arch transformer_iwslt_de_en --share-decoder-input-output-embed --optimizer cpu_adam --adam-betas '(0.9, 0.98)' --clip-norm 0.0 --lr 5e-4 --lr-scheduler inverse_sqrt --warmup-updates 4 --dropout 0.3 --weight-decay 0.0001 --criterion label_smoothed_cross_entropy --label-smoothing 0.1 --max-tokens 4096 --max-update 4 --ddp-backend fully_sharded --fp16 --fp16-init-scale 4 --cpu-offload --log-format json --log-interval 1 --encoder-layers 6 --decoder-layers 6 --encoder-embed-dim 5120 --decoder-embed-dim 5120 --encoder-ffn-embed-dim 5120 --decoder-ffn-embed-dim 5120 --encoder-attention-heads 40 --decoder-attention-heads 40
```

```
Traceback (most recent call last):
File "/home/christopher/gitrepos/fairseq/env/bin/fairseq-train", line 33, in
sys.exit(load_entry_point('fairseq', 'console_scripts', 'fairseq-train')())
File "/home/christopher/gitrepos/fairseq/fairseq_cli/train.py", line 528, in cli_main
distributed_utils.call_main(cfg, main)
File "/home/christopher/gitrepos/fairseq/fairseq/distributed/utils.py", line 344, in call_main
torch.multiprocessing.spawn(
File "/home/christopher/gitrepos/fairseq/env/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 230, in spawn
return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
File "/home/christopher/gitrepos/fairseq/env/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
while not context.join():
File "/home/christopher/gitrepos/fairseq/env/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 150, in join
raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
File "/home/christopher/gitrepos/fairseq/env/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
fn(i, *args)
File "/home/christopher/gitrepos/fairseq/fairseq/distributed/utils.py", line 328, in distributed_main
main(cfg, **kwargs)
File "/home/christopher/gitrepos/fairseq/fairseq_cli/train.py", line 188, in main
valid_losses, should_stop = train(cfg, trainer, task, epoch_itr)
File "/usr/lib/python3.9/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/home/christopher/gitrepos/fairseq/fairseq_cli/train.py", line 317, in train
valid_losses, should_stop = validate_and_save(
File "/home/christopher/gitrepos/fairseq/fairseq_cli/train.py", line 414, in validate_and_save
checkpoint_utils.save_checkpoint(
File "/home/christopher/gitrepos/fairseq/fairseq/checkpoint_utils.py", line 50, in save_checkpoint
trainer.consolidate_optimizer() # TODO(SS): do we need this if no_save_optimizer_state
File "/home/christopher/gitrepos/fairseq/fairseq/trainer.py", line 380, in consolidate_optimizer
st = self.model.gather_full_optim_state_dict(
File "/home/christopher/gitrepos/fairseq/env/lib/python3.9/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 2267, in gather_full_optim_state_dict
state, singleton_state = self._gather_optim_state(sd.pop("state"))
File "/home/christopher/gitrepos/fairseq/env/lib/python3.9/site-packages/fairscale/nn/data_parallel/fully_sharded_data_parallel.py", line 2211, in _gather_optim_state
assert (
AssertionError: Only flatten param or a single non-shared param is supported: len=0
```

My environment is below.
- PyTorch: 1.10.2
- Fairseq: 1.0.0a0+4a7835b
- Fairscale: 0.4.5
- GeForce Titan RTX x2

`params` is empty because `encoder` and `decoder` the transformer model of are both sharded,
https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L414

Then `to_be_flatten` is empty, and therefore self._fsdp_wrapped_module is empty.
https://github.com/facebookresearch/fairscale/blob/main/fairscale/nn/data_parallel/fully_sharded_data_parallel.py#L442

As a result, non_shared_params is empty, which seems to satisfy the assert condition.

Is there any workaround?