Document WebGPU feasibility and browser training benchmarks

[ale18V]

What & why

WebGPU is feasible for DISCO and ready for opt-in support across Local, Federated, and Decentralized topologies (Closes #575). This builds on top of the merged Tiny Shakespeare task (#1244).

• Integration: Added @tensorflow/[email protected] to webapp. WebGPU is available as an opt-in via ?backend=webgpu with automated fallback to WebGL on failure.
• Full End-to-End Validation & Convergence:
  • Local (Titanic): 10/10 epochs in 11.45s, IndexedDB model saving verified.
  • Federated (Titanic, 3 peers: 1 WebGPU + 2 Node): 5 rounds, reached 87.2% validation accuracy. Mixed-backend aggregation (browser WebGPU + Node CPU/C++) verified.
  • Decentralized (MNIST, 3 peers: 1 WebGPU + 2 Node): 10 rounds over WebRTC data channels, reached 100% validation accuracy (loss $5.72 \times 10^{-7}$).
  • Federated nanoGPT (Shakespeare, 3 peers: 1 WebGPU + 2 Node): 8 rounds, 312 aggregated gradient updates cluster-wide, loss monotonically converged from 10.83 down to 5.72 (perplexity reduced from 50,257 to ~350).
• Correctness & TFJS WebGPU Bug Root Causes:
  • argMax Rank Mismatch: The GPT prediction shape mismatch previously hypothesized to be linked to closed upstream issue WebGPU: TransposeProgram WGSL shader fails with type mismatch on Chrome 141+ (Tint IR rewrite) tensorflow/tfjs#8638 was actually caused by TFJS's WebGPU kernel ArgMinMaxProgram (argminmax_webgpu.ts:30), which forces scalar reduction outputs from [] to [1], causing probs.argMax() on a 1D vector to return rank-1 rather than a rank-0 scalar. Fixed in DISCO's GPT.#predictSingle() using .asScalar().
  • tf.multinomial WGSL Shader Bug: For single-token sampling (batch=1, samples=1), WebGPU's MultinomialProgram WGSL shader evaluates UV coordinates coords / outShape as (0.0, 0.0). The shader's PRNG function evaluates to 0.0 unconditionally, causing sampling to always collapse to token index 0 (top-1 mode token \n). Fixed in DISCO by implementing CPU-assisted top-$k$ sampling.
  • (Note: Neither of these two WebGPU backend bugs currently has an open issue on tensorflow/tfjs; both are isolated and ready to be reported upstream with reproduction testcases.)
• Recommendation: Keep WebGL as the default in production and enable WebGPU via opt-in (?backend=webgpu), as implemented.

The tested Linux setup used these Chrome GPU flags:

--enable-unsafe-webgpu --use-angle=vulkan --enable-features=Vulkan --disable-software-rasterizer

Technical documentation

Comprehensive benchmark timings, multi-peer convergence logs, and WebGPU kernel analyses are documented in:

• docs/WEBGPU_FULL_RUNS.md
• docs/WEBGPU_RESEARCH.md

[JulienVig]

On my MacBook Air with M2 chip:

Model	M2 WebGL / WebGPU	M2 speedup	AMD 860M speedup (docs)
Titanic	22.3 / 6.3 ms	3.57×	1.38×
MNIST	63.7 / 26.0 ms	2.45×	1.08×
LUS	26.2 / 16.8 ms	1.55×	1.49×
Dog	37.7 / 20.0 ms	1.89×	1.61×
CIFAR MobileNet	179.8 / 40.2 ms	4.47×	1.21×
GPT nano	1153.7 / 176.3 ms	6.54×	3.24×