Skip to content

[bf16] add bf16 support#3882

Merged
FrankLeeeee merged 6 commits intomainfrom
feature/bf16
Jun 5, 2023
Merged

[bf16] add bf16 support#3882
FrankLeeeee merged 6 commits intomainfrom
feature/bf16

Conversation

@ver217
Copy link
Copy Markdown
Contributor

@ver217 ver217 commented Jun 1, 2023

📌 Checklist before creating the PR

  • I have created an issue for this PR for traceability
  • The title follows the standard format: [doc/gemini/tensor/...]: A concise description
  • I have added relevant tags if possible for us to better distinguish different PRs

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

closed #3839 , closed #3203 and closed #2887 .

📝 What does this PR do?

Summarize your work here.
if you have any plots/diagrams/screenshots/tables, please attach them here.

  1. Add bf16 support for fused adam kernel
  2. Add bf16 support for base adam optimizers, include FusedAdam, CPUAdam and HybridAdam
  3. Implement a helper class MixedPrecisionMixin, by which we can implement mixed precision optimizers easily
  4. Add bf16 support for low level zero, gemini and legacy zero
  5. Update related booster plugins

💥 Checklist before requesting a review

  • I have linked my PR to an issue (instruction)
  • My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
  • I have performed a self-review of my code
  • I have added thorough tests.
  • I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

  • 🌝 Yes, I do.
  • 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

ver217 added 6 commits May 25, 2023 16:59
* [bf16] fused adam kernel support bf16

* [test] update fused adam kernel test

* [test] update fused adam test
…evel zero (#3869)

* [bf16] add mixed precision mixin

* [bf16] low level zero optim support bf16

* [text] update low level zero test

* [text] fix low level zero grad acc test
* [bf16] gemini support bf16

* [test] update gemini bf16 test

* [doc] update gemini docstring
* [zero] init context support bf16

* [zero] legacy zero support bf16

* [test] add zero bf16 test

* [doc] add bf16 related docstring for legacy zero
@ver217 ver217 added the enhancement New feature or request label Jun 1, 2023
@github-actions
Copy link
Copy Markdown
Contributor

github-actions Bot commented Jun 1, 2023

The code coverage for the changed files is 87%.

Click me to view the complete report
Name                                                         Stmts   Miss  Cover
--------------------------------------------------------------------------------
colossalai/amp/naive_amp/mixed_precision_mixin/__init__.py       4      0   100%
colossalai/amp/naive_amp/mixed_precision_mixin/base.py          20      5    75%
colossalai/amp/naive_amp/mixed_precision_mixin/bf16.py          15      1    93%
colossalai/amp/naive_amp/mixed_precision_mixin/fp16.py          49      3    94%
colossalai/booster/plugin/gemini_plugin.py                     109      9    92%
colossalai/booster/plugin/low_level_zero_plugin.py              94      8    91%
colossalai/nn/optimizer/cpu_adam.py                             68      3    96%
colossalai/nn/optimizer/fused_adam.py                           55     10    82%
colossalai/nn/optimizer/hybrid_adam.py                          63      3    95%
colossalai/zero/gemini/gemini_ddp.py                           402     67    83%
colossalai/zero/gemini/gemini_optimizer.py                     195     32    84%
colossalai/zero/legacy/init_ctx/init_context.py                147      4    97%
colossalai/zero/legacy/sharded_model/_utils.py                  57     24    58%
colossalai/zero/legacy/sharded_model/sharded_model_v2.py       299    108    64%
colossalai/zero/legacy/sharded_optim/sharded_optim_v2.py       215     34    84%
colossalai/zero/low_level/low_level_optim.py                   314     20    94%
tests/test_optimizer/test_adam_kernel.py                        87      1    99%
tests/test_optimizer/test_adam_optim.py                         49      0   100%
tests/test_zero/test_gemini/test_optim.py                      130      1    99%
tests/test_zero/test_legacy/test_zero_engine.py                 74      4    95%
tests/test_zero/test_low_level/test_grad_acc.py                 88      1    99%
tests/test_zero/test_low_level/test_zero1_2.py                  90      1    99%
--------------------------------------------------------------------------------
TOTAL                                                         2624    339    87%

@FrankLeeeee FrankLeeeee merged commit ae02d4e into main Jun 5, 2023
@ver217 ver217 deleted the feature/bf16 branch June 6, 2023 03:24
FrankLeeeee pushed a commit to FrankLeeeee/ColossalAI that referenced this pull request Jun 7, 2023
* [bf16] add bf16 support for fused adam (hpcaitech#3844)

* [bf16] fused adam kernel support bf16

* [test] update fused adam kernel test

* [test] update fused adam test

* [bf16] cpu adam and hybrid adam optimizers support bf16 (hpcaitech#3860)

* [bf16] implement mixed precision mixin and add bf16 support for low level zero (hpcaitech#3869)

* [bf16] add mixed precision mixin

* [bf16] low level zero optim support bf16

* [text] update low level zero test

* [text] fix low level zero grad acc test

* [bf16] add bf16 support for gemini (hpcaitech#3872)

* [bf16] gemini support bf16

* [test] update gemini bf16 test

* [doc] update gemini docstring

* [bf16] add bf16 support for plugins (hpcaitech#3877)

* [bf16] add bf16 support for legacy zero (hpcaitech#3879)

* [zero] init context support bf16

* [zero] legacy zero support bf16

* [test] add zero bf16 test

* [doc] add bf16 related docstring for legacy zero
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEATURE]: support BF16 mixed precision training [FEATURE]: BF16 support for LLM [FEATURE]: will zero+bf16 be supported recently ?

2 participants