mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-01 09:07:51 +00:00
[misc] update pre-commit and run all files (#4752)
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
This commit is contained in:
@@ -18,17 +18,24 @@ from colossalai.legacy.global_variables import tensor_parallel_env as env
|
||||
def get_depth_from_env() -> int:
|
||||
try:
|
||||
depth = env.depth_3d
|
||||
assert depth > 0, 'DEPTH must be greater than zero'
|
||||
assert depth > 0, "DEPTH must be greater than zero"
|
||||
return depth
|
||||
|
||||
except KeyError as e:
|
||||
raise EnvironmentError('DEPTH is not found in the current environment, '
|
||||
'please make sure that you have used the correct process group initializer')
|
||||
except KeyError:
|
||||
raise EnvironmentError(
|
||||
"DEPTH is not found in the current environment, "
|
||||
"please make sure that you have used the correct process group initializer"
|
||||
)
|
||||
|
||||
|
||||
def get_parallel_mode_from_env(group):
|
||||
assert group in [INPUT_GROUP_3D, WEIGHT_GROUP_3D, OUTPUT_GROUP_3D, INPUT_X_WEIGHT_3D, OUTPUT_X_WEIGHT_3D], \
|
||||
f'{group} is not valid for 3D tensor parallelism.'
|
||||
assert group in [
|
||||
INPUT_GROUP_3D,
|
||||
WEIGHT_GROUP_3D,
|
||||
OUTPUT_GROUP_3D,
|
||||
INPUT_X_WEIGHT_3D,
|
||||
OUTPUT_X_WEIGHT_3D,
|
||||
], f"{group} is not valid for 3D tensor parallelism."
|
||||
return getattr(env, group)
|
||||
|
||||
|
||||
@@ -44,12 +51,10 @@ def dbg_check_shape(tensor: Tensor, shape: tuple):
|
||||
rank = gpc.get_global_rank()
|
||||
if rank == 0:
|
||||
print(tensor.shape)
|
||||
assert tensor.shape == shape, \
|
||||
'{} does not match {}'.format(tensor.shape, shape)
|
||||
assert tensor.shape == shape, "{} does not match {}".format(tensor.shape, shape)
|
||||
|
||||
|
||||
class AsyncGradientBucket(object):
|
||||
|
||||
def __init__(self):
|
||||
self.bucket = OrderedDict()
|
||||
|
||||
|
Reference in New Issue
Block a user