mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 15:33:51 +08:00
9548d48f48
* Abstract accelerator (step 2) * more flex op_builder path for both installation and runtime * add SpatialInferenceBuilder into cuda_accelerator.py * use reflection to make cuda_accelerator adapt to CUDA op builder change automatically * clean up deepspeed/__init__.py * add comments in cuda_accelerator for no torch path * Update deepspeed/env_report.py Change env_report.py according to suggestion Co-authored-by: Michael Wyatt <mrwyattii@gmail.com> * reduce the range of try...except for better code clarity * Add porting for deepspeed/ops/random_ltd/dropping_utils.py * move accelerator to top directory and create symlink under deepspeed Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com> Co-authored-by: Michael Wyatt <mrwyattii@gmail.com> Co-authored-by: Jeff Rasley <jerasley@microsoft.com>
31 lines
1.2 KiB
Python
31 lines
1.2 KiB
Python
import sys
|
|
import os
|
|
import pkgutil
|
|
import importlib
|
|
|
|
# List of all available op builders from deepspeed op_builder
|
|
|
|
try:
|
|
import op_builder # noqa: F401
|
|
op_builder_dir = "op_builder"
|
|
except ImportError:
|
|
op_builder_dir = "deepspeed.ops.op_builder"
|
|
|
|
op_builder_module = importlib.import_module(op_builder_dir)
|
|
__op_builders__ = []
|
|
|
|
this_module = sys.modules[__name__]
|
|
|
|
# reflect all builder names into variable definition such as 'TransformerBuilder = "TransformerBuilder"'
|
|
for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]):
|
|
# avoid self references
|
|
if module_name != 'all_ops' and module_name != 'builder' and module_name != 'builder_names':
|
|
module = importlib.import_module("{}.{}".format(op_builder_dir, module_name))
|
|
for member_name in module.__dir__():
|
|
if member_name.endswith(
|
|
'Builder'
|
|
) and member_name != "OpBuilder" and member_name != "CUDAOpBuilder":
|
|
# assign builder name to variable with same name
|
|
# the following is equivalent to i.e. TransformerBuilder = "TransformerBuilder"
|
|
this_module.__dict__[member_name] = member_name
|