File size: 5,578 Bytes
2facf1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | import random
import numpy as np
import torch
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def get_device(use_cuda: bool = True) -> torch.device:
"""Get the appropriate device (GPU or CPU)."""
if use_cuda:
if torch.cuda.is_available():
return torch.device("cuda")
elif torch.backends.mps.is_available():
return torch.device("mps")
else:
print("No compatible GPU found. Falling back to CPU.")
return torch.device("cpu")
# Adapted from https://github.com/linkedin/Liger-Kernel/blob/main/test/utils.py
@torch.no_grad()
def verbose_allclose(
received: torch.Tensor,
expected: torch.Tensor,
rtol=1e-05,
atol=1e-08,
max_print=5
) -> list[str]:
"""
Assert that two tensors are element-wise equal within a tolerance, providing detailed information about mismatches.
Parameters:
received (torch.Tensor): Tensor we actually got.
expected (torch.Tensor): Tensor we expected to receive.
rtol (float): Relative tolerance; relative to expected
atol (float): Absolute tolerance.
max_print (int): Maximum number of mismatched elements to print.
Raises:
AssertionError: If the tensors are not all close within the given tolerance.
"""
# Check if the shapes of the tensors match
if received.shape != expected.shape:
return [f"SIZE MISMATCH. received shape: {received.shape}, expected shape: {expected.shape}"]
# Calculate the difference between the tensors
diff = torch.abs(received - expected)
# Determine the tolerance
tolerance = atol + rtol * torch.abs(expected)
# Find tolerance mismatched elements
tol_mismatched = diff > tolerance
# Find nan mismatched elements
nan_mismatched = torch.logical_xor(torch.isnan(received), torch.isnan(expected))
# Find +inf mismatched elements
posinf_mismatched = torch.logical_xor(torch.isposinf(received), torch.isposinf(expected))
# Find -inf mismatched elements
neginf_mismatched = torch.logical_xor(torch.isneginf(received), torch.isneginf(expected))
# Find all mismatched elements
mismatched = torch.logical_or(
torch.logical_or(tol_mismatched, nan_mismatched),
torch.logical_or(posinf_mismatched, neginf_mismatched),
)
mismatched_indices = torch.nonzero(mismatched)
# Count the number of mismatched elements
num_mismatched = mismatched.count_nonzero().item()
# Generate detailed information if there are mismatches
if num_mismatched >= 1:
mismatch_details = [f"Number of mismatched elements: {num_mismatched}"]
for index in mismatched_indices[:max_print]:
i = tuple(index.tolist())
mismatch_details.append(f"ERROR AT {i}: {received[i]} {expected[i]}")
if num_mismatched > max_print:
mismatch_details.append(f"... and {num_mismatched - max_print} more mismatched elements.")
return mismatch_details
return []
@torch.no_grad()
def verbose_allequal(received: torch.Tensor, expected: torch.Tensor, max_print: int=5):
"""
Assert that two tensors are element-wise perfectly equal, providing detailed information about mismatches.
Parameters:
received (torch.Tensor): Tensor we actually got.
expected (torch.Tensor): Tensor we expected to receive.
max_print (int): Maximum number of mismatched elements to print.
Returns:
Empty string if tensors are equal, otherwise detailed error information
"""
mismatched = torch.not_equal(received, expected)
mismatched_indices = torch.nonzero(mismatched)
# Count the number of mismatched elements
num_mismatched = mismatched.count_nonzero().item()
# Generate detailed information if there are mismatches
if num_mismatched >= 1:
mismatch_details = [f"Number of mismatched elements: {num_mismatched}"]
for index in mismatched_indices[:max_print]:
i = tuple(index.tolist())
mismatch_details.append(f"ERROR AT {i}: {received[i]} {expected[i]}")
if num_mismatched > max_print:
mismatch_details.append(f"... and {num_mismatched - max_print} more mismatched elements.")
return mismatch_details
return []
def match_reference(data, output, reference: callable, rtol=1e-05, atol=1e-08):
"""
Convenient "default" implementation for tasks' `check_implementation` function.
"""
output_mla, output_kv = output
# To fit in memory for the big test
output_mla = output_mla.cpu()
output_kv = output_kv.cpu()
config, x, kv_cache = data
expected_mla, expected_kv = reference((config, x, kv_cache))
output_mla = output_mla.cuda()
output_kv = output_kv.cuda()
reasons_mla = verbose_allclose(output_mla, expected_mla, rtol=rtol, atol=atol)
reasons_kv = verbose_allclose(output_kv, expected_kv, rtol=rtol, atol=atol)
if len(reasons_mla) > 0:
return "mismatch found on MLA output! custom implementation doesn't match reference: " + " ".join(reasons_mla)
if len(reasons_kv) > 0:
return "mismatch found on KV cache output! custom implementation doesn't match reference: " + " ".join(reasons_kv)
return ''
def make_match_reference(reference: callable, **kwargs):
def wrapped(data, output):
return match_reference(data, output, reference=reference, **kwargs)
return wrapped |