Skip to content
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

add 8bit quant all reduce support #11

Merged
merged 1 commit into from
Aug 2, 2024
Merged

add 8bit quant all reduce support #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion open_diloco/train_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,9 @@ class HvConfig(BaseConfig):
announce_maddrs: list[str] | None = None
matchmaking_time: float | None = None
averaging_timeout: float | None = None
hivemind_compression: Literal["none", "fp16", "scaled-fp16"] = "none"
hivemind_compression: Literal["none", "fp16", "scaled-fp16", "uniform8bit", "quantile8bit", "blockwise8bit"] = (
"none"
)
all_reduce_strategy: AllReduceStrategy = AllReduceStrategy.WAIT_FOR_ALL
timeout_waiting_for_peers: float | None = None
skip_load_from_peers: bool = False
Expand Down
20 changes: 17 additions & 3 deletions open_diloco/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,10 +108,24 @@ def get_compression_kwargs(hivemind_compression: str) -> dict:

ret_kwargs["grad_compression"] = NoCompression()
ret_kwargs["state_averaging_compression"] = NoCompression()
elif hivemind_compression == "uniform8bit":
from hivemind import Uniform8BitQuantization

ret_kwargs["grad_compression"] = Uniform8BitQuantization()
ret_kwargs["state_averaging_compression"] = Uniform8BitQuantization()
elif hivemind_compression == "quantile8bit":
from hivemind import Quantile8BitQuantization

ret_kwargs["grad_compression"] = Quantile8BitQuantization()
ret_kwargs["state_averaging_compression"] = Quantile8BitQuantization()

elif hivemind_compression == "blockwise8bit":
from hivemind import BlockwiseQuantization

ret_kwargs["grad_compression"] = BlockwiseQuantization()
ret_kwargs["state_averaging_compression"] = BlockwiseQuantization()
else:
raise ValueError(
f"Invalid hivemind_compression: {hivemind_compression}. Please choose 'none', 'fp16', or 'scaled-fp16'."
)
Comment on lines -112 to -114
Copy link
Member Author

Choose a reason for hiding this comment

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

the config is already catching mis naming and propose suggestion thanks to literal support in pydantic

Copy link
Member

Choose a reason for hiding this comment

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

ah ok great!

raise ValueError(f"Invalid hivemind_compression: {hivemind_compression}")
return ret_kwargs


Expand Down
Loading