github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[hotfix] removed unused flag (#5242)

pull/5278/head
Frank Lee 2024-01-09 14:57:07 +08:00 committed by GitHub
parent d202cc28c0
commit 9102d655ab
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/accelerator/npu_accelerator.py
View File

@ -7,11 +7,8 @@ import torch.distributed as dist
from .base_accelerator import BaseAccelerator from .base_accelerator import BaseAccelerator
IS_NPU_AVAILABLE = False
try: try:
import torch_npu # noqa import torch_npu # noqa
IS_NPU_AVAILABLE = True
except ImportError: except ImportError:
pass pass

4
colossalai/booster/plugin/gemini_plugin.py
View File

@ -12,7 +12,7 @@ from torch.optim import Optimizer
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from colossalai.accelerator import IS_NPU_AVAILABLE, get_accelerator from colossalai.accelerator import get_accelerator
from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO, GeneralCheckpointIO
from colossalai.checkpoint_io.utils import ( from colossalai.checkpoint_io.utils import (
get_model_base_filenames, get_model_base_filenames,
@ -362,7 +362,7 @@ class GeminiPlugin(DPPluginBase):
) -> None: ) -> None:
super().__init__() super().__init__()
assert precision in SUPPORTED_PRECISION, f"precision {precision} is not supported" assert precision in SUPPORTED_PRECISION, f"precision {precision} is not supported"
if IS_NPU_AVAILABLE: if get_accelerator().name == "npu":
assert placement_policy == "static", "NPU only supports static placement policy" assert placement_policy == "static", "NPU only supports static placement policy"
self.gemini_config = dict( self.gemini_config = dict(
chunk_config_dict=chunk_config_dict, chunk_config_dict=chunk_config_dict,