Web demo: weight tying + WebRTC resilience (relay fallback, keepalive)

web/public/app.js

@@ -13,6 +13,7 @@ const STUN = [
   { urls: "turn:openrelay.metered.ca:443", username: "openrelayproject", credential: "openrelayproject" },
   { urls: "turns:openrelay.metered.ca:443?transport=tcp", username: "openrelayproject", credential: "openrelayproject" },
 ];
+let rtcConfig = { iceServers: STUN };          // /mode can override (own TURN etc.)
 
 const ui = {
   status: document.getElementById("status"),
@@ -129,6 +130,10 @@ function connectSignaling() {
       log(`${nmeOf(msg.id)} left`); cleanupPeer(msg.id); names.delete(msg.id); updatePeers();
     } else if (msg.type === "signal") {
       await onSignal(msg.from, msg.data);
+    } else if (msg.type === "ping") {          // keepalive (proxies close idle sockets)
+      ws.send(JSON.stringify({ type: "pong" }));
+    } else if (msg.type === "relay") {         // payload relayed via the server (WS fallback)
+      onGrad(msg.from, bufFromB64(msg.data));
     }
   };
 }
@@ -157,10 +162,18 @@ function addJoinRequest(id, name) {
 }
 
 function newPC(peerId) {
-  const pc = new RTCPeerConnection({ iceServers: STUN });
+  const pc = new RTCPeerConnection(rtcConfig);
   pc.onicecandidate = (e) => { if (e.candidate) signal(peerId, { candidate: e.candidate }); };
   pc.onconnectionstatechange = () => {
     if (pc.connectionState === "failed") { cleanupPeer(peerId); scheduleReconnect(peerId); }
+    // "disconnected" often self-heals; give it a grace period, then treat as failed
+    if (pc.connectionState === "disconnected")
+      setTimeout(() => {
+        if (pcs.get(peerId) === pc && pc.connectionState === "disconnected") {
+          log(`${nmeOf(peerId)} connection did not recover — redialing`);
+          cleanupPeer(peerId); scheduleReconnect(peerId);
+        }
+      }, 4000);
   };
   pcs.set(peerId, pc);
   return pc;
@@ -171,14 +184,37 @@ function newPC(peerId) {
 // them), redial with backoff. Only the higher-numbered peer initiates, so both
 // sides don't collide (offer glare). If the peer truly left, names loses them
 // and the retries stop on their own.
+// ---- WebSocket relay fallback (PairDrop's WSPeer pattern) ---------------------
+// A peer still in the room after 5 failed WebRTC attempts is reachable through
+// the signaling server even though no direct path exists (symmetric NATs, no
+// TURN). Install a pseudo data channel that relays binary over the WS instead
+// of dropping them. Devices that actually left are removed via peer-left.
+function b64FromBuf(buf) {
+  const u = new Uint8Array(buf); let s = "";
+  for (let i = 0; i < u.length; i += 0x8000) s += String.fromCharCode(...u.subarray(i, i + 0x8000));
+  return btoa(s);
+}
+function bufFromB64(s) { return Uint8Array.from(atob(s), c => c.charCodeAt(0)).buffer; }
+function makeRelayChannel(peerId) {
+  return {
+    isRelay: true,
+    get readyState() { return ws && ws.readyState === 1 && names.has(peerId) ? "open" : "closed"; },
+    get bufferedAmount() { return ws ? ws.bufferedAmount : 0; },
+    send(buf) { ws.send(JSON.stringify({ type: "relay", to: peerId, data: b64FromBuf(buf) })); },
+    close() { chans.delete(peerId); },
+  };
+}
+
 const reconnectTimers = new Map();             // peerId -> attempt count
 function scheduleReconnect(peerId, attempt = 0) {
   if (!names.has(peerId) || chans.has(peerId) || reconnectTimers.has(peerId)) return;
-  if (attempt >= 5) {                          // 5 tries, then they're out of the group
-    log(`${nmeOf(peerId)} unreachable after 5 reconnect attempts — removed from the group`);
-    cleanupPeer(peerId);
-    names.delete(peerId);
-    updatePeers();
+  if (attempt >= 5) {                          // no direct path — fall back to server relay
+    if (names.has(peerId)) {
+      log(`no direct WebRTC path to ${nmeOf(peerId)} after 5 attempts — relaying through the server instead`);
+      chans.set(peerId, makeRelayChannel(peerId));
+      updatePeers(); wake();
+      ui.start.disabled = false;
+    }
     return;
   }
   const delay = 1500 * Math.pow(2, attempt);
@@ -817,6 +853,7 @@ function buildModel(cfg) {
   // to create their own room (they become its host) or join one by code.
   try {
     const mode = await (await fetch("mode")).json();
+    if (mode.rtc) rtcConfig = mode.rtc;                  // operator-provided ICE config
     if (mode.forceRooms && !room()) {
       const lobby = document.getElementById("lobby");
       for (const c of document.querySelectorAll(".card")) if (c !== lobby) c.style.display = "none";
@@ -887,6 +924,15 @@ function buildModel(cfg) {
   }
   updatePeers();
   connectSignaling();
+  // phones suspend the page (screen off, app switch) — reconnect the moment we
+  // come back, and on network changes (wifi <-> cellular), PairDrop-style
+  document.addEventListener("visibilitychange", () => {
+    if (!document.hidden && (!ws || ws.readyState > 1)) { log("page woke — reconnecting"); connectSignaling(); }
+  });
+  if (navigator.connection && navigator.connection.addEventListener)
+    navigator.connection.addEventListener("change", () => {
+      if (!ws || ws.readyState > 1) { log("network changed — reconnecting"); connectSignaling(); }
+    });
   // dataset: stream the chosen dataset from HuggingFace; built-in corpus if offline
   loadDataset((ui.cfgData.value || "").trim());
   ui.start.onclick = async () => {

web/public/transformer.js

@@ -169,7 +169,9 @@
         Wv: add(`b${l}.Wv`, mk(c * c, 0.08)), Wo: add(`b${l}.Wo`, mk(c * c, 0.08)),
         W1: add(`b${l}.W1`, mk(c * hidden, 0.08)), W2: add(`b${l}.W2`, mk(hidden * c, 0.08)),
       });
-    m.Wu = add("Wu", mk(c * vocabSize(), 0.08));
+    // weight-tied unembedding: logits use embᵀ (no separate Wu). Halves the
+    // vocab-sized parameters — and with a 16k vocab that's ~half of ALL
+    // parameters, so gradients over the wire shrink ~2× too.
     m.nParams = params.reduce((a, p) => a + p.length, 0);
     return m;
   }
@@ -273,7 +275,7 @@
       cache.blocks.push(cb);
     }
     const lf = lnFwd(x, BT, C); cache.lnf = lf; cache.xf = x;
-    const logits = await vmm(lf.y, m.Wu, BT, C, vocab, ctx);
+    const logits = await vmm(lf.y, TC.transpose(m.emb, vocab, C), BT, C, vocab, ctx);  // tied: embᵀ
     // cross-entropy + dlogits
     let loss = 0;
     const dlogits = new Float32Array(BT * vocab);
@@ -298,9 +300,10 @@
     const BT = B * T, hd = C / heads, mm = TC.matmul, tr = TC.transpose;
     const g = m.params.map(p => new Float32Array(p.length));
     const gi = Object.fromEntries(m.names.map((n, i) => [n, i]));
-    // unembed
-    g[gi.Wu] = mm(tr(cache.lnf.y, BT, C), cache.dlogits, C, BT, vocab);
-    let dx = lnBwd(mm(cache.dlogits, tr(m.Wu, C, vocab), BT, vocab, C), cache.lnf.y, cache.lnf.sig, BT, C);
+    // tied unembed: logits = lnf @ embᵀ, so the unembedding gradient flows
+    // straight into emb — dlogitsᵀ @ lnf is V×C, emb's own shape
+    g[gi.emb] = mm(tr(cache.dlogits, BT, vocab), cache.lnf.y, vocab, BT, C);
+    let dx = lnBwd(mm(cache.dlogits, m.emb, BT, vocab, C), cache.lnf.y, cache.lnf.sig, BT, C);
     const scale = 1 / Math.sqrt(hd);
     for (let l = layers - 1; l >= 0; l--) {
       const bl = m.blocks[l], cb = cache.blocks[l];

web/server.js

@@ -30,13 +30,28 @@ const PUB = path.join(__dirname, "public");
 const TYPES = { ".html": "text/html", ".js": "text/javascript",
                 ".css": "text/css", ".json": "application/json" };
 
+// keepalive: proxies (incl. HuggingFace's) close idle WebSockets — ping every
+// 30s keeps the pipe warm and detects dead clients within a minute
+setInterval(() => {
+  for (const room of rooms.values())
+    for (const [pid, v] of room.peers) {
+      if (v.ws.isAlive === false) { v.ws.terminate(); continue; }
+      v.ws.isAlive = false;
+      send(v.ws, { type: "ping" });
+    }
+}, 30000);
+
 const server = http.createServer((req, res) => {
   let p = decodeURIComponent(req.url.split("?")[0]);
   if (p === "/mode") {                 // deployment flags for the client
     res.writeHead(200, { "Content-Type": "application/json" });
     // DAISY_FORCE_ROOMS=1 (set on the HF Space): no LAN auto-grouping — every
-    // visitor creates their own private room and invites devices by link
-    return res.end(JSON.stringify({ forceRooms: !!process.env.DAISY_FORCE_ROOMS }));
+    // visitor creates their own private room and invites devices by link.
+    // DAISY_RTC_CONFIG: optional JSON RTCPeerConnection config (own TURN etc.,
+    // PairDrop-style server-provided rtcConfig).
+    let rtc = null;
+    try { if (process.env.DAISY_RTC_CONFIG) rtc = JSON.parse(process.env.DAISY_RTC_CONFIG); } catch (e) {}
+    return res.end(JSON.stringify({ forceRooms: !!process.env.DAISY_FORCE_ROOMS, rtc }));
   }
   if (p === "/") p = "/index.html";
   const file = path.join(PUB, path.normalize(p));
@@ -89,6 +104,14 @@ wss.on("connection", (ws, req) => {
     if (msg.type === "signal" && msg.to && room.peers.has(id)) {   // relay WebRTC signaling
       const target = room.peers.get(msg.to);
       if (target) send(target.ws, { type: "signal", from: id, data: msg.data });
+    } else if (msg.type === "pong") {                              // keepalive reply
+      ws.isAlive = true;
+    } else if (msg.type === "relay" && msg.to && room.peers.has(id)) {
+      // PairDrop-style WS fallback: when two devices can't form a direct
+      // WebRTC path (symmetric NATs, no TURN), the training payloads flow
+      // through here instead of dropping the peer
+      const target = room.peers.get(msg.to);
+      if (target) send(target.ws, { type: "relay", from: id, data: msg.data });
     } else if (msg.type === "admit" && id === room.host) {         // host verdict on a joiner
       const p = room.pending.get(msg.id);
       if (!p) return;