github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-07 20:10:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'hpcaitech:main' into feature/fp8_comm

Browse Source
This commit is contained in:
Hanks
2024-07-04 20:34:41 +08:00
committed by GitHub
parent f5a52e1600 7997683aac
commit 6991819a97
63 changed files with 265 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
colossalai/quantization/bnb.py
View File

@@ -1,17 +1,25 @@
# adapted from Hugging Face accelerate/utils/bnb.py accelerate/utils/modeling.py
import importlib.metadata
import logging
import torch
import torch.nn as nn
from packaging.version import Version
from .bnb_config import BnbQuantizationConfig
try:
import bitsandbytes as bnb
IS_4BIT_BNB_AVAILABLE = bnb.__version__ >= "0.39.0"
IS_8BIT_BNB_AVAILABLE = bnb.__version__ >= "0.37.2"
try:
# in case lower version of bitsandbytes does not have __version__ attribute
BNB_VERSION = Version(bnb.__version__)
except AttributeError:
BNB_VERSION = Version(importlib.metadata.version("bitsandbytes"))
IS_4BIT_BNB_AVAILABLE = BNB_VERSION >= Version("0.39.0")
IS_8BIT_BNB_AVAILABLE = BNB_VERSION >= Version("0.37.2")
except ImportError:
pass