| |
| """Bring-up check of ``tt.l1res`` (OPT round 4): the real encoder weights, the round-3 encoder vs the L1-resident one, |
| per bottleneck of layers 3-4, every intermediate (conv1, conv2, conv3, identity, output) compared bit for bit. |
| Both encoders get the same layer-2 output (the round-3 encoder's). |
| |
| bin/devrun -t 900 -- python -u code/scripts/debug_l1res.py |
| """ |
| from __future__ import annotations |
|
|
| import sys |
|
|
| import numpy as np |
|
|
|
|
| def main(): |
| import torch |
| import ttnn |
|
|
| from tt_bevdet.device import close_device, open_device |
| from tt_bevdet.reference.config import default_config |
| from tt_bevdet.reference.weights import BEVDetWeights |
| from tt_bevdet.tests import paths |
| from tt_bevdet.tt.conv_tune import apply_tune |
| from tt_bevdet.tt.image_encoder import TtImageEncoder |
| from tt_bevdet.tt.l1res import L1ResBlock, block_mc, ensure |
| from tt_bevdet.tt.params import DEFAULT_PRECISION, BEVDetParams |
| from tt_bevdet.ttaw.ops.conv import is_dram |
|
|
| def bits(t): |
| if t.memory_config().is_sharded(): |
| t = ttnn.to_memory_config(t, ttnn.DRAM_MEMORY_CONFIG) |
| return ttnn.to_torch(ttnn.from_device(t)).to(torch.bfloat16).view(torch.int16).numpy().reshape( |
| -1, int(t.shape[-1])) |
|
|
| def cmp(tag, a, b): |
| x, y = bits(a), bits(b) |
| print(f" {tag}: {int(np.count_nonzero(x != y))} differ; l1 mc " |
| f"{'dram' if is_dram(b) else b.memory_config().shard_spec.shape if b.memory_config().is_sharded() else 'l1'}", |
| flush=True) |
|
|
| cfg = default_config() |
| params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg) |
| dev = open_device() |
| try: |
| kw = dict(precision=DEFAULT_PRECISION, stem_in="kin", conv_l1="all", backbone_res="next12", mm1x1=True) |
| ref = TtImageEncoder(params, **kw) |
| l1 = TtImageEncoder(params, l1res=True, **kw) |
| apply_tune(ref.convs()) |
| apply_tune(l1.convs()) |
| n, (h, w) = cfg.n_cams, cfg.input_hw |
| rng = np.random.default_rng(7) |
| crop = rng.integers(0, 256, (n, h, w, 3), dtype=np.uint8) |
| din = ttnn.from_torch(torch.from_numpy(crop.reshape(1, 1, n * h, w * 3)), dtype=ttnn.uint8, |
| layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG) |
| x = ref.kin_stem(din) |
| for blk in ref.stages.stages[0] + ref.stages.stages[1]: |
| x = blk(x) |
| xr = xl = x |
| gy = int(dev.compute_with_storage_grid_size().y) |
| for s in (2, 3): |
| for rb, lb in zip(ref.stages.stages[s], l1.stages.stages[s]): |
| b = rb |
| lbb: L1ResBlock = lb |
| print(f"{b.name}: in ref {'dram' if is_dram(xr.tensor) else 'l1'} / l1 " |
| f"{'dram' if is_dram(xl.tensor) else 'l1'}", flush=True) |
| y_r = b.conv1(xr) |
| if is_dram(xl.tensor): |
| y_l = lbb.blk.conv1(xl) |
| else: |
| y_l = lbb._mm(ensure(xl, block_mc(xl.rows, xl.channels, 8, gy))[0], lbb.w1, |
| b.conv3.in_channels, True, gy) |
| cmp("conv1", y_r.tensor, y_l.tensor) |
| z_r = b.conv2(y_r) |
| getattr(lbb.blk.conv2, "conv").slicing = "l1" if is_dram(y_l.tensor) else "auto" |
| z_l = lbb.blk.conv2(y_l) |
| cmp("conv2", z_r.tensor, z_l.tensor) |
| zs, mv = ensure(z_l, block_mc(z_l.rows, z_l.channels, 8, gy)) |
| print(" conv2 moved:", mv, flush=True) |
| y3_r = b.conv3(z_r) |
| y3_l = lbb._mm(zs, lbb.w3, b.conv3.out_channels, False, gy) |
| cmp("conv3", y3_r.tensor, y3_l.tensor) |
| if b.downsample is not None: |
| d_r = b.downsample(xr) |
| xd = xl if is_dram(xl.tensor) else xl.with_tensor( |
| ttnn.to_memory_config(xl.tensor, ttnn.DRAM_MEMORY_CONFIG)) |
| d_l = lbb.blk.downsample(xd) |
| cmp("downsample", d_r.tensor, d_l.tensor) |
| xr = b(xr) |
| xl = lb(xl) |
| cmp("out", xr.tensor, xl.tensor) |
| finally: |
| close_device(dev) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
