| import pytest |
| import torch |
|
|
| import kernels |
|
|
| dt = kernels.get_kernel("phanerozoic/det-train", version=1, trust_remote_code=True) |
|
|
| requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") |
|
|
|
|
| @requires_cuda |
| @pytest.mark.kernels_ci |
| @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) |
| def test_sharding_invariance_bitwise(dtype): |
| """The core contract: any K-sharding on canonical block boundaries |
| composes to the bitwise-identical result.""" |
| torch.manual_seed(0) |
| M, N, K = 64, 96, 8192 |
| x = (torch.randn(M, K, device="cuda") * |
| torch.exp2(torch.randint(-20, 21, (M, K), device="cuda").float())).to(dtype) |
| w = (torch.randn(N, K, device="cuda") * |
| torch.exp2(torch.randint(-20, 21, (N, K), device="cuda").float())).to(dtype) |
| full = dt.det_gemm(x, w) |
| for shards in (2, 4, 8): |
| bs = K // shards |
| assert bs % 1024 == 0 |
| states = [dt.det_gemm_partial(x[:, i*bs:(i+1)*bs].contiguous(), |
| w[:, i*bs:(i+1)*bs].contiguous()) |
| for i in range(shards)] |
| y = dt.det_finalize(*dt.det_combine(states)) |
| assert torch.equal(y, full), f"{shards}-way shard differs" |
| |
| st = [dt.det_gemm_partial(x[:, i*2048:(i+1)*2048].contiguous(), |
| w[:, i*2048:(i+1)*2048].contiguous()) for i in range(4)] |
| left = dt.det_combine([dt.det_combine(st[:2]), dt.det_combine(st[2:])]) |
| seq = dt.det_combine(st) |
| assert torch.equal(left[0], seq[0]) and torch.equal(left[1], seq[1]) |
|
|
|
|
| @requires_cuda |
| @pytest.mark.kernels_ci |
| def test_repeat_bitwise_and_accuracy(): |
| torch.manual_seed(1) |
| x = torch.randn(128, 4096, device="cuda") |
| w = torch.randn(256, 4096, device="cuda") |
| y1, y2 = dt.det_gemm(x, w), dt.det_gemm(x, w) |
| assert torch.equal(y1, y2) |
| ref = (x.double() @ w.double().t()) |
| rel = ((y1.double() - ref).abs().max() / ref.abs().max()).item() |
| assert rel < 2 ** -30, rel |
|
|
|
|
| @requires_cuda |
| @pytest.mark.kernels_ci |
| def test_training_bitwise_repeatable_and_shard_equal(): |
| """Two full training runs are bitwise identical; a run whose forward is |
| computed via 4-way K-sharded composition equals the unsharded run.""" |
| def train(sharded: bool): |
| torch.manual_seed(7) |
| l1 = dt.DetLinear(2048, 256, device="cuda") |
| opt = torch.optim.SGD(l1.parameters(), lr=1e-3) |
| x = torch.randn(32, 2048, device="cuda") |
| t = torch.randn(32, 256, device="cuda") |
| losses = [] |
| for _ in range(50): |
| if sharded: |
| st = [dt.det_gemm_partial(x[:, i*1024:(i+1)*1024].contiguous(), |
| l1.weight[:, i*1024:(i+1)*1024].contiguous()) |
| for i in range(2)] |
| y = dt.det_finalize(*dt.det_combine(st)) |
| |
| y2 = _DetProxy.apply(x, l1.weight, y) |
| else: |
| y2 = l1(x) |
| loss = (y2 - t).square().mean() |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| opt.step() |
| losses.append(loss.detach().clone()) |
| return torch.stack(losses) |
|
|
| class _DetProxy(torch.autograd.Function): |
| @staticmethod |
| def forward(ctx, x, W, y): |
| ctx.save_for_backward(x, W) |
| return y |
| @staticmethod |
| def backward(ctx, dY): |
| x, W = ctx.saved_tensors |
| dX = dt.det_gemm(dY.contiguous(), W.t().contiguous()) |
| dW = dt.det_gemm(dY.t().contiguous(), x.t().contiguous()) |
| return dX, dW, None |
| globals()['_DetProxy'] = _DetProxy |
|
|
| a = train(False) |
| b = train(False) |
| assert torch.equal(a, b), "repeat runs must be bitwise identical" |
| c = train(True) |
| assert torch.equal(a, c), "sharded-forward run must equal unsharded bitwise" |
|
|
|
|
| @requires_cuda |
| @pytest.mark.kernels_ci |
| def test_det_sum_matches_permutation(): |
| torch.manual_seed(3) |
| x = torch.randn(64, 4096, device="cuda") * torch.exp2( |
| torch.randint(-15, 16, (64, 4096), device="cuda").float()) |
| a = dt.det_sum(x, dim=-1) |
| idx = torch.randperm(4096, device="cuda") |
| b = dt.det_sum(x[:, idx], dim=-1) |
| assert torch.equal(a, b), "det_sum must be permutation-invariant" |
|
|
