74 lines
No EOL
2.4 KiB
Python
74 lines
No EOL
2.4 KiB
Python
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
import gc
|
|
import random
|
|
from contextlib import contextmanager
|
|
|
|
import os
|
|
import pytest
|
|
import numpy as np
|
|
import torch
|
|
import torch._dynamo as dynamo
|
|
|
|
|
|
@contextmanager
|
|
def set_seed(seed: int = 0):
|
|
np.random.seed(42)
|
|
torch.manual_seed(42)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False
|
|
torch.use_deterministic_algorithms(True)
|
|
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"
|
|
yield
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_dyno_state():
|
|
cache_limit = dynamo.config.cache_size_limit
|
|
try:
|
|
dynamo.config.cache_size_limit = 512
|
|
dynamo.reset()
|
|
torch.cuda.synchronize()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
yield {}
|
|
except RuntimeError as err:
|
|
raise err
|
|
finally:
|
|
dynamo.config.cache_size_limit = cache_limit
|
|
torch.cuda.synchronize()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def assert_all_close(a: torch.Tensor, b: torch.Tensor, rtol=0, atol=1e-1) -> None:
|
|
"""
|
|
Check that all elements of tensors a and b are within provided thresholds.
|
|
"""
|
|
assert a.shape == b.shape, f"Shapes don't match: {a.shape} != {b.shape}"
|
|
assert a.dtype == b.dtype, f"Dtypes don't match: {a.dtype} != {b.dtype}"
|
|
assert a.device == b.device, f"Devices don't match: {a.device} != {b.device}"
|
|
max_abs_diff = torch.max(torch.abs(a - b))
|
|
rel_diff = torch.abs(a / b)
|
|
max_rel_diff = torch.max(rel_diff)
|
|
mismatch_elements = torch.sum(torch.abs(a - b) > atol + rtol * torch.abs(b))
|
|
nb_elements = torch.numel(a)
|
|
msg = (
|
|
f"Differences: "
|
|
f"{max_abs_diff:.3f} (max abs), "
|
|
f"{max_rel_diff:.3f} (max rel), "
|
|
f"{mismatch_elements}/{nb_elements} (mismatch elements)"
|
|
)
|
|
assert torch.allclose(a, b, rtol=rtol, atol=atol), msg |