mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-06 11:32:10 +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:
Hongxin Liu
GitHub
@@ -80,7 +80,6 @@ def is_using_sequence():
|
||||
|
||||
|
||||
class model_branch_context(object):
|
||||
|
||||
def __enter__(self):
|
||||
self.env_status = env.save()
|
||||
|
||||
@@ -98,16 +97,14 @@ def _calc_l2_norm(grads):
|
||||
|
||||
if fused_optim is None:
|
||||
from colossalai.kernel.op_builder import FusedOptimBuilder
|
||||
|
||||
fused_optim = FusedOptimBuilder().load()
|
||||
|
||||
norm = 0.0
|
||||
if len(grads) > 0:
|
||||
dummy_overflow_buf = torch.cuda.IntTensor([0])
|
||||
norm, _ = multi_tensor_applier(
|
||||
fused_optim.multi_tensor_l2norm,
|
||||
dummy_overflow_buf,
|
||||
[grads],
|
||||
False # no per-parameter norm
|
||||
fused_optim.multi_tensor_l2norm, dummy_overflow_buf, [grads], False # no per-parameter norm
|
||||
)
|
||||
return norm
|
||||
|
||||
@@ -121,7 +118,7 @@ def _calc_lp(grads, norm_type):
|
||||
|
||||
|
||||
def _move_norm_to_cuda(norm: Union[float, torch.Tensor]) -> Union[float, torch.Tensor]:
|
||||
if torch.is_tensor(norm) and norm.device.type != 'cuda':
|
||||
if torch.is_tensor(norm) and norm.device.type != "cuda":
|
||||
norm = norm.to(torch.cuda.current_device())
|
||||
return norm
|
||||
|
||||
@@ -141,11 +138,11 @@ def _compute_local_lp(params: List[ColoParameter], norm_type: float) -> float:
|
||||
if len(params) == 0:
|
||||
return 0.0
|
||||
grads = [p.grad for p in params]
|
||||
use_cuda_kernel = grads[0].device.type == 'cuda'
|
||||
use_cuda_kernel = grads[0].device.type == "cuda"
|
||||
if norm_type == inf:
|
||||
local_lp = max([g.abs().max() for g in grads])
|
||||
elif norm_type == 2.0 and use_cuda_kernel:
|
||||
local_lp = _calc_l2_norm(grads)**norm_type
|
||||
local_lp = _calc_l2_norm(grads) ** norm_type
|
||||
else:
|
||||
local_lp = _calc_lp(grads, norm_type)
|
||||
if isinstance(local_lp, torch.Tensor):
|
||||
@@ -202,8 +199,8 @@ def _compute_grad_lp(parameters, norm_type: float = 2.0) -> float:
|
||||
assert isinstance(p, ColoParameter)
|
||||
if grad_dtype is None:
|
||||
grad_dtype = p.grad.dtype
|
||||
assert p.grad.dtype == grad_dtype, f'Expected all grads are {grad_dtype}, got {p.grad.dtype}'
|
||||
if p.grad.device.type == 'cuda':
|
||||
assert p.grad.dtype == grad_dtype, f"Expected all grads are {grad_dtype}, got {p.grad.dtype}"
|
||||
if p.grad.device.type == "cuda":
|
||||
cuda_grad_params.append(p)
|
||||
else:
|
||||
cpu_grad_params.append(p)
|
||||
@@ -221,7 +218,7 @@ def compute_grad_norm(parameters, norm_type: float = 2.0) -> float:
|
||||
norm_type = float(norm_type)
|
||||
total_norm = _compute_grad_lp(parameters, norm_type)
|
||||
if norm_type != inf:
|
||||
total_norm = total_norm**(1 / norm_type)
|
||||
total_norm = total_norm ** (1 / norm_type)
|
||||
return total_norm
|
||||
|
||||
|
||||
@@ -235,14 +232,15 @@ def _clip_grad_norm(parameters, max_norm: float, total_norm: float) -> None:
|
||||
for p in parameters:
|
||||
if p.grad is None:
|
||||
continue
|
||||
if p.grad.device.type == 'cuda':
|
||||
if p.grad.device.type == "cuda":
|
||||
cuda_grads.append(p.grad.detach())
|
||||
else:
|
||||
cpu_grads.append(p.grad.detach())
|
||||
if len(cuda_grads) > 0:
|
||||
dummy_overflow_buf = torch.cuda.IntTensor([0])
|
||||
multi_tensor_applier(fused_optim.multi_tensor_scale, dummy_overflow_buf, [cuda_grads, cuda_grads],
|
||||
clip_coef)
|
||||
multi_tensor_applier(
|
||||
fused_optim.multi_tensor_scale, dummy_overflow_buf, [cuda_grads, cuda_grads], clip_coef
|
||||
)
|
||||
for g in cpu_grads:
|
||||
g.mul_(clip_coef)
|
||||
|
||||
@@ -284,16 +282,17 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
|
||||
for param in parameters:
|
||||
if param.grad is not None:
|
||||
# Make sure the grads are in fp32
|
||||
assert param.grad.dtype == torch.float, \
|
||||
f'expected gradient to be dtype torch.float, but got {param.grad.type()}'
|
||||
if hasattr(param, 'colo_attr') and param.colo_attr.sharded_data_tensor.is_sharded:
|
||||
assert (
|
||||
param.grad.dtype == torch.float
|
||||
), f"expected gradient to be dtype torch.float, but got {param.grad.type()}"
|
||||
if hasattr(param, "colo_attr") and param.colo_attr.sharded_data_tensor.is_sharded:
|
||||
has_zero_shared_param = True
|
||||
params.append(param)
|
||||
|
||||
if len(params) == 0:
|
||||
enable_cuda_kernels = False
|
||||
else:
|
||||
enable_cuda_kernels = params[0].grad.device.type == 'cuda'
|
||||
enable_cuda_kernels = params[0].grad.device.type == "cuda"
|
||||
# Norm parameters.
|
||||
max_norm = float(max_norm)
|
||||
norm_type = float(norm_type)
|
||||
@@ -307,15 +306,13 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
|
||||
total_norm_cuda = torch.cuda.FloatTensor([float(total_norm)])
|
||||
# Take max across all model-parallel GPUs.
|
||||
if gpc.is_initialized(ParallelMode.MODEL) and gpc.get_world_size(ParallelMode.MODEL) > 1:
|
||||
dist.all_reduce(total_norm_cuda,
|
||||
op=dist.ReduceOp.MAX,
|
||||
group=gpc.get_group(ParallelMode.MODEL),
|
||||
async_op=False)
|
||||
dist.all_reduce(
|
||||
total_norm_cuda, op=dist.ReduceOp.MAX, group=gpc.get_group(ParallelMode.MODEL), async_op=False
|
||||
)
|
||||
if has_zero_shared_param:
|
||||
dist.all_reduce(total_norm_cuda,
|
||||
op=dist.ReduceOp.MAX,
|
||||
group=gpc.get_group(ParallelMode.DATA),
|
||||
async_op=False)
|
||||
dist.all_reduce(
|
||||
total_norm_cuda, op=dist.ReduceOp.MAX, group=gpc.get_group(ParallelMode.DATA), async_op=False
|
||||
)
|
||||
total_norm = total_norm_cuda[0].item()
|
||||
else:
|
||||
tensor_parallel_grads = []
|
||||
@@ -323,17 +320,17 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
|
||||
zero_sharded_grads = []
|
||||
for p in params:
|
||||
if is_model_parallel_parameter(p):
|
||||
reductor = (gpc.get_world_size(ParallelMode.TENSOR) / getattr(p, NUM_PARTITIONS))**(1 / norm_type)
|
||||
reductor = (gpc.get_world_size(ParallelMode.TENSOR) / getattr(p, NUM_PARTITIONS)) ** (1 / norm_type)
|
||||
tensor_parallel_grads.append(p.grad.data / reductor)
|
||||
elif hasattr(p, 'colo_attr') and p.colo_attr.sharded_data_tensor.is_sharded:
|
||||
elif hasattr(p, "colo_attr") and p.colo_attr.sharded_data_tensor.is_sharded:
|
||||
zero_sharded_grads.append(p.grad.data)
|
||||
else:
|
||||
no_tensor_parallel_grads.append(p.grad.data)
|
||||
|
||||
if norm_type == 2.0 and enable_cuda_kernels:
|
||||
tensor_parallel_norm = _calc_l2_norm(tensor_parallel_grads)**norm_type
|
||||
no_tensor_parallel_norm = _calc_l2_norm(no_tensor_parallel_grads)**norm_type
|
||||
zero_sharded_norm = _calc_l2_norm(zero_sharded_grads)**norm_type
|
||||
tensor_parallel_norm = _calc_l2_norm(tensor_parallel_grads) ** norm_type
|
||||
no_tensor_parallel_norm = _calc_l2_norm(no_tensor_parallel_grads) ** norm_type
|
||||
zero_sharded_norm = _calc_l2_norm(zero_sharded_grads) ** norm_type
|
||||
else:
|
||||
tensor_parallel_norm = _calc_lp(tensor_parallel_grads, norm_type)
|
||||
no_tensor_parallel_norm = _calc_lp(no_tensor_parallel_grads, norm_type)
|
||||
@@ -358,7 +355,7 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
|
||||
total_norm = tensor_parallel_norm + no_tensor_parallel_norm
|
||||
if gpc.is_initialized(ParallelMode.PIPELINE) and gpc.get_world_size(ParallelMode.PIPELINE) > 1:
|
||||
dist.all_reduce(total_norm, op=dist.ReduceOp.SUM, group=gpc.get_group(ParallelMode.PIPELINE))
|
||||
total_norm = total_norm**(1.0 / norm_type)
|
||||
total_norm = total_norm ** (1.0 / norm_type)
|
||||
if torch.is_tensor(total_norm):
|
||||
total_norm = total_norm.item()
|
||||
|
||||
@@ -397,13 +394,14 @@ def count_zeros_fp32(parameters):
|
||||
# Sum across all model-parallel GPUs.
|
||||
ops = []
|
||||
ops.append(
|
||||
dist.all_reduce(total_num_zeros, op=dist.ReduceOp.SUM, group=gpc.get_group(ParallelMode.TENSOR), async_op=True))
|
||||
dist.all_reduce(total_num_zeros, op=dist.ReduceOp.SUM, group=gpc.get_group(ParallelMode.TENSOR), async_op=True)
|
||||
)
|
||||
if gpc.is_initialized(ParallelMode.PIPELINE):
|
||||
ops.append(
|
||||
dist.all_reduce(total_num_zeros,
|
||||
op=dist.ReduceOp.SUM,
|
||||
group=gpc.get_group(ParallelMode.PIPELINE),
|
||||
async_op=True))
|
||||
dist.all_reduce(
|
||||
total_num_zeros, op=dist.ReduceOp.SUM, group=gpc.get_group(ParallelMode.PIPELINE), async_op=True
|
||||
)
|
||||
)
|
||||
|
||||
for req in ops:
|
||||
req.wait()
|
||||
@@ -420,8 +418,9 @@ def copy_tensor_parallel_attributes(src_tensor, dst_tensor):
|
||||
|
||||
|
||||
def param_is_not_tensor_parallel_duplicate(param):
|
||||
return (hasattr(param, IS_TENSOR_PARALLEL) and getattr(param, IS_TENSOR_PARALLEL)) or (gpc.get_local_rank(
|
||||
ParallelMode.TENSOR) == 0)
|
||||
return (hasattr(param, IS_TENSOR_PARALLEL) and getattr(param, IS_TENSOR_PARALLEL)) or (
|
||||
gpc.get_local_rank(ParallelMode.TENSOR) == 0
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
|
||||
Reference in New Issue
Block a user