* Callable option for optimizer and scheduler

* Add unit test

* Formatting

* Disable debug prints

* Use base optimizer to construct lr scheduler

* Formatting

* Remove dead import

Co-authored-by: NAlex Muzio <Alex.Muzio@microsoft.com>
Co-authored-by: NAmmar Ahmad Awan <ammar.awan@microsoft.com>
Co-authored-by: NConglong Li <conglong.li@gmail.com>
Co-authored-by: NFelipe Cruz Salinas <Andres.Cruz@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NReza Yazdani <reyazda@microsoft.com>
Co-authored-by: NSamyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: NShaden Smith <shaden.smith@microsoft.com>
Co-authored-by: NYoung Jin Kim <youki@microsoft.com>
Co-authored-by: Nbapatra <bapatra@microsoft.com>
Co-authored-by: NSamyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: NShaden Smith <shaden.smith@microsoft.com>
Co-authored-by: NYoung Jin Kim <youki@microsoft.com>

* Use correct input size for splits

* Use smarter partitioning

* fix empty grad zero tests

* dont clear grads in stage 1 code path

* prevent none grads from being reduced

* Make round robin gradient partitioning configurable (default False)

* Use the correct default

* Log config setting

* zero_param_shapes: switch to round_robin_fp16_groups

* add test

* old torch workaround

* fix reference counting in backward over multiple forwards

* test + cleanup

* Disable copy stream

* Format fixes

* Remove debug codes

* Remove debug codes

* Fix indent

Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* add live zero checkpoint to fp32 consolidation version

* some more docs

* zero2 model states uses a different filename

* fix

* make debug mode cli configurable

* copy the script only on node 0 process 0

* validate that we have the right number of files

* revamp _get_zero_param_shapes, instrument with easier debug

* correct assertion

* rename API; add even simpler API

* style

* docs improve

* update the docs

* revert the unpartitioned_params detection and report as it's most likely persistent params
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* post_init to be run only by a child module

* better solution

* add test

* safer attr name

* wants half()

* improve doc
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* [zero3] params_to_reduce isn't always there

Trying to port HF's Electra model's to Deepspeed I'm getting this on the very first backward step (with some extra debug):

```
Incrementing with parameter id 42
------ Before allocating allgather param name=generator_lm_head.weight id=41 shape=torch.Size([1]) status=ZeroParamStatus.NOT_AVAILABLE partition size=327680
------allgather param with name=generator_lm_head.weight id=41 shape=torch.Size([1]) status=ZeroParamStatus.NOT_AVAILABLE partition size=327680
------ Before allocating allgather param name=generator_lm_head.bias id=42 shape=torch.Size([1]) status=ZeroParamStatus.NOT_AVAILABLE partition size=5120
------allgather param with name=generator_lm_head.bias id=42 shape=torch.Size([1]) status=ZeroParamStatus.NOT_AVAILABLE partition size=5120
Backward name=generator_lm_head.weight id=41 shape=torch.Size([5120, 64])
Inside reduce ipg buckets. name=generator_lm_head.weight id=41 shape=torch.Size([5120, 64]), ipg elements 0, reduce bucket size 4096
Params in ipg bucket []
Reducing []
GOT 1
torch.Size([4096])
Traceback (most recent call last):
  File "examples/pytorch/language-modeling/run_mlm.py", line 533, in <module>
    main()
  File "examples/pytorch/language-modeling/run_mlm.py", line 484, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/mnt/nvme1/code/huggingface/transformers-ds-zero_to_fp32-tests/src/transformers/trainer.py", line 1269, in train
    tr_loss += self.training_step(model, inputs)
  File "/mnt/nvme1/code/huggingface/transformers-ds-zero_to_fp32-tests/src/transformers/trainer.py", line 1778, in training_step
    loss = self.deepspeed.backward(loss)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/engine.py", line 1188, in backward
    self.optimizer.backward(loss)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/zero/stage3.py", line 2964, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/stas/anaconda3/envs/py38-pt19/lib/python3.8/site-packages/torch/autograd/__init__.py", line 147, in backward
    Variable._execution_engine.run_backward(
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/zero/stage3.py", line 1867, in reduce_partition_and_remove_grads
    self.reduce_ready_partitions_and_remove_grads(param, i)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/zero/stage3.py", line 2212, in reduce_ready_partitions_and_remove_grads
    self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/zero/stage3.py", line 1897, in reduce_independent_p_g_buckets_and_remove_grads
    self.reduce_ipg_grads()
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/zero/stage3.py", line 2193, in reduce_ipg_grads
    self.average_tensor(reduction_list, params_to_reduce)
  File "/mnt/nvme1/code/github/00optimize/DeepSpeed-zero-init-child-only-post_init/deepspeed/runtime/zero/stage3.py", line 1972, in average_tensor
    params_to_reduce[0].reduce_gradients_at_owner(
```

Is it always that `params_to_reduce` is populated?

If I add this check the problem goes away it seems.

* real fix

* contiguous gradients should be set to True by default

* Set contiguous gradients to True by default

Features such as reduce_scatter depends on contiguous gradients being True. This is also the preferred default configuration.

Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* undo noise

* another

Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

Co-authored-by: NSamyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* largest_partitioned_params calculation fix

largest partitioned params was getting calculated incorrectly

* Update stage3.py

* Update stage3.py

* formatting fix

* changing sub-group size default to 1e9
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* fix missed subclassed partitioning bug

* fix on exit
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* fix Reduce Scatter default value

* Update constants.py
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* Align fp16 param wap buffers

* Integrating swap buffer manager for fp16 params

* Support swapping misaligned fp16 parameters

* Support swap into unaligned fp16 buffer

* Align fp16 param wap buffers

* Integrating swap buffer manager for fp16 params

* Support swapping misaligned fp16 parameters

* Round robin partitioning to improve ZeRO-2 Offload CPU copy

* Formatting fixes

* Fix index issues in debug dumps

* Remove debug prints

* Code cleanup

* Remove unintended stage3.py changes

* Add TODO

Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* Unused parameters assert should be disabled by default

* Fix message

* Invert assert logic in unit test

* Change option for ignoring unused parameters
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* Add additional conditions when checking types of output from the model

* Add test

* Modify test to use torch.tensor as well
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* Add find_unused_parameters option

As unused parameters in modules may not be expected sometimes, 
add an explicit error msg when it occurred and an option to avoid the error: https://github.com/microsoft/DeepSpeed/issues/707

* Add find_unused_parameters option

As unused parameters in modules may not be expected sometimes, 
add an explicit error msg when it occurred and an option to avoid the error: https://github.com/microsoft/DeepSpeed/issues/707

* Fix syntax error

* Fix yapf error

* Fix yapf error

* Fix yapf error

* Fix yapf error

* Move stage2 find_unused_parameters to config file

* Add stage2 find_unused_parameters

* Add stage2 find_unused_parameters

* Add stage2_find_unused_parameters option

* Change error msg to reflect zero_optimization config change

* Fix yapf error

* Fix yapf errors

* Change find_unused_parameters option name

* Change find_unused_parameters option name

* Change find_unused_parameters option name

* Change find_unused_parameters option name

* Change find_unused_parameters option name

* Add UnusedParametersModel for test option find_unused_parameters

* Add unit test for stage2 find_unused_parameters

* Add cpu-adam compatible check

* Remove dups import

* Trim spaces

* Fix yapf errors

* Trim spaces

* Add False Positive test check

* Fix find_unused_parameters test

* Trim spaces

* Fix yapf error

* Use amp autocast in ZeRO3 linear

* Fix typo

* Handle specific exceptions

* CI breaks on torch.distributed

* Add autocast unit test

* Format fixes

* Fix skip logic
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* Fix issue where gradient_predivide_factor was called as a func.

`gradient_predivide_factor` is a `float`, hence shouldn't be called as func.
This crashes when `reduce_scatter` flag is set to `False`.