(parameters)
| 124 | |
| 125 | |
| 126 | def count_zeros_fp32(parameters): |
| 127 | |
| 128 | if isinstance(parameters, torch.Tensor): |
| 129 | parameters = [parameters] |
| 130 | |
| 131 | # Filter parameters based on: |
| 132 | # - grad should not be none |
| 133 | # - parameter should not be shared |
| 134 | # - should not be a replica due to tensor model parallelism |
| 135 | total_num_zeros = 0.0 |
| 136 | for param in parameters: |
| 137 | grad_not_none = param.grad is not None |
| 138 | is_not_shared = param_is_not_shared(param) |
| 139 | is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param) |
| 140 | if grad_not_none and is_not_shared and is_not_tp_duplicate: |
| 141 | grad = param.grad.detach() |
| 142 | num_zeros = grad.numel() - torch.count_nonzero(grad) |
| 143 | total_num_zeros = num_zeros + total_num_zeros |
| 144 | |
| 145 | # Sum across all model-parallel GPUs. |
| 146 | torch.distributed.all_reduce( |
| 147 | total_num_zeros, |
| 148 | op=torch.distributed.ReduceOp.SUM, |
| 149 | group=mpu.get_model_parallel_group(), |
| 150 | ) |
| 151 | total_num_zeros = total_num_zeros.item() |
| 152 | |
| 153 | return total_num_zeros |
no test coverage detected