Skip to content

Commit 8fc2555

Browse files
danielsobradodanielsobrado
committed
feat: add mixture of experts (moe) animation
1 parent 6665978 commit 8fc2555

File tree

11 files changed

+506
-0
lines changed

11 files changed

+506
-0
lines changed

‎moe-animation/index.html‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
<!doctype html>
2+
<html lang="en">
3+
4+
<head>
5+
<meta charset="UTF-8" />
6+
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
7+
<title>Mixture of Experts Animation</title>
8+
</head>
9+
10+
<body>
11+
<div id="root"></div>
12+
<script type="module" src="/src/main.jsx"></script>
13+
</body>
14+
15+
</html>

‎moe-animation/package.json‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
{
2+
"name": "moe-animation",
3+
"version": "1.0.0",
4+
"type": "module",
5+
"scripts": {
6+
"dev": "vite",
7+
"build": "vite build",
8+
"preview": "vite preview"
9+
},
10+
"dependencies": {
11+
"gsap": "^3.12.2",
12+
"react": "^18.2.0",
13+
"react-dom": "^18.2.0",
14+
"three": "^0.158.0"
15+
},
16+
"devDependencies": {
17+
"@types/react": "^18.2.15",
18+
"@types/react-dom": "^18.2.7",
19+
"@vitejs/plugin-react": "^4.0.3",
20+
"autoprefixer": "^10.4.16",
21+
"postcss": "^8.4.31",
22+
"tailwindcss": "^3.3.5",
23+
"vite": "^5.0.0"
24+
}
25+
}

‎moe-animation/postcss.config.js‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
export default {
2+
plugins: {
3+
tailwindcss: {},
4+
autoprefixer: {},
5+
},
6+
}

‎moe-animation/src/App.jsx‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import React, { useState } from 'react';
2+
import RoutingPanel from './RoutingPanel';
3+
import GatingPanel from './GatingPanel';
4+
import LoadBalancingPanel from './LoadBalancingPanel';
5+
6+
export default function App() {
7+
const [numExperts, setNumExperts] = useState(8);
8+
const [topK, setTopK] = useState(2);
9+
const [batchSize, setBatchSize] = useState(10);
10+
const [expertLoads, setExpertLoads] = useState(Array(8).fill(0));
11+
12+
const handleGenerate = () => {
13+
// Simulate load updates
14+
const newLoads = [...expertLoads];
15+
for (let i = 0; i < batchSize; i++) {
16+
// Randomly assign to topK experts
17+
for (let k = 0; k < topK; k++) {
18+
const idx = Math.floor(Math.random() * numExperts);
19+
newLoads[idx]++;
20+
}
21+
}
22+
setExpertLoads(newLoads);
23+
};
24+
25+
// Reset loads when expert count changes
26+
React.useEffect(() => {
27+
setExpertLoads(Array(numExperts).fill(0));
28+
}, [numExperts]);
29+
30+
return (
31+
<div className="min-h-screen bg-slate-950 p-8 font-sans text-slate-200">
32+
<header className="max-w-7xl mx-auto mb-8 flex justify-between items-end">
33+
<div>
34+
<h1 className="text-4xl font-black text-transparent bg-clip-text bg-gradient-to-r from-neon-blue to-neon-purple mb-2">
35+
Mixture of Experts
36+
</h1>
37+
<p className="text-slate-400 text-lg">
38+
Visualizing Sparse Gating & Expert Routing
39+
</p>
40+
</div>
41+
<div className="flex gap-4">
42+
<div className="bg-slate-900 p-4 rounded-lg border border-slate-800">
43+
<label className="block text-xs font-bold text-slate-500 uppercase mb-1">Experts</label>
44+
<div className="flex gap-2">
45+
{[4, 8, 16].map(n => (
46+
<button
47+
key={n}
48+
onClick={() => setNumExperts(n)}
49+
className={`px-3 py-1 rounded text-sm font-bold transition-colors ${numExperts === n ? 'bg-neon-blue text-black' : 'bg-slate-800 hover:bg-slate-700'
50+
}`}
51+
>
52+
{n}
53+
</button>
54+
))}
55+
</div>
56+
</div>
57+
58+
<div className="bg-slate-900 p-4 rounded-lg border border-slate-800">
59+
<label className="block text-xs font-bold text-slate-500 uppercase mb-1">Top-K</label>
60+
<div className="flex gap-2">
61+
{[1, 2].map(k => (
62+
<button
63+
key={k}
64+
onClick={() => setTopK(k)}
65+
className={`px-3 py-1 rounded text-sm font-bold transition-colors ${topK === k ? 'bg-neon-pink text-black' : 'bg-slate-800 hover:bg-slate-700'
66+
}`}
67+
>
68+
{k}
69+
</button>
70+
))}
71+
</div>
72+
</div>
73+
</div>
74+
</header>
75+
76+
<main className="max-w-7xl mx-auto grid grid-cols-1 lg:grid-cols-3 gap-8">
77+
{/* Main Visualization */}
78+
<div className="lg:col-span-2 aspect-video">
79+
<RoutingPanel
80+
numExperts={numExperts}
81+
topK={topK}
82+
batchSize={batchSize}
83+
onGenerate={handleGenerate}
84+
/>
85+
</div>
86+
87+
{/* Stats / Info Panel */}
88+
<div className="space-y-6">
89+
<div className="bg-slate-900 p-6 rounded-xl border border-slate-800">
90+
<h2 className="text-xl font-bold text-white mb-4 flex items-center gap-2">
91+
<span className="w-2 h-8 bg-neon-green rounded-full"></span>
92+
How it Works
93+
</h2>
94+
<div className="space-y-4 text-sm text-slate-400">
95+
<p>
96+
<strong className="text-white">1. The Router:</strong> A learned gating network that predicts which experts are best suited for each token.
97+
</p>
98+
<p>
99+
<strong className="text-white">2. Sparse Activation:</strong> Instead of using all parameters (Dense), we only use the Top-{topK} experts. This saves massive compute.
100+
</p>
101+
<p>
102+
<strong className="text-white">3. Load Balancing:</strong> Ideally, tokens are spread evenly. If one expert gets too many, it becomes a bottleneck (Expert Collapse).
103+
</p>
104+
</div>
105+
106+
<div className="mt-6">
107+
<GatingPanel numExperts={numExperts} topK={topK} />
108+
</div>
109+
</div>
110+
111+
<LoadBalancingPanel numExperts={numExperts} loads={expertLoads} />
112+
</div>
113+
</main>
114+
</div>
115+
);
116+
}

‎moe-animation/src/GatingPanel.jsx‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import React, { useEffect, useRef } from 'react';
2+
3+
export default function GatingPanel({ numExperts, topK }) {
4+
const canvasRef = useRef(null);
5+
6+
useEffect(() => {
7+
const canvas = canvasRef.current;
8+
if (!canvas) return;
9+
const ctx = canvas.getContext('2d');
10+
11+
// Simulate some logits
12+
const logits = Array.from({ length: numExperts }, () => Math.random() * 5);
13+
// Softmax
14+
const exp = logits.map(x => Math.exp(x));
15+
const sum = exp.reduce((a, b) => a + b, 0);
16+
const probs = exp.map(x => x / sum);
17+
18+
// Find top-k indices
19+
const indices = probs.map((p, i) => ({ p, i }))
20+
.sort((a, b) => b.p - a.p)
21+
.slice(0, topK)
22+
.map(x => x.i);
23+
24+
// Draw
25+
const width = canvas.width;
26+
const height = canvas.height;
27+
const barWidth = width / numExperts;
28+
29+
ctx.clearRect(0, 0, width, height);
30+
31+
probs.forEach((p, i) => {
32+
const isSelected = indices.includes(i);
33+
const x = i * barWidth;
34+
const h = p * height * 0.8;
35+
const y = height - h;
36+
37+
// Bar
38+
ctx.fillStyle = isSelected ? '#00f3ff' : '#334155';
39+
ctx.fillRect(x + 2, y, barWidth - 4, h);
40+
41+
// Label
42+
ctx.fillStyle = '#94a3b8';
43+
ctx.font = '10px monospace';
44+
ctx.fillText(`E${i}`, x + barWidth / 2 - 6, height - 5);
45+
46+
// Value
47+
if (isSelected) {
48+
ctx.fillStyle = '#ffffff';
49+
ctx.fillText(p.toFixed(2), x + 2, y - 5);
50+
}
51+
});
52+
53+
}, [numExperts, topK]);
54+
55+
return (
56+
<div className="bg-slate-900 p-4 rounded-lg border border-slate-800">
57+
<h3 className="text-sm font-bold text-slate-400 mb-2 uppercase">Gating Probabilities</h3>
58+
<canvas ref={canvasRef} width={300} height={100} className="w-full h-24" />
59+
</div>
60+
);
61+
}

‎moe-animation/src/LoadBalancingPanel.jsx‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import React from 'react';
2+
3+
export default function LoadBalancingPanel({ numExperts, loads }) {
4+
// loads is an array of numbers representing tokens processed
5+
const maxLoad = Math.max(...(loads || [1]));
6+
7+
return (
8+
<div className="bg-slate-900 p-6 rounded-xl border border-slate-800">
9+
<h2 className="text-xl font-bold text-white mb-4">Expert Load</h2>
10+
<div className="space-y-2">
11+
{Array.from({ length: numExperts }).map((_, i) => {
12+
const load = loads ? loads[i] : 0;
13+
const percentage = (load / maxLoad) * 100 || 0;
14+
15+
return (
16+
<div key={i} className="flex items-center gap-2">
17+
<span className="text-xs font-mono w-6 text-slate-500">E{i}</span>
18+
<div className="flex-1 h-2 bg-slate-800 rounded-full overflow-hidden">
19+
<div
20+
className={`h-full transition-all duration-500 ${percentage > 90 ? 'bg-red-500' : 'bg-neon-green'
21+
}`}
22+
style={{ width: `${percentage}%` }}
23+
></div>
24+
</div>
25+
<span className="text-xs text-slate-500 w-8 text-right">{load}</span>
26+
</div>
27+
);
28+
})}
29+
</div>
30+
<p className="text-xs text-slate-500 mt-4">
31+
*Red indicates potential bottleneck (Expert Collapse risk)
32+
</p>
33+
</div>
34+
);
35+
}

0 commit comments

Comments
 (0)