| import torch |
| import sys |
| import os |
|
|
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) |
|
|
| from arbitor.kernel.ternary_scale import TernaryScaleTensor, TScaleType |
|
|
|
|
| def test_ternary_sign_values(): |
| """T ∈ {-1, 0, 1} — ternary values are pure sign.""" |
| lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) |
| T = lin._get_T() |
| unique = set(T.detach().flatten().tolist()) |
| assert unique.issubset({-1, 0, 1}), f"T has non-ternary values: {unique}" |
| print(" PASS test_ternary_sign_values") |
|
|
|
|
| def test_scale_positive(): |
| """S = 2^E > 0 — scales are always positive.""" |
| lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) |
| S = lin._get_S() |
| assert (S > 0).all(), "Some S values are not positive" |
| assert torch.isfinite(S).all(), "Some S values are non-finite" |
| print(" PASS test_scale_positive") |
|
|
|
|
| def test_effective_weight_polarity(): |
| """W = S * T = {-S, 0, +S} — effective weight is product of sign and scale.""" |
| lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) |
| T = lin._get_T().float() |
| S = lin._get_S() |
| W = S * T |
| T_binary = T.sign().to(torch.int8) |
| assert (W == S * T_binary.float()).all(), "W should equal S * sign(T)" |
| print(" PASS test_effective_weight_polarity") |
|
|
|
|
| def test_e_update_changes_magnitude(): |
| """Modifying E changes W magnitude, not T polarity.""" |
| lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) |
| T_before = lin._get_T().clone() |
| E_orig = lin.E.clone() |
| lin.E[0] = E_orig[0] + 1 |
| W_after = lin.dequantize() |
| T_after = lin._get_T() |
| assert (T_before == T_after).all(), "T changed when only E was modified" |
| S_after = lin._get_S() |
| assert (W_after == S_after * T_after.float()).all(), "W should reflect new S" |
| lin.E = E_orig |
| print(" PASS test_e_update_changes_magnitude") |
|
|
|
|
| def test_state_dict_no_float_weights(): |
| """Model state_dict has no float weight tensors.""" |
| from arbitor.main import ARBModel |
| model = ARBModel( |
| enable_image=False, enable_audio=False, enable_vq=False, |
| enable_graph=False, enable_memory_modules=False, enable_moe=False, |
| ) |
| sd = model.state_dict() |
| for key, tensor in sd.items(): |
| if "T_packed" in key: |
| assert tensor.dtype == torch.uint8, f"{key} should be uint8, got {tensor.dtype}" |
| elif "T_accum" in key or "E_accum" in key or "E" in key or "group_lr" in key: |
| assert tensor.dtype == torch.int8, f"{key} should be int8, got {tensor.dtype}" |
| elif "corr_accum" in key or "step_counter" in key: |
| assert tensor.dtype == torch.int64, f"{key} should be int64, got {tensor.dtype}" |
| elif "bias" in key: |
| assert tensor.dtype == torch.int32, f"{key} should be int32, got {tensor.dtype}" |
| print(" PASS test_state_dict_no_float_weights") |
|
|
