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

Co-authored-by: NLogan Adams <114770087+loadams@users.noreply.github.com>

* Load z3 checkpoints for inference

* PR feedback

* Fix API bugs

* Fix typo

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

* remove duplicate check for pp and zero stage

* remove line

* use ZeroStageEnum

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

* Option to exclude frozen weights for checkpoint save

* Extend unit test

* Support PP training

Co-authored-by: NLogan Adams <114770087+loadams@users.noreply.github.com>

Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NLogan Adams <114770087+loadams@users.noreply.github.com>

* Add ZeRO 1 support to PP for BF16.

* Switching enum.

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

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

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

* zero3 checkpoint frozen params

* Remove debug prints

* Move to cpu

* WIP

* WIP

* WIP

* Cleanup

* Cleanup

* Extend unit test for frozen params

* API fix

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

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

* Integrate accelerator abstraction interface into deepspeed/

* Fix error message in fp16/fused_optimizer

* fix error message in fp16/unfused_optimizer.py

* assign get_accelerator().pin_memory() result to input Tensor name

* no need to check cuda and whether nvtx supported

* move try-except into inner most block

* call Event() and Stream() in get_accelerator() for data type

* Make Stream and Event as properties of abstract interface so they can be used as data type in deepspeed

* Apply op_builder backend api change from #2705 from @jeffra

* fix tests where Builder NAME is used

* keep original ...Builder.NAME interface instead of ...Builder().NAME interface

* fix builder closure for installation

* fix randomltd builder

* add comments to clarify create_op_builder and get_op_builder

* fix compatibility with pip install -e
Co-authored-by: NCheng Li <pistasable@gmail.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

taking gradient accumulation steps into account for throughput calculation
Co-authored-by: NAlexander Jipa <azzhipa@amazon.com>
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

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

* parallelize layer checkpoints across data parallel groups

* use partition_uniform to determine start/end index values

* formatting fix

* config: add option for parallel write of layer checkpoints in pipeline stage

* yapf fixes

* enable parallel layer write according to config param

* avoid extraneous makedir when rank 0 writes all layers
Co-authored-by: NOlatunji Ruwase <olruwase@microsoft.com>

Co-authored-by: NArpan Jain <t-arpanjain@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

* enable checkpoint engine

* seprated nebula config

* add __init__.py for nebula importing

* linter fix

* fix: ds_config is None

* fix: ds config

* fix: get sd loader fix

* align the API with torch raw code

* linter fix

* remove duplicate tag params

* make checkpoint_engine as required args

* fix args

* extract parameters out to config

* fix: load state dict

* separate load engine

* linter fix

* extract checkpoint engine to abstract calss

* linter fix

* construct function args fix

* add docs for dev/customers

* linter fix

* remove load engine

* print->log_dist

* linter fix

* add tag flag to distinguish the loading order
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: NAmmar Ahmad Awan <ammar.awan@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>

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

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

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

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

* fix pp

* better fix

* CL+PP

* add TODO

Co-authored-by: NThomas <thomas@Thomass-MacBook-Pro.local>
Co-authored-by: NShaden Smith <shaden.smith@microsoft.com>
Co-authored-by: NJeff Rasley <jerasley@microsoft.com>
Co-authored-by: NTunji Ruwase <olruwase@microsoft.com>
Co-authored-by: NStas Bekman <stas00@users.noreply.github.com>

* Add flexibility of pipeline module and engine

* Separate PRs

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