TensorRT: NHWC transformer trunk fails to build at large maxBatchSize (Error Code 10)

[memory1116]

Summary

Since d0c49c67 ("Default TensorRT transformer trunks to NHWC", #1218) made
trtTransformerNHWC = true the default in v1.17, transformer models fail to build a
TensorRT engine when maxBatchSize is large. With numSearchThreads = 576 the build
fails every time, so KataGo cannot start at all unless the user knows to set
trtTransformerNHWC = false.

NCHW builds fine at the same batch sizes on the same machine, so this is specific to the
NHWC trunk layout.

Error

IBuilder::buildSerializedNetwork: Error Code 10: Internal Error
(Could not find any implementation for node
 {ForeignNode[model.blocks.0.normactconvp.conv.Wnhwc + ONNXTRT_Broadcast...model.act_trunkfinal]}.
 In computeCosts at /_src/optimizer/common/tactic/optimizer.cpp:4265)

The process then aborts (exit code 134).

Reproduction

katago benchmark -model 28b.bin.gz -config default_gtp.cfg \
  -v 60000 -t 576 \
  -override-config numNNServerThreadsPerModel=16,trtTransformerNHWC=true

3/3 runs failed. Same config with trtTransformerNHWC=false: 3/3 succeeded.

The failure depends on maxBatchSize

maxBatchSize is derived from the thread count, and it determines whether the build
succeeds. Evidence from the trtcache directory — NCHW produced plan-cache entries for
every batch size tried, NHWC only for the small ones:

maxBatchSize	NCHW	NHWC
b16	built	built
b64	built	built
b576	built	never builds

This also explains why the failure initially looked intermittent: at 16 NN server threads
it reproduced roughly 1 in 10 runs, which is near a batch-size boundary where success
depends on the exact graph. At b576 it is deterministic.

Environment

• KataGo v1.17.1 (also reproduced on v1.17.1 + Serialize TensorRT engine builds across GPUs #1225)
• TensorRT 10.16.1.11
• 8× RTX 4090 D (sm89), Driver 595.71.05, CUDA 13.0
• Ubuntu 22.04, g++ 11.4
• Model: b11c768h12nbt3tflrs-fson-silu (nbt transformer, 70.4M params), FP16
• numNNServerThreadsPerModel = 16, 2 threads per GPU

Notes on what it is not

• Not the JIT tactic source restriction. Adding a switch to drop
  setTacticSources(kJIT_CONVOLUTIONS) did not help: 3/3 runs still failed with
  Error Code 10. Those lines are also unchanged since v1.16.5.
• Not a startup-concurrency problem. It reproduces with a single NN server thread
  building on one GPU.
• Not runtime-related. Once a plan cache exists for a batch size that does build, NHWC
  inference runs fine — the failure is confined to buildSerializedNetwork.

Suggested handling

The FP32 pinning may be involved: the failing node is a fused region spanning
conv.Wnhwc through act_trunkfinal, and the ONNX path sets
kOBEY_PRECISION_CONSTRAINTS with RMSNorm/head layers pinned to FP32. A hard constraint
over a fused NHWC region may leave TensorRT with no valid tactic. Possible directions:

1. Fall back to NCHW automatically when the NHWC build fails, instead of aborting.
2. Relax the pinned layers from kOBEY to kPREFER for the NHWC path (needs accuracy
   validation).
3. Reconsider NHWC as the default until the large-batch case is resolved, since users
   running high thread counts hit this on first startup with no obvious diagnosis.

nhwc_b576_fail_1.log
nhwc_b576_fail_2.log
nhwc_b576_fail_3.log

[memory1116]

BTW: the "28b" model is actually the latest 11b transformer weight.

[lightvector]

@memory1116 Thanks for the report. Thoughts about @zsqdx 's patch here #1229 ?

I wonder if there is a way to auto-adjust the workspace limit because it seems like if we set a fixed number here then the next user that comes along with a larger net or wants to test an even larger batch size might just hit the limit again.

[zsqdx]

Yes — TensorRT already has a device-dependent automatic default here. If we don't call setMemoryPoolLimit(kWORKSPACE, ...), the workspace limit defaults to the target GPU's total memory. It is only a tactic-selection ceiling, not a preallocation; each context still allocates only what its selected plan needs.

I updated #1229 to use that default. A fixed 2 GiB limit fixes b576 but fails again at b1024; the automatic version builds b576, b1024, and b2048. The exact 16-NN / 8-GPU b576 benchmark and testgpuerror both pass, and two b576 contexts use about 5.47 GiB total, essentially the same as with the 2 GiB cap.

TensorRT docs: https://docs.nvidia.com/deeplearning/tensorrt/latest/reference/troubleshooting-faq.html#how-do-i-choose-the-optimal-workspace-size