dela
584f7e9e73
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
# For a list of all contributors, visit:
|
|
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
|
|
|
|
import logging
|
|
import warnings
|
|
|
|
import torch
|
|
|
|
from ._config import FLA_CI_ENV
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_abs_err(x, y):
|
|
return (x.detach() - y.detach()).flatten().abs().max().item()
|
|
|
|
|
|
def get_err_ratio(x, y):
|
|
err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item()
|
|
base = (x.detach()).flatten().square().mean().sqrt().item()
|
|
return err / (base + 1e-8)
|
|
|
|
|
|
def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6):
|
|
abs_atol = get_abs_err(ref, tri)
|
|
error_rate = get_err_ratio(ref, tri)
|
|
msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {error_rate:.6f}"
|
|
logger.info(msg)
|
|
if abs_atol <= err_atol:
|
|
return
|
|
assert not torch.isnan(ref).any(), f"{prefix}: NaN detected in ref"
|
|
assert not torch.isnan(tri).any(), f"{prefix}: NaN detected in tri"
|
|
if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)):
|
|
if error_rate > ratio:
|
|
warnings.warn(msg)
|
|
else:
|
|
assert error_rate < ratio, msg
|