Calculate l2 norm of parameters
(model)
| 47 | |
| 48 | |
| 49 | def calc_params_l2_norm(model): |
| 50 | """Calculate l2 norm of parameters""" |
| 51 | args = get_args() |
| 52 | if not isinstance(model, list): |
| 53 | model = [model] |
| 54 | # Remove duplicate params. |
| 55 | params_data = [] |
| 56 | for model_ in model: |
| 57 | for param in model_.parameters(): |
| 58 | is_not_shared = param_is_not_shared(param) |
| 59 | is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param) |
| 60 | if is_not_shared and is_not_tp_duplicate: |
| 61 | if args.bf16: |
| 62 | params_data.append(param.data.float()) |
| 63 | else: |
| 64 | params_data.append(param.data) |
| 65 | # Calculate norm |
| 66 | dummy_overflow_buf = torch.cuda.IntTensor([0]) |
| 67 | norm, _ = multi_tensor_applier( |
| 68 | amp_C.multi_tensor_l2norm, |
| 69 | dummy_overflow_buf, |
| 70 | [params_data], |
| 71 | False, # no per-parameter norm |
| 72 | ) |
| 73 | norm_2 = norm * norm |
| 74 | # Sum across all model-parallel GPUs. |
| 75 | torch.distributed.all_reduce( |
| 76 | norm_2, op=torch.distributed.ReduceOp.SUM, group=mpu.get_model_parallel_group() |
| 77 | ) |
| 78 | return norm_2.item() ** 0.5 |
| 79 | |
| 80 | |
| 81 | def average_losses_across_data_parallel_group(losses): |
no test coverage detected