Diff Coverage

def _all_reduce_offloaded(
    tensor: torch.Tensor, op: dist.ReduceOp, group: object,
) -> torch.Tensor:
    """Reduce CPU-offloaded values on the backend's communication device."""
    original_device = tensor.device
    if original_device.type == "cpu":
        device_type = {"nccl": "cuda", "hccl": "npu"}.get(dist.get_backend(group))
        if device_type is not None:
            device_module = torch.get_device_module(device_type)
            tensor = tensor.to(torch.device(device_type, device_module.current_device()))
    dist.all_reduce(tensor, op=op, group=group)
    return tensor.to(original_device)


def _get_grad_obj(param: torch.nn.Parameter) -> Optional[torch.Tensor]:
    """Return the gradient object for *param*.

        local_norm = _compute_local_norm(grads, norm_type, device)
        if mesh_id is not None:
            mesh = mesh_cache[mesh_id]
            for dim in shard_dims:
                local_norm = _all_reduce_offloaded(local_norm, reduce_op, mesh.get_group(dim))
        group_norms.append(local_norm)
    if not group_norms:
        if norm_type == -math.inf:
            return torch.tensor(float("inf"), device=device)

        local_val = _compute_local_norm(grads, norm_type, device)
        if mesh_id is not None:
            mesh = mesh_cache[mesh_id]
            for dim in shard_dims:
                local_val = _all_reduce_offloaded(local_val, dist.ReduceOp.SUM, mesh.get_group(dim))
        total.add_(local_val)
    return total

        else:
            local_p = torch.tensor(0.0, device=device, dtype=torch.float32)

        for group in sig_groups[sig]:
            local_p = _all_reduce_offloaded(local_p, dist.ReduceOp.SUM, group)

        total_p = total_p + local_p

    return total_p

        buf = torch.cat(chunks)

        for pdim, reduce_op, needs_avg in partial_info:
            group = mesh.get_group(pdim)
            buf = _all_reduce_offloaded(buf, reduce_op, group)
            if needs_avg:
                buf /= dist.get_world_size(group=group)

        # Extract views for params with actual gradients.