colossalai.utils.common

colossalai.utils.common.print_rank_0(msg, logger=None)

Print messages and save logs(optional). This is executed only if you are the rank-0 gpu.

Parameters

• msg (str) – A string message to output.

• logger (colossalai.logging.DistributedLogger, optional) – The logger to record the message, defaults to None.

colossalai.utils.common.sync_model_param(model, parallel_mode)

Make sure data parameters are consistent during Data Parallel Mode.

Parameters

• model (torch.nn.Module) – A pyTorch model on whose parameters you check the consistency.

• parallel_mode (colossalai.context.ParallelMode) – Parallel mode to be checked.

Note

The parallel_mode should be concluded in ParallelMode. More details about ParallelMode could be found in parallel_mode

colossalai.utils.common.clip_grad_norm_fp32(parameters, max_norm, norm_type=2)

Clips gradient norm of an iterable of parameters whose gradients are in fp32.

This is adapted from torch.nn.utils.clip_grad.clip_grad_norm_() and added functionality to handle model parallel parameters.

Note

the gradients are modified in place.

Parameters

• parameters (Iterable[torch.tensor] or torch.tensor) – An iterable of Tensors or a single Tensor that will have gradients normalized.

• max_norm (Union[float, int]) – Max norm of the gradients.

• norm_type (Union[float, int, 'inf']) – Type of the used p-norm. Can be 'inf' for infinity norm.

Returns

Total norm of the parameters.

Return type

float