35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
import torch
|
|
from typing import Iterable
|
|
|
|
|
|
def calc_diff(x: torch.Tensor, y: torch.Tensor):
|
|
x, y = x.double(), y.double()
|
|
denominator = (x * x + y * y).sum()
|
|
sim = 2 * (x * y).sum() / denominator
|
|
return 1 - sim
|
|
|
|
|
|
def count_bytes(*tensors):
|
|
total = 0
|
|
for t in tensors:
|
|
if isinstance(t, (tuple, list)):
|
|
total += count_bytes(*t)
|
|
elif t is not None:
|
|
total += t.numel() * t.element_size()
|
|
return total
|
|
|
|
def check_signal(num_local_expert, max_m, block_m, threshold, signal, masked_m):
|
|
ceil_div = lambda a, b: (a + b - 1) // b
|
|
|
|
expert_len = max_m // block_m
|
|
for expert in range(num_local_expert):
|
|
mask = masked_m[expert]
|
|
start = expert * expert_len
|
|
end = expert * expert_len + expert_len
|
|
valid_len = ceil_div(mask, block_m)
|
|
for i in range(start, end):
|
|
if i < start + valid_len:
|
|
assert signal[i] == threshold, f'{i=}, {signal[i]=}, {threshold=}'
|
|
else:
|
|
assert signal[i] == 0, f'{i=}, {signal[i]=}'
|