import math import torch import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) from arbitor.kernel import ternary_scale as tscale from arbitor.kernel.ternary_scale import TernaryScaleTensor, TScaleType, TILE_SIZE, GROUP_SIZES from arbitor.optim.sign_sgd import SignSGD from arbitor.components import StickyZoneSTE from arbitor.config import VOCAB, CTX, SPECIAL_VOCAB from arbitor.main import ARBModel from arbitor.components import LossComponents from arbitor.kernel.ternary_scale import TernaryRMSNorm from arbitor.sequencers import ByteEmbedding def _cuda_available(min_gib=10): """Check CUDA is available with enough GPU memory (min_gib GiB).""" if not torch.cuda.is_available(): return False free, total = torch.cuda.mem_get_info() if total < min_gib * 1e9: return False return True # ─── TernaryScaleTensor Tests ─── def test_tscale_shape(): lin = TernaryScaleTensor(32, 16) x = torch.randn(2, 10, 32) out = lin(x) assert out.shape == (2, 10, 16), f"Shape: {out.shape}" print(" PASS test_tscale_shape") def test_tscale_ternary_output(): lin = TernaryScaleTensor(32, 16, threshold=0.05) T = lin._get_T() unique = set(T.detach().flatten().tolist()) assert unique.issubset({-1, 0, 1}), f"Non-ternary values in T: {unique}" print(" PASS test_tscale_ternary_output") def test_tscale_T64_per_element_s(): lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T64) dq = lin.dequantize() assert dq.shape == (16, 32), f"Dequantize shape: {dq.shape}" print(" PASS test_tscale_T64_per_element_s") def test_tscale_T32_group_s(): lin = TernaryScaleTensor(96, 16, tscale_type=TScaleType.T32) dq = lin.dequantize() gpr = lin.E.shape[0] // lin.out_dim assert gpr > 0, f"Groups per row: {gpr}" assert dq.shape == (16, 96), f"Dequantize shape: {dq.shape}" print(" PASS test_tscale_T32_group_s") def test_tscale_to_switching(): lin = TernaryScaleTensor(96, 16, tscale_type=TScaleType.T64) dq_before = lin.dequantize() assert lin.tscale_type == TScaleType.T64 lin.tscale_to(TScaleType.T32) assert lin.tscale_type == TScaleType.T32 dq_after = lin.dequantize() assert dq_before.shape == dq_after.shape lin.tscale_to(TScaleType.T4) assert lin.tscale_type == TScaleType.T4 dq_t4 = lin.dequantize() assert dq_t4.shape == dq_before.shape print(" PASS test_tscale_to_switching") def test_tscale_cast_alias(): lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T64) result = lin.tscale_cast(TScaleType.T8) assert result is lin, "tscale_cast should return self" assert lin.tscale_type == TScaleType.T8 print(" PASS test_tscale_cast_alias") def test_tscale_gradient_flow(): lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) x = torch.randn(2, 10, 32) x.requires_grad_(True) out = lin(x) out.sum().backward() assert x.grad is not None, "No gradient on input" print(" PASS test_tscale_gradient_flow") def test_tscale_all_types_forward(): for tscale_type in TScaleType: lin = TernaryScaleTensor(96, 16, tscale_type=tscale_type) x = torch.randn(2, 4, 96) out = lin(x) assert out.shape == (2, 4, 16), f"{tscale_type.name}: shape {out.shape}" assert torch.isfinite(out).all(), f"{tscale_type.name}: non-finite output" print(" PASS test_tscale_all_types_forward") def test_tscale_dequantize(): lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) w_eff = lin.dequantize() assert w_eff.shape == (16, 32), f"Shape: {w_eff.shape}" assert torch.isfinite(w_eff).all() print(" PASS test_tscale_dequantize") def test_tscale_effective_bpw(): lin64 = TernaryScaleTensor(384, 384, tscale_type=TScaleType.T64) lin4 = TernaryScaleTensor(384, 384, tscale_type=TScaleType.T4) assert lin4.effective_bpw > lin64.effective_bpw, "T4 (gs=4) should have higher BPW than T64 (gs=64)" print(f" T64 BPW: {lin64.effective_bpw:.2f}, T4 BPW: {lin4.effective_bpw:.2f}") print(" PASS test_tscale_effective_bpw") def test_tscale_model_integration(): if not _cuda_available(): print(" SKIP test_tscale_model_integration (need CUDA + >10GB GPU)") return for tscale_type in [TScaleType.T64, TScaleType.T32, TScaleType.T8]: model = ARBModel(tscale_type=tscale_type).to("cuda") x = torch.randint(0, VOCAB, (2, 10), device="cuda") logits, losses, _, _ = model(x, targets=x[:, 3:]) assert losses is not None losses.total.backward() print(" PASS test_tscale_model_integration") def test_tscale_runtime_switch(): if not _cuda_available(): print(" SKIP test_tscale_runtime_switch (need CUDA + >10GB GPU)") return model = ARBModel(tscale_type=TScaleType.T64).to("cuda") x = torch.randint(0, VOCAB, (1, 10), device="cuda") logits64, _, _, _ = model(x) for module in model.modules(): if isinstance(module, TernaryScaleTensor): module.tscale_to(TScaleType.T4) logits4, _, _, _ = model(x) assert torch.isfinite(logits4).all(), "Non-finite after tscale.to(T4)" assert logits4.shape == logits64.shape, "Shape mismatch after tscale switch" print(" PASS test_tscale_runtime_switch") # ─── SignSGD Tests ─── def test_sign_sgd_step(): model = torch.nn.Linear(10, 5) optimizer = SignSGD(model.parameters(), lr=0.01) x = torch.randn(2, 10) loss = model(x).sum() loss.backward() w_before = model.weight.clone() optimizer.step() assert not torch.equal(model.weight, w_before), "Weights did not change" print(" PASS test_sign_sgd_step") def test_sign_sgd_no_momentum(): model = torch.nn.Linear(10, 5) optimizer = SignSGD(model.parameters(), lr=0.01) assert len(optimizer.state) == 0, "SignSGD should have no state (no momentum)" print(" PASS test_sign_sgd_no_momentum") def test_sign_sgd_memory(): model = torch.nn.Linear(100, 100) optimizer = SignSGD(model.parameters(), lr=0.01) mem = optimizer.get_memory_mb() assert mem > 0, "Memory should be positive" print(f" SignSGD memory: {mem:.2f} MB") print(" PASS test_sign_sgd_memory") def test_sign_sgd_with_tscale_model(): if not _cuda_available(): print(" SKIP test_sign_sgd_with_tscale_model (need CUDA + >10GB GPU)") return model = ARBModel(tscale_type=TScaleType.T32).to("cuda") x = torch.randint(0, VOCAB, (2, 10), device="cuda") logits, losses, _, _ = model(x, targets=x[:, 3:]) losses.total.backward() model._ternary_update_memory() print(" PASS test_sign_sgd_with_tscale_model") def test_sign_sgd_weight_decay(): model = torch.nn.Linear(10, 5) optimizer = SignSGD(model.parameters(), lr=0.01, weight_decay=0.01) x = torch.randn(2, 10) loss = model(x).sum() loss.backward() w_before = model.weight.clone() optimizer.step() w_diff = (model.weight - w_before).abs().sum().item() assert w_diff > 0, "Weights should change with weight_decay" print(" PASS test_sign_sgd_weight_decay") # ─── TileLang PyTorch Reference Tests ─── def test_dequant_gemm_pytorch_ref(): import importlib.util kernel_path = os.path.join(os.path.dirname(__file__), "..", "tilelang", "kernels", "dequant_gemm.py") if not os.path.exists(kernel_path): print(" SKIP test_dequant_gemm_pytorch_ref (tilelang reference file missing)") return spec = importlib.util.spec_from_file_location("dequant_gemm", kernel_path) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) dequant_gemm_pytorch_ref = mod.dequant_gemm_pytorch_ref M, N, K, group_size = 4, 8, 96, 12 signs = torch.randint(-1, 2, (N, K), dtype=torch.int8) exponents = torch.randint(-3, 4, (N, K // group_size), dtype=torch.int8) x = torch.randn(M, K, dtype=torch.float16) output = dequant_gemm_pytorch_ref(signs, exponents, x, group_size) assert output.shape == (M, N), f"Shape: {output.shape}" assert torch.isfinite(output).all(), "Non-finite output" print(" PASS test_dequant_gemm_pytorch_ref") def test_dequant_gemm_matches_manual(): import importlib.util import torch.nn.functional as F kernel_path = os.path.join(os.path.dirname(__file__), "..", "tilelang", "kernels", "dequant_gemm.py") if not os.path.exists(kernel_path): print(" SKIP test_dequant_gemm_matches_manual (tilelang reference file missing)") return spec = importlib.util.spec_from_file_location("dequant_gemm", kernel_path) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) dequant_gemm_pytorch_ref = mod.dequant_gemm_pytorch_ref M, N, K, group_size = 2, 4, 48, 12 signs = torch.randint(-1, 2, (N, K), dtype=torch.int8) exponents = torch.randint(-3, 4, (N, K // group_size), dtype=torch.int8) x = torch.randn(M, K, dtype=torch.float16) result = dequant_gemm_pytorch_ref(signs, exponents, x, group_size) exp_expanded = exponents.repeat_interleave(group_size, dim=1) pos_mask = exp_expanded >= 0 two_pow = torch.where(pos_mask, (1 << exp_expanded.to(torch.int32)).to(torch.float16), (1 >> (-exp_expanded.to(torch.int32))).to(torch.float16)) w = signs.to(torch.float16) * two_pow expected = x @ w.t() assert torch.allclose(result, expected, atol=1e-3), "PyTorch ref mismatch" print(" PASS test_dequant_gemm_matches_manual") # ─── Integration: SignSGD + TernaryScaleTensor training step ─── def test_full_training_step(): if not _cuda_available(): print(" SKIP test_full_training_step (need CUDA + >10GB GPU)") return model = ARBModel(tscale_type=TScaleType.T32).to("cuda") x = torch.randint(0, VOCAB, (2, 10), device="cuda") logits, losses, _, _ = model(x, targets=x[:, 3:]) losses.total.backward() model._ternary_update_memory() logits2, losses2, _, _ = model(x, targets=x[:, 3:]) assert torch.isfinite(losses2.total), "Non-finite loss after step" print(" PASS test_full_training_step") def test_multiple_steps_converge(): if not _cuda_available(): print(" SKIP test_multiple_steps_converge (need CUDA + >10GB GPU)") return model = ARBModel(tscale_type=TScaleType.T32).to("cuda") x = torch.randint(0, VOCAB, (4, 10), device="cuda") losses = [] for step in range(50): logits, losses_out, _, _ = model(x, targets=x[:, 3:]) loss_val = losses_out.total loss_val.backward() model._ternary_update_memory(accum_threshold=3) losses.append(loss_val.item()) assert torch.isfinite(torch.tensor(losses)).all(), "Non-finite loss during training" print(f" Loss range: {min(losses):.4f} – {max(losses):.4f} over 50 steps") print(" PASS test_multiple_steps_converge") def test_cuda_triton_correctness_linear(): if not torch.cuda.is_available() or not tscale._HAS_TRITON: print(" SKIP test_cuda_triton_correctness_linear (CUDA/Triton unavailable)") return from arbitor.kernel.ternary_scale import TernaryRMSNorm, _triton_ternary_embed from arbitor.main import ByteEmbedding ATOL = 2e-3 for tt in [TScaleType.T4, TScaleType.T6, TScaleType.T8, TScaleType.T16, TScaleType.T32, TScaleType.T64]: lin_cpu = TernaryScaleTensor(32, 16, tscale_type=tt) x = torch.randn(4, 4, 32, requires_grad=True) lin_gpu = TernaryScaleTensor(32, 16, tscale_type=tt).cuda() lin_gpu.load_state_dict(lin_cpu.state_dict()) cpu_out = lin_cpu(x) grad_out = torch.randn_like(cpu_out) cpu_out.backward(grad_out) cpu_grad_x = x.grad.clone() x_gpu = x.detach().clone().cuda().requires_grad_(True) gpu_out = lin_gpu(x_gpu) gpu_out.backward(grad_out.cuda()) gpu_grad_x = x_gpu.grad.clone() fwd_diff = (cpu_out - gpu_out.cpu()).abs().max().item() bwd_diff = (cpu_grad_x - gpu_grad_x.cpu()).abs().max().item() assert fwd_diff < ATOL, f"{tt.name} fwd_diff={fwd_diff}" assert bwd_diff < ATOL, f"{tt.name} bwd_diff={bwd_diff}" print(" PASS test_cuda_triton_correctness_linear") def test_cuda_triton_correctness_rmsnorm(): if not torch.cuda.is_available() or not tscale._HAS_TRITON: print(" SKIP test_cuda_triton_correctness_rmsnorm (CUDA/Triton unavailable)") return from arbitor.kernel.ternary_scale import TernaryRMSNorm for tt in [TScaleType.T4, TScaleType.T6, TScaleType.T8, TScaleType.T16, TScaleType.T32, TScaleType.T64]: norm_cpu = TernaryRMSNorm(256, tscale_type=tt) x = torch.randn(2, 4, 256, requires_grad=True) cpu_out = norm_cpu(x) cpu_out.sum().backward() cpu_grad_x = x.grad.clone() norm_gpu = TernaryRMSNorm(256, tscale_type=tt).cuda() norm_gpu.load_state_dict(norm_cpu.state_dict()) x_gpu = x.detach().clone().cuda().requires_grad_(True) gpu_out = norm_gpu(x_gpu) gpu_out.sum().backward() gpu_grad_x = x_gpu.grad.clone() fwd_diff = (cpu_out - gpu_out.cpu()).abs().max().item() bwd_diff = (cpu_grad_x - gpu_grad_x.cpu()).abs().max().item() assert fwd_diff < 1e-5, f"{tt.name} rmsnorm fwd_diff={fwd_diff}" assert bwd_diff < 1e-5, f"{tt.name} rmsnorm bwd_diff={bwd_diff}" print(" PASS test_cuda_triton_correctness_rmsnorm") def test_cuda_triton_correctness_embedding(): if not torch.cuda.is_available() or not tscale._HAS_TRITON: print(" SKIP test_cuda_triton_correctness_embedding (CUDA/Triton unavailable)") return from arbitor.main import ByteEmbedding for tt in [TScaleType.T4, TScaleType.T6, TScaleType.T8, TScaleType.T16, TScaleType.T32, TScaleType.T64]: emb_cpu = ByteEmbedding(tscale_type=tt) x = torch.tensor([0, 1, 2, 5, 10]) cpu_out = emb_cpu(x) cpu_out.sum().backward() emb_gpu = ByteEmbedding(tscale_type=tt).cuda() emb_gpu.load_state_dict(emb_cpu.state_dict()) x_gpu = x.cuda() gpu_out = emb_gpu(x_gpu) gpu_out.sum().backward() fwd_diff = (cpu_out - gpu_out.cpu()).abs().max().item() assert fwd_diff < 1e-5, f"{tt.name} embed fwd_diff={fwd_diff}" if hasattr(emb_cpu, '_hook_grad_T_sign') and hasattr(emb_gpu, '_hook_grad_T_sign'): gs_match = (emb_gpu._hook_grad_T_sign.cpu() == emb_cpu._hook_grad_T_sign).float().mean().item() assert gs_match > 0.99, f"{tt.name} embed grad_sign match={gs_match}" print(" PASS test_cuda_triton_correctness_embedding") def test_cuda_triton_correctness_update_E(): if not torch.cuda.is_available() or not tscale._HAS_TRITON: print(" SKIP test_cuda_triton_correctness_update_E (CUDA/Triton unavailable)") return for tt in [TScaleType.T4, TScaleType.T6, TScaleType.T8, TScaleType.T16, TScaleType.T32, TScaleType.T64]: lin_cpu = TernaryScaleTensor(32, 16, tscale_type=tt) lin_gpu = TernaryScaleTensor(32, 16, tscale_type=tt).cuda() lin_gpu.load_state_dict(lin_cpu.state_dict()) x_cpu = torch.randn(4, 4, 32, requires_grad=True) x_gpu = x_cpu.detach().clone().cuda().requires_grad_(True) cpu_out = lin_cpu(x_cpu) cpu_out.sum().backward() E_cpu = lin_cpu.E.clone() corr_cpu = lin_cpu.corr_accum.clone() step_cpu = lin_cpu.step_counter.clone() gpu_out = lin_gpu(x_gpu) gpu_out.sum().backward() E_gpu = lin_gpu.E.clone() corr_gpu = lin_gpu.corr_accum.clone() step_gpu = lin_gpu.step_counter.clone() # E is fixed after init; BigInt corr_accum carries the continuous scale adjustment. E_diff = (E_cpu.float() - E_gpu.cpu().float()).abs().max().item() assert E_diff < 0.01, f"{tt.name} CPU-GPU E update mismatch: {E_diff}" corr_diff = (corr_cpu - corr_gpu.cpu()).abs().max().item() assert corr_diff == 0, f"{tt.name} CPU-GPU corr_accum update mismatch: {corr_diff}" assert int(step_cpu.item()) == int(step_gpu.cpu().item()) == 1, \ f"{tt.name} CPU-GPU step_counter mismatch: cpu={step_cpu.item()} gpu={step_gpu.cpu().item()}" print(" PASS test_cuda_triton_correctness_update_E") def test_cuda_triton_correctness_ternary_step(): if not torch.cuda.is_available() or not tscale._HAS_TRITON: print(" SKIP test_cuda_triton_correctness_ternary_step (CUDA/Triton unavailable)") return for tt in [TScaleType.T4, TScaleType.T6, TScaleType.T8, TScaleType.T16, TScaleType.T32, TScaleType.T64]: lin_cpu = TernaryScaleTensor(32, 16, tscale_type=tt) lin_gpu = TernaryScaleTensor(32, 16, tscale_type=tt).cuda() lin_gpu.load_state_dict(lin_cpu.state_dict()) x_cpu = torch.randn(4, 4, 32, requires_grad=True) x_gpu = x_cpu.detach().clone().cuda().requires_grad_(True) cpu_out = lin_cpu(x_cpu) cpu_out.sum().backward() lin_cpu.ternary_step(accum_threshold=3) T_cpu = lin_cpu._get_T().clone() corr_cpu = lin_cpu.corr_accum.clone() gpu_out = lin_gpu(x_gpu) gpu_out.sum().backward() lin_gpu.ternary_step(accum_threshold=3) T_gpu = lin_gpu._get_T().clone() corr_gpu = lin_gpu.corr_accum.clone() T_match = (T_cpu == T_gpu.cpu()).float().mean().item() corr_match = (corr_cpu == corr_gpu.cpu()).float().mean().item() assert T_match == 1.0, f"{tt.name} T_match={T_match}" assert corr_match == 1.0, f"{tt.name} corr_match={corr_match}" print(" PASS test_cuda_triton_correctness_ternary_step") def test_cuda_triton_tscale_path(): if not torch.cuda.is_available() or not tscale._HAS_TRITON: print(" SKIP test_cuda_triton_tscale_path (CUDA/Triton unavailable)") return lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32).cuda() x = torch.randn(2, 4, 32, device="cuda", requires_grad=True) out = lin(x) assert out.is_cuda, "Triton path should produce CUDA output" assert out.shape == (2, 4, 16), f"Shape: {out.shape}" grad_out = torch.randn_like(out) out.backward(grad_out) assert x.grad is not None and x.grad.is_cuda, "CUDA grad_x missing" assert lin.corr_accum.abs().sum().item() > 0, \ "Triton path should stream updates into int64 corr_accum" assert int(lin.step_counter.item()) == 1, "Triton path should advance the BigInt step counter" assert not hasattr(lin, "_hook_grad_T_sign"), \ "Triton path should not retain full weight-shaped grad-sign hooks" assert not hasattr(lin, "_hook_grad_2d") and not hasattr(lin, "_hook_x_2d"), \ "Triton path should not retain fp32 grad/x views" torch.cuda.synchronize() assert not hasattr(lin, "_hook_grad_T_sign"), \ "No retained grad-sign hook should remain after streaming backward" assert lin.T_packed.is_cuda and lin.E.is_cuda, "Ternary buffers moved off CUDA after update" lin_force = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32).cuda() lin_force._hook_grad_2d = torch.ones(2, 16, device="cuda") lin_force._hook_x_2d = torch.ones(2, 32, device="cuda") lin_force.update_E() forced_T = lin_force._get_T() assert forced_T.is_cuda, "Unpacked CUDA ternary state should stay on CUDA" assert lin_force.corr_accum.abs().sum().item() > 0, "Forced CUDA hook should update BigInt corr_accum" assert int(lin_force.step_counter.item()) == 1, "Forced CUDA hook should advance the BigInt step counter" print(" PASS test_cuda_triton_tscale_path") def test_small_ternary_training_loss_finite(): if not torch.cuda.is_available(): print(" SKIP test_small_ternary_training_loss_finite (CUDA unavailable)") return model = ARBModel( enable_image=False, enable_audio=False, enable_vq=False, enable_graph=False, enable_memory_modules=False, enable_moe=False, tscale_type=TScaleType.T32, ).cuda() x = torch.randint(0, VOCAB, (1, 4), device="cuda") _, losses, _, _ = model(x, targets=x[:, 3:]) assert torch.isfinite(losses.total), "Small ternary training loss is non-finite" model._ternary_update_memory(accum_threshold=3, update_scales=True, loss_components=losses) leftovers = [ name for name, module in model.named_modules() if any(hasattr(module, hook) for hook in ("_hook_grad_T_sign", "_hook_grad_2d", "_hook_x_2d")) ] assert not leftovers, f"Ternary update left stale hooks: {leftovers[:5]}" print(" PASS test_small_ternary_training_loss_finite") def test_ternary_update_rejects_nonfinite_loss(): import warnings model = ARBModel( enable_image=False, enable_audio=False, enable_vq=False, enable_graph=False, enable_memory_modules=False, enable_moe=False, tscale_type=TScaleType.T32, ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") lc = LossComponents(lm=torch.tensor(float("nan"))) model._ternary_update_memory(loss_components=lc) assert len(w) > 0, "Expected a warning for non-finite loss" assert "Non-finite loss" in str(w[0].message), f"Unexpected warning: {w[0].message}" print(" PASS test_ternary_update_rejects_nonfinite_loss") # ─── Phase 12: E Gradient Field + Statistical Metrics Tests ─── def test_e_rms_weighted_delta(): lin = TernaryScaleTensor(32, 16, tscale_type=TScaleType.T32) grad = torch.randn(4, 16) x = torch.randn(4, 32) raw_grad = grad.T @ x # compute RMS per group gpr = (32 + lin.group_size - 1) // lin.group_size rms_per_group = [] for g in range(gpr): start = g * lin.group_size end = min(start + lin.group_size, 32) group = raw_grad[:, start:end] rms = group.pow(2).mean().sqrt().item() rms_per_group.append(rms) rms = rms_per_group[0] score = (raw_grad * lin._get_T().float()).sum().item() delta = - (1 if score > 0 else -1 if score < 0 else 0) * max(1, min(3, round(math.log2(1 + rms)))) assert 1 <= abs(delta) <= 4, f"delta magnitude {abs(delta)} out of range" print(" PASS test_e_rms_weighted_delta") def test_e_rms_vs_sign_only(): # RMS-weighted delta differs for different gradient magnitudes even when sign is same raw_low = torch.ones(16, 32) * 0.1 raw_high = torch.ones(16, 32) * 10.0 T = torch.ones(16, 32) rms_low = raw_low.pow(2).mean().sqrt() rms_high = raw_high.pow(2).mean().sqrt() delta_low = max(1, min(3, round(math.log2(1 + rms_low.item())))) delta_high = max(1, min(3, round(math.log2(1 + rms_high.item())))) assert delta_low != delta_high, "RMS delta should differ for different magnitudes" assert delta_low < delta_high, "Higher RMS should give larger delta" print(" PASS test_e_rms_vs_sign_only") def test_e_zscore_normalization(): comp_a_rms = torch.tensor([10.0, 12.0, 8.0, 11.0]) comp_b_rms = torch.tensor([1.0, 1.2, 0.8, 1.1]) z_a = (comp_a_rms - comp_a_rms.mean()) / (comp_a_rms.std() + 1e-8) z_b = (comp_b_rms - comp_b_rms.mean()) / (comp_b_rms.std() + 1e-8) assert abs(z_a.mean().item()) < 1e-6, f"z_a mean not ~0: {z_a.mean().item()}" assert abs(z_b.mean().item()) < 1e-6, f"z_b mean not ~0: {z_b.mean().item()}" assert abs(z_a.std().item() - 1.0) < 0.1, f"z_a std not ~1: {z_a.std().item()}" assert abs(z_b.std().item() - 1.0) < 0.1, f"z_b std not ~1: {z_b.std().item()}" print(" PASS test_e_zscore_normalization") def test_e_zscore_zero_std(): rms_flat = torch.ones(8) * 5.0 z = torch.where(rms_flat.std() > 1e-8, (rms_flat - rms_flat.mean()) / (rms_flat.std()), torch.zeros_like(rms_flat)) assert torch.isfinite(z).all(), "z-scores should be finite when std=0" assert (z == 0).all(), "z-scores should be zero when std=0" print(" PASS test_e_zscore_zero_std") def test_group_lr_registration(): tst = TernaryScaleTensor(32, 16) assert hasattr(tst, "corr_accum") assert tst.corr_accum.dtype == torch.int64 assert tst.corr_accum.shape == tst.E.shape assert int(tst.step_counter.item()) == 0 be = ByteEmbedding() assert hasattr(be, "corr_accum") assert be.corr_accum.dtype == torch.int64 assert be.corr_accum.shape[0] > 0 assert int(be.step_counter.item()) == 0 rms = TernaryRMSNorm(256) assert hasattr(rms, "E") print(" PASS test_group_lr_registration") def test_group_lr_effect(): delta = torch.tensor(4, dtype=torch.int8) group_lr_high = torch.tensor(8, dtype=torch.int8) group_lr_low = torch.tensor(1, dtype=torch.int8) eff_high = delta.to(torch.int16) * group_lr_high.to(torch.int16) // 8 eff_low = delta.to(torch.int16) * group_lr_low.to(torch.int16) // 8 assert eff_high.item() == 4, f"high LR should give full delta, got {eff_high.item()}" assert eff_low.item() == 0, f"low LR should give 0 delta, got {eff_low.item()}" print(" PASS test_group_lr_effect") def test_group_lr_dynamic_update(): group_lr = torch.ones(4, dtype=torch.int8) rms_prev = torch.tensor([1.0, 5.0, 3.0, 2.0]) rms_curr = torch.tensor([2.0, 3.0, 3.0, 1.0]) rms_growth = rms_curr - rms_prev updated = torch.clamp(group_lr.to(torch.int16) + (rms_growth > 0).to(torch.int16) - (rms_growth < 0).to(torch.int16), 1, 8).to(torch.int8) assert updated[0].item() == 2, f"RMS increased -> LR should increase, got {updated[0].item()}" assert updated[1].item() == 1, f"RMS decreased -> LR should decrease, got {updated[1].item()}" assert updated[2].item() == 1, f"RMS unchanged -> LR unchanged, got {updated[2].item()}" # clamp boundaries too_high = torch.clamp(torch.tensor([100], dtype=torch.int16), 1, 8) too_low = torch.clamp(torch.tensor([-100], dtype=torch.int16), 1, 8) assert too_high.item() == 8, f"clamp max, got {too_high.item()}" assert too_low.item() == 1, f"clamp min, got {too_low.item()}" print(" PASS test_group_lr_dynamic_update") def test_e_stats_cpu_fallback(): N, K, group_size = 16, 32, 12 grad = torch.randn(4, N) x = torch.randn(4, K) raw_grad = grad.T @ x gpr = (K + group_size - 1) // group_size rms_vals = [] for g in range(gpr): start = g * group_size end = min(start + group_size, K) group = raw_grad[:, start:end] rms = group.pow(2).mean().sqrt() rms_vals.append(rms.item()) assert all(torch.isfinite(torch.tensor(rms_vals))), "finite check" assert all(1 <= max(1, min(3, round(math.log2(1 + r)))) <= 3 for r in rms_vals), "clamp range" print(" PASS test_e_stats_cpu_fallback") def test_e_per_component_routing(): if not _cuda_available(): print(" SKIP test_e_per_component_routing (CUDA)") return model = ARBModel(enable_image=False, enable_audio=False, enable_vq=False, enable_graph=False, enable_memory_modules=False, enable_moe=False).cuda() x = torch.randint(0, VOCAB, (1, 4), device="cuda") for step in range(3): _, lc, _, _ = model(x, targets=x[:, 3:]) model._ternary_update_memory(loss_components=lc) assert True # no crash = pass print(" PASS test_e_per_component_routing") def test_ensure_group_lr_backward_compat(): tst = TernaryScaleTensor(32, 16) assert hasattr(tst, "corr_accum") assert hasattr(tst, "step_counter") be = ByteEmbedding() assert hasattr(be, "corr_accum") assert hasattr(be, "step_counter") rms = TernaryRMSNorm(256) assert hasattr(rms, "E") print(" PASS test_ensure_group_lr_backward_compat") # ─── Main ─── if __name__ == "__main__": tests = [ test_tscale_shape, test_tscale_ternary_output, test_tscale_T64_per_element_s, test_tscale_T32_group_s, test_tscale_to_switching, test_tscale_cast_alias, test_tscale_gradient_flow, test_tscale_all_types_forward, test_tscale_dequantize, test_tscale_effective_bpw, test_tscale_model_integration, test_tscale_runtime_switch, test_sign_sgd_step, test_sign_sgd_no_momentum, test_sign_sgd_memory, test_sign_sgd_with_tscale_model, test_sign_sgd_weight_decay, test_dequant_gemm_pytorch_ref, test_dequant_gemm_matches_manual, test_cuda_triton_correctness_linear, test_cuda_triton_correctness_rmsnorm, test_cuda_triton_correctness_embedding, test_cuda_triton_correctness_update_E, test_cuda_triton_correctness_ternary_step, test_cuda_triton_tscale_path, test_small_ternary_training_loss_finite, test_ternary_update_rejects_nonfinite_loss, test_full_training_step, test_multiple_steps_converge, test_e_rms_weighted_delta, test_e_rms_vs_sign_only, test_e_zscore_normalization, test_e_zscore_zero_std, test_group_lr_registration, test_group_lr_effect, test_group_lr_dynamic_update, test_e_stats_cpu_fallback, test_e_per_component_routing, test_ensure_group_lr_backward_compat, ] print("Running TernaryScale + SignSGD + TileLang Phase 2 tests...\n") passed = 0 failed = 0 for test in tests: try: test() passed += 1 except Exception as e: print(f" FAIL {test.__name__}: {e}") import traceback traceback.print_exc() failed += 1 print(f"\n{passed} passed, {failed} failed out of {len(tests)} tests")