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

Co-authored-by: Nyaozhewei <zheweiy@berkeley.edu>
Co-authored-by: Nxiaoxiawu <yxiaoxiawu@microsoft.com>
Co-authored-by: NConglong Li <conglong.li@gmail.com>
Co-authored-by: NXiaoxia (Shirley) Wu <94406484+xiaoxiawu-microsoft@users.noreply.github.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NMichael Wyatt <michaelwyatt@microsoft.com>

Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NMichael Wyatt <michaelwyatt@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>

* Split parameter offload from z3

* Format fixes

* Bug fixes

* Cleanup

* Remove dead code

* Retain prefetched params until last use

* Unit tests fixes

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

Co-authored-by: NJeff Rasley <jerasley@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>

* Fix do not updated sparse grads

* Remove call .data for sparse grads
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

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

* DeepSpeed needs to start cleaning up

* Remove debug prints

* [pipe] prevent deadlock with multiple evals sequence

* style

* style

* style

* align DSE commit w. latest master
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

* GatheredParameters - accept any iterable

* torch tensor is an iterable, so can't use collections.abc.Iterable

* fix

* Fix OOM and type mismatch

* Toggle prefetching

* Disable z3 prefetching for inference (temp workaround)

* Fix zero3 tracing issues

* Remove debug prints

* Enable prefetch for inference

* Code clarity

* Invalidate trace cache

* Trace cache invalidation when needed
Separate nvme prefetch from all-gather prefetch

* Track last used step id

* Use debug name in error message

* Construct param trace from module trace
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

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

Co-authored-by: NReza Yazdani <reyazda@microsoft.com>
Co-authored-by: NReza Yazdani <44502768+RezaYazdaniAminabadi@users.noreply.github.com>

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

* Fix zero3 tracing issues

* Remove debug prints

* Code clarity

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: NJeff Rasley <jerasley@microsoft.com>

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