github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-04-29 04:05:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] improve compatibility (#6165)

Browse Source
This commit is contained in:
Hongxin Liu 2024-12-23 18:57:08 +08:00 committed by GitHub
parent 5f82bfa636
commit 8b0ed61490
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/utils/safetensors.py
View File

@ -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