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 (optional) – Python logger object, 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 nn.model on whose parameters you check the consistency

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

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 that the gradients are modified in place.

Parameters

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

  • max_norm (float or int) – Max norm of the gradients

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

Returns

Total norm of the parameters (viewed as a single vector).

Return type

float