New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

create an engine for all things pipelining and deprecate custom mixed precision #91

Merged

siddharth9820 merged 2 commits into develop from deprecate-mixed-prec-refactor-pp

Jul 1, 2024

Collaborator

siddharth9820 commented Jul 1, 2024 •

edited

Loading

This PR introduces a major refactor of our pipeline parallelism backend.

axonn.inter_layer.AxoNN_Inter_Layer_Engine is being introduced as a class to handle all things pipeline parallelism.
Our custom mixed precision implementation is being deprecated. Instead, the aforementioned class has an argument computation_dtype which can be set to torch.float16, torch.bfloat16, or torch.float32.
Also deprecating cpu_offload optimizer for now. We will revisit it later.


          create an engine for all things pipelining

f118a46

siddharth9820 mentioned this pull request

update to be in line with latest axonn axonn-ai/distrib-dl-tutorial#8

Merged


          refactor tests

9c52d2b

siddharth9820 merged commit ade0c1a into develop

4 checks passed

siddharth9820 deleted the deprecate-mixed-prec-refactor-pp branch

July 1, 2024 15:54

siddharth9820 mentioned this pull request

Deprecate mixed precision support #90

Closed

3 tasks

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet