* unit test, remove exception, add notes

* Move param_shapes to model files

* Remove hard-coded constants

* Conditioned to zero optimizer

* Add zero checkpoint merging

* Print checkpoint version

* Reshape zero_* ckpt files

* Merge zero* files contraction

* Utils for 3D contraction reshaping

* Remove bogus import

* Support bf16_zero ckpts

* Add param slice mappings

* Load universal checkpoints

* Per group mappings from Stas

* Hack to load bf16 zero files

* Param attributes

* WIP

* Fix api bug

* Update lp with local/remote hp

* Disable vocab padding handling

* Update z2 checkpoint

* Remove debug prints

* Remove debug prints; Rebase unit test

* Add reshape assert

* Padding

* Typo

* Catch nonexistent checkpoint path

* Cleanup

* Restore checkpoint state comparisons

* Add torch version guards

* More precise avoidance of false positives.
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* Shards expert parameter groups
* Do upscaling, optimizer and deletion of fp32 grads one-by-one on each parameter group in zero-2
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

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

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

Co-authored-by: NMichael Wyatt <michaelwyatt@microsoft.com>

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

Co-authored-by: NQuentin Anthony <qganthony@yahoo.com>
Co-authored-by: NAmmar Ahmad Awan <ammar.awan@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

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

* bf16 updates

* Got bf16 working

* fp32 reduction; flattened tensors

* bf16+zero_stage_1 first cut

* finish zero_stage 1 sharding

* Matching fp16 with debugging codes

* Matching loss with fp16

* Fix gradient clipping

* bf16 gradient clipping fix
bf16 checkpoint save/load

* Unscale grad norm

* Fix grad norm scaling

* Enable loading fp16_zero_1 into bf16_zero_1 engine and vice versa

* Fix clip_grad key error

* Reduce tied weight gradients

* Fix grad norm for moe

* Reduce specified gradients

* Use O(n) instead of O(n^2)

* Remove optimizer restriction for bf16

* Link bf16 & fp32 params

* Clip gradients of last stage tied weights

* Simplify tied weights reduction logic

* Also clip all tp rank parameters

* lp to hp mapping

* Link lp/hp/optim state; Refresh links after checkpoint load

* Remove debug print

* Remove debug print

* Simplify zero_grad logic

* fp32 accessors

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

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

Co-authored-by: Nyaozhewei <zheweiy@berkeley.edu>
Co-authored-by: NReza Yazdani <reyazda@microsoft.com>

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

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

Co-authored-by: NReza Yazdani <reyazda@microsoft.com>
Co-authored-by: NZhewei Yao <zheweiy@berkeley.edu>
Co-authored-by: NAmmar Ahmad Awan <ammar.awan@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NSamyam Rajbhandari <samyamr@microsoft.com>

[ZeRO] Default disable elastic ckpt in stage 1+2 and reduce CPU memory overhead during ckpt load (#1525)
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

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

* fp16 allreduce

* Undo sparse sum in nan check

* communication_data_type instead of fp32_allreduce and fp16_allreduce

* sparse_allreduce with fp32 or fp16 data type

* FIx communication_data_type checks

* Allow only torch data types for communication_data_type

* Fix Zero assert messages
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* Enforce nccl/rccl alignment of start location of each shard

* Making yapf happy
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

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

* [squash] Staging autotuning v4
Co-authored-by: NCheng Li <pistasable@gmail.com>
Co-authored-by: NMinjia Zhang <minjiaz@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* add new extra, guard xgboost, cleanup dead files (#268)

* Fix autotuning docs (#1553)

* fix docs

* rewording the goal

* fix typos

* fix typos (#1556)

* fix typos

* fix format

* fix bug (#1557)

* fix bug
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NMinjia Zhang <minjiaz@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* Changes for bfloat16 Zero2

* Cleaned up additional comments and debugging code

* Adapted fp16_master_weights_and_grads option to cover BF16

* Reverted fp16_master_weights_and_gradients extension to BFloat16 and minor cleanup

* Fixed formatting and variable naming errors recognized in testing

* Added relevant unit tests for bfloat16 with ZeRO-2

* Updates conditions for skipping BFloat16 unit tests

* Added check for NCCL inconsistent version naming convention

* Update skip message for Bfloat16 tests to mention additional checks
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* Add unit test to check moe+zero checkpoints
* Fix zero stage2 checkpoint loading logic to deal with experts related state dicts.

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

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

* Fix typos in docs/

* Fix typos in code comments and output strings

* Fix typos in the code itself

* Fix typos in tests/
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NShaden Smith <shaden.smith@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NSamyam Rajbhandari <samyamr@microsoft.com>
Co-authored-by: NReza Yazdani <44502768+RezaYazdaniAminabadi@users.noreply.github.com>
Co-authored-by: Neltonzheng <eltonz@microsoft.com>
Co-authored-by: NStas Bekman <stas00@users.noreply.github.com>

* Added 4-byte alignment on NCCL/RCCL

* pre-commit formatting fixes

* Fix for checkpoint loading with optimizer partitioning

* Better assert print

* Added unit tests for nccl/rccl 4-byte alignment

* bug

* Updated alignment to implicit
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* 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>

* 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