Skip to content

Automatic mixed precision#2359

Merged
ndryden merged 52 commits into
LBANN:developfrom
ndryden:amp
Nov 9, 2023
Merged

Automatic mixed precision#2359
ndryden merged 52 commits into
LBANN:developfrom
ndryden:amp

Conversation

@ndryden

@ndryden ndryden commented Oct 19, 2023

Copy link
Copy Markdown
Collaborator

This is an initial pass at AMP. It appears correct in the cases I've tested.

It has some particularly pointy bits in batchnorm and in gradient unscaling.

@benson31 benson31 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm still working on the changes model.cpp and the amp.{cpp,cu} files. I suppose I should also read the .py file too. But just a few small things.

Comment thread applications/vision/resnet.py
Comment thread applications/vision/lenet.py
Comment thread include/lbann/layers/transform/evaluation.hpp Outdated
Comment thread src/layers/regularizers/batch_normalization_builder.cpp Outdated
Comment thread include/lbann/utils/print_helpers.hpp Outdated
Comment thread src/utils/cudnn.cpp Outdated
Comment thread include/lbann/layers/regularizers/batch_normalization.hpp
Comment thread src/utils/amp.cu Outdated

@benson31 benson31 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Still looking good after the rebase/recent commits. I touched the memory_profile callback and added an example of how to use the SwitchDispatcher. I found it a bit clunkier than I remembered to use, but I think I have some ideas to improve it in future H2 versions (in particular, I want to make it possible to just use a lambda as the functor).

Also, please clang-format the whole PR.

@ndryden ndryden marked this pull request as ready for review October 30, 2023 18:34
@ndryden ndryden force-pushed the amp branch 2 times, most recently from dc5c648 to b5b91c8 Compare November 1, 2023 17:04
Comment thread src/models/model.cpp Outdated
ndryden and others added 26 commits November 9, 2023 11:26
Co-authored-by: Tom Benson <benson31@llnl.gov>
Co-authored-by: Tom Benson <benson31@llnl.gov>
DistData also includes the datatype, and so spuriously indicated
things were sharded when mixing datatypes.
Co-authored-by: Tom Benson <benson31@llnl.gov>
Co-authored-by: Tom Benson <benson31@llnl.gov>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants