[BUG] ZeRO-3 can't handle a new Parameter in forward

Describe the bug

This model defines a new nn.Parameter in forward:

https://github.com/huggingface/transformers/blob/e923917cd975d6768d90eb49fdab6468b33b214f/src/transformers/models/m2m_100/modeling_m2m_100.py#L133-L134

and of course deepspeed is not equipped for that scenario and fails:

E               if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
E           AttributeError: 'Parameter' object has no attribute 'ds_status'
E               embed_pos = self.embed_positions(input_ids, inputs_embeds)
E             File "/home/stas/anaconda3/envs/py38-pt110/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1120, in _call_impl
E               result = forward_call(*input, **kwargs)
E             File "/home/stas/anaconda3/envs/py38-pt110/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
E               return func(*args, **kwargs)
E             File "/mnt/nvme1/code/huggingface/transformers-ds-model-zoo-2/src/transformers/models/m2m_100/modeling_m2m_100.py", line 175, in forward
E               self.make_weights(max_pos + self.offset, self.embedding_dim, self.padding_idx)
E             File "/mnt/nvme1/code/huggingface/transformers-ds-model-zoo-2/src/transformers/models/m2m_100/modeling_m2m_100.py", line 134, in make_weights
E               self.weights.requires_grad = False
E             File "/home/stas/anaconda3/envs/py38-pt110/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1168, in __getattr__
E               return _parameters[name]
E             File "/mnt/nvme1/code/github/00optimize/deepspeed/deepspeed/runtime/zero/stage3.py", line 150, in __getitem__
E               if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
E           AttributeError: 'Parameter' object has no attribute 'ds_status'

this param is not known by DS yet it tries to get it for some reason with ZeROOrderedDict.__getitem__ because the sub-module class was created with zero.Init.

It's important to note that it doesn't create a totally new nn.Parameter but it resizes the old one that it created at __init__ but assigns a new variable.

The same code works everywhere else but ZeRO-3.

We have this the same in all positional embedding classes. It creates a param at init, which is all good, but if at run-time a longer positional embedding is required it creates a new param of that longer length, replacing the original one.

We do the same with normal embeddings, but those get resized if needed before the first forward

I guess until now the tests didn't happen to try to extend the length of embeddings, but this one I have added to the model zoo deepspeed tests did trigger this scenario.

@jeffra started working on a solution here: https://github.com/microsoft/DeepSpeed/pull/1606 but it appears to have fallen between the cracks.

So I thought I'd document the original issue properly.

To reproduce

git clone https://github.com/huggingface/transformers
cd transformers
# this should work
CUDA_VISIBLE_DEVICES=0 RUN_SLOW=1 pyt tests/deepspeed/test_model_zoo.py -k test_zero_to_fp32_zero3_trans_m2m_100

# now swapping for a different model and now the test fails
perl -pi -e 's|stas/tiny-m2m_100|hf-internal-testing/tiny-random-m2m_100|' tests/deepspeed/test_model_zoo.py
CUDA_VISIBLE_DEVICES=0 RUN_SLOW=1 pyt tests/deepspeed/test_model_zoo.py -k test_zero_to_fp32_zero3_trans_m2m_100

To bypass the test and going straight for the script that fails:

deepspeed --num_nodes 1 --num_gpus 1 --master_port 10999 \
examples/pytorch/translation/run_translation.py --train_file \
tests/fixtures/tests_samples/wmt_en_ro/train.json --source_lang en \
--target_lang ro --model_name_or_path hf-internal-testing/tiny-random-m2m_100 \
--do_train --max_train_samples 4 --per_device_train_batch_size 2 \
--num_train_epochs 1 --fp16 --report_to none --overwrite_output_dir \
--deepspeed tests/deepspeed/ds_config_zero3.json --output_dir /tmp/tmp2i4fmejh \
--save_steps 1

the key difference between the model that works and one that doesn't - is that the first one has a longish positional embedding tensor and it doesn't get re-created during the first forward and all is good and then model that it fails with has a very short positional embedding and which gets then re-created at forward.

Start with tests/deepspeed/test_model_zoo.py and reproduce test_zero_to_fp32_zero3_trans_m2m_100 using the commands in the issue. Trace the failure through deepspeed/runtime/zero/stage3.py and ZeROOrderedDict.getitem; done means the short positional embedding can be recreated during forward under ZeRO-3 without the missing ds_status error and the reproduction test passes.