File size: 7,307 Bytes
411abf3 63ed286 411abf3 63ed286 411abf3 63ed286 411abf3 63ed286 411abf3 63ed286 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """One resolver for CPU, CUDA, ROCm and MPS (a named backend that is missing falls back with one warning), batch sizes from
free memory and the graph's size (reproducing the measured calibration), and the network's arithmetic agreeing with
the CPU on whichever accelerator this machine has."""
import logging
import types
import numpy as np
import pytest
import torch
from fly_trader.brain import device, plastic
from fly_trader.train.fly_selector import SCALE, FlyNet
from tests.test_fly_plastic import _graph
ACCELERATORS = [d for d in ("cuda", "mps") if device._available(d)] # cuda: an AMD GPU too, on a ROCm build
def test_resolve_honours_the_name_and_falls_back_with_one_warning(caplog, monkeypatch):
device.reset()
assert device.resolve("cpu").type == "cpu" and device.resolve("CPU").type == "cpu"
assert device.resolve("auto").type == (ACCELERATORS[0] if ACCELERATORS else "cpu")
missing = next((d for d in ("cuda", "mps") if d not in ACCELERATORS), None)
if missing:
with caplog.at_level(logging.WARNING, logger="fly_trader.brain.device"):
assert device.resolve(missing).type == "cpu" and device.resolve(missing + ":0").type == "cpu"
assert device.resolve(missing).type == "cpu"
assert sum("not available" in r.message for r in caplog.records) == 2 # once per name: cached
with caplog.at_level(logging.WARNING, logger="fly_trader.brain.device"):
assert device.resolve("tpu").type == "cpu"
assert any("not cpu, cuda, rocm, mps or auto" in r.message for r in caplog.records)
monkeypatch.setattr(device.config, "DEVICE", "cpu"); device.reset()
assert device.resolve().type == "cpu" # the default reads config.DEVICE
device.reset()
def test_rocm_names_the_amd_gpu_a_rocm_build_drives_through_torch_cuda(caplog, monkeypatch):
monkeypatch.setattr(torch.version, "hip", "7.2.53211", raising=False) # a ROCm build with two AMD GPUs
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "device_count", lambda: 2)
device.reset()
assert device.resolve("rocm") == torch.device("cuda") and device.resolve("ROCm:1") == torch.device("cuda", 1)
assert device.resolve("hip:1") == torch.device("cuda", 1) and device.resolve("cuda:1") == torch.device("cuda", 1)
assert device.resolve("auto") == torch.device("cuda")
with caplog.at_level(logging.WARNING, logger="fly_trader.brain.device"):
assert device.resolve("rocm:5") == torch.device("cuda")
assert any("only 2 ROCm device(s) present; using rocm:0" in r.message for r in caplog.records)
gpu = types.SimpleNamespace(name="AMD Radeon 8060S Graphics", gcnArchName="gfx1151:xnack-", is_integrated=True,
total_memory=96 * 2**30)
monkeypatch.setattr(torch.cuda, "get_device_properties", lambda i: gpu)
assert device.describe(torch.device("cuda", 0)) == "rocm:0 (AMD Radeon 8060S Graphics, gfx1151, integrated, 96 GB)"
monkeypatch.setattr(torch.version, "hip", None); device.reset() # a CUDA build: rocm is missing
with caplog.at_level(logging.WARNING, logger="fly_trader.brain.device"):
assert device.resolve("rocm").type == "cpu" and device.resolve("cuda").type == "cuda"
assert any("not a ROCm build" in r.message and "rocm72" in r.message for r in caplog.records)
device.reset()
def test_rows_for_follows_memory_and_the_graph_and_reproduces_the_calibration():
edges, n = 1_039_659, 41_756 # the FAFB v783 sub-graph
free = int((19.5 * 2**30 + device.FIXED_BYTES) / device.BUDGET_SHARE) # what 2,048 rows measured at
assert 1800 <= device.rows_for(edges, n, ceiling=4096, free=free) <= 2300
small, big = (device.rows_for(edges, n, ceiling=4096, free=gb * 2**30) for gb in (8, 24))
assert device.MIN_ROWS <= small < big # an 8 GB card gets fewer rows than a 24 GB one
assert device.rows_for(edges, n, ceiling=4096, free=8 * 2**30, training=True) < small # training keeps every step: fewer still
assert device.rows_for(edges, n, ceiling=256, free=10**13) == 256 # the caller's ceiling
assert device.rows_for(edges, n, ceiling=4096, free=0) == device.MIN_ROWS # never zero
assert device.rows_for(100, 80, ceiling=500, free=2**30) == 500 # a toy graph fits whole
def test_the_memory_override_caps_the_budget(monkeypatch):
monkeypatch.setenv(device.MEMORY_GB_ENV, "2")
assert device.free_bytes(torch.device("cpu")) == 2 * 2**30
monkeypatch.setenv(device.MEMORY_GB_ENV, "lots")
assert device.free_bytes(torch.device("cpu")) > 2 * 2**30 # unparseable: measured instead
def test_describe_batch_rows_and_seeding():
assert device.describe(torch.device("cpu")).startswith("cpu (")
net = FlyNet(_graph(), obs_dim=5, device="cpu")
assert net.batch_rows(ceiling=2048) == 2048 and net.batch_rows(ceiling=64, training=True) == 64 # a toy graph: the ceilings bind
device.seed_all(3); a = torch.randn(3); device.seed_all(3)
assert torch.equal(a, torch.randn(3))
device.empty_cache(torch.device("cpu")) # a no-op that must not raise
@pytest.mark.parametrize("acc", ACCELERATORS or [pytest.param("none", marks=pytest.mark.skip(reason="no accelerator on this machine"))])
def test_the_network_and_its_plasticity_agree_with_the_cpu_on(acc):
torch.manual_seed(0); g = _graph()
cpu = FlyNet(g, obs_dim=5, device="cpu").eval(); other = FlyNet(g, obs_dim=5, device=acc).eval()
other.load_state_dict({k: v.to(acc) for k, v in cpu.state_dict().items()})
x = torch.randn(64, 5)
with torch.no_grad():
pa = cpu.forward_parts_all_h(x); pb = other.forward_parts_all_h(x.to(acc))
for a, b in zip(pa, pb):
assert torch.allclose(a, b.cpu(), atol=2e-4, rtol=1e-3) # scatter, k-WTA and the decoder agree
assert ((pa[2] > 0) == (pb[2].cpu() > 0)).float().mean() > 0.98 # the same Kenyon cells win
learn = np.ones((1, cpu.n_mbon), dtype=bool); B = 16; keys = [(0.0, f"m{i}", "ev") for i in range(B)]
out = []
for net, (y, u, k, _) in ((cpu, pa), (other, pb)):
bank = plastic.PlasticBank(net, [(1e-3, 1.0)], SCALE, learn=learn, read=learn)
bank.estimate_nu(y[:B, 0], u[:B], k[:B])
tags = plastic.Tags(keys, np.zeros(B), y[:B, 0], u[:B], k[:B], torch.ones(1, B, device=bank.dev))
bank.update(tags, torch.linspace(-0.05, 0.05, B, device=bank.dev), 60.0)
pred, _ = bank.predict(y[:B, 0], u[:B], k[:B], s=np.zeros(B, dtype=np.int64))
out.append((bank.D.detach().cpu(), pred.detach().cpu(), float(bank.drift()[0])))
assert torch.allclose(out[0][0], out[1][0], atol=1e-7, rtol=1e-3) # the learned change
assert torch.allclose(out[0][1], out[1][1], atol=2e-4, rtol=1e-3) # and the prediction it feeds
assert out[0][2] == pytest.approx(out[1][2], rel=1e-3, abs=1e-7)
|