mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-27 19:36:13 +00:00
[hotfix] improve compatibility (#6165)
This commit is contained in:
parent
5f82bfa636
commit
8b0ed61490
@ -1,5 +1,6 @@
|
|||||||
# a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214
|
# a python safetensors serializer modified from https://github.com/huggingface/safetensors/blob/41bd1acf38ad28ac559522d40596c6c802f79453/safetensors/src/tensor.rs#L214
|
||||||
import json
|
import json
|
||||||
|
import warnings
|
||||||
from dataclasses import asdict, dataclass
|
from dataclasses import asdict, dataclass
|
||||||
from typing import Dict, List, Optional, Tuple
|
from typing import Dict, List, Optional, Tuple
|
||||||
|
|
||||||
@ -8,8 +9,10 @@ from safetensors.torch import _TYPES, load_file, safe_open
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
from tensornvme.async_file_io import AsyncFileWriter
|
from tensornvme.async_file_io import AsyncFileWriter
|
||||||
except ModuleNotFoundError:
|
except Exception:
|
||||||
raise ModuleNotFoundError("Please install tensornvme to use NVMeOptimizer")
|
warnings.warn(
|
||||||
|
"Please install the latest tensornvme to use async save. pip install git+https://github.com/hpcaitech/TensorNVMe.git"
|
||||||
|
)
|
||||||
_TYPES_INV = {v: k for k, v in _TYPES.items()}
|
_TYPES_INV = {v: k for k, v in _TYPES.items()}
|
||||||
import io
|
import io
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user