Spaces:

nanotron
/

ultrascale-playbook

Running

App Files Files Community

thomwolf HF staff commited on Sep 4, 2024

Commit

878bd55

1 Parent(s): f2c15d5

update

Browse files

Files changed (5) hide show

dist/index.html +138 -76
dist/main.bundle.js +0 -0
dist/main.bundle.js.map +0 -0
src/index.html +138 -76
src/memory.js +158 -81

dist/index.html CHANGED Viewed

@@ -9,6 +9,26 @@
     <title>FineWeb: decanting the web for the finest text data at scale</title>
     <link rel="stylesheet" href="style.css">
     <style>
         #graph svg {
             font-family: sans-serif;
         }
@@ -68,61 +88,101 @@
         <aside>We are extremely thankful to the whole <a href="https://distill.pub/">distill.pub</a> team for creating the template on which we based this blog post.</aside>
-        <div>
-            <label for="a">Attention Heads (a):</label>
-            <input type="range" id="a" name="a" min="1" max="128" value="8">
-            <input type="number" id="a_input" value="8" min="1" max="128">
-            <br>
-            <label for="b">Micro Batch Size (b):</label>
-            <input type="range" id="b" name="b" min="1" max="53248" value="32">
-            <input type="number" id="b_input" value="32" min="1" max="53248">
-            <br>
-            <label for="h">Hidden Dimension Size (h):</label>
-            <input type="range" id="h" name="h" min="1" max="16384" value="512">
-            <input type="number" id="h_input" value="512" min="128" max="16384">
-            <br>
-            <label for="h_ff">Feedforward Dimension Size (h_ff):</label>
-            <input type="range" id="h_ff" name="h_ff" min="1" max="65536" value="2048">
-            <input type="number" id="h_ff_input" value="2048" min="512" max="65536">
-            <br>
-            <label for="L">Number of Layers (L):</label>
-            <input type="range" id="L" name="L" min="1" max="126" value="12">
-            <input type="number" id="L_input" value="12" min="1" max="126">
-            <br>
-            <label for="s">Sequence Length (s):</label>
-            <input type="range" id="s" name="s" min="1" max="128000" value="128">
-            <input type="number" id="s_input" value="128" min="64" max="128000">
-            <br>
-            <label for="v">Vocabulary Size (v):</label>
-            <input type="range" id="v" name="v" min="1000" max="100000" value="30522">
-            <input type="number" id="v_input" value="30522" min="1000" max="100000">
-            <br>
-            <label for="mixed">Mixed Precision:</label>
-            <input type="checkbox" id="mixed" name="mixed" checked>
-            <br>
-            <label for="recomputation">Recomputation:</label>
-            <select id="recomputation" name="recomputation">
-                <option value="none">None</option>
-                <option value="selective">Selective</option>
-                <option value="full">Full</option>
-            </select>
-            <br>
-            <label for="ff_activation">FF Activation:</label>
-            <select id="ff_activation" name="ff_activation">
-                <option value="relu">ReLU</option>
-                <option value="gelu">GELU</option>
-                <option value="swiglu">SwiGLU</option>
-            </select>
-            <br>
-            <label for="presets">Presets:</label>
-            <select id="presets" name="presets">
-                <option value="Tiny">Tiny</option>
-                <option value="8B">8B</option>
-                <option value="70B">70B</option>
-                <option value="405B">405B</option>
-            </select>
         </div>
-        <div id="graph" style="position: relative; width: 960px; height: 500px;"></div>
     <p><strong>TLDR:</strong> This blog covers a discussion on processing and evaluating data quality at scale, the 🍷 FineWeb
         recipe (listing and explaining all of our design choices), and the process followed to create its 📚 FineWeb-Edu subset.</p>
@@ -357,13 +417,13 @@ m_{act} =  2 * L* seq * bs * h * (34 +   \frac{5*n_{heads}*seq}{h})
 <p>Most frameworks these days use FlashAttention (TODO: see later) which makes the attention computation less memory intensive through kernel fusion, thus most trainings use the <code>full</code> settings.</p>
-<p>We can save some GPU memory with activation recomputation but this only delays by a bit the next bottleneck: as hinted earlier for LLM training there is usually a sweet spot for the GBST and we need to work out the training configuration backward from there. However, you can’t choose MBS to be an arbitrary large number on your GPU; at some point you will run out of GPU memory again since you need to store at least some of the activations in memory.</p>
 <p>There is a useful trick to compensate for that: <strong>gradient accumulation</strong> (<em>GradAcc</em>). With gradient accumulation we will split our batch in micro-batch, do forward and backward passes repeatedly on each micro-batch, compute the gradients, and, as the name suggests, sum the gradients step by step before doing a final optimizer step.</p>
-<p>We call the <code>micro batch size</code> (MBS) the batch size for each forward pass on a single node (the number of samples flowing through the model in one forward pass). We’ll refer to the overall batch size between each optimizer step as the <code>global batch size</code> (GBS). If we do one optimizer step each 8 forward/backward pass, the <code>global batch size</code> will be 8 times the <code>micro batch size</code>.</p>
-<p>What we now call <code>global batch size</code> thus corresponds to what we’ve called up to now just <code>batch size</code> for simplicity (we now make the terms more precise to avoid ambiguity).</p>
 <p>With gradient accumulation the global batch size can be computed as follows:</p>
@@ -377,7 +437,7 @@ BS = GBS=MBS * GradAcc
 <p>This is actually a bummer since the forward/backward passes for each micro-batch could actually totally be run in parallel. They are independent from each other and the only changing parameter are the input samples.</p>
-<p>Here comes data parallelism to solve exactly this problem! Let’s take a look, you say? Okay sure!</p>
 <h3>Data Parallelism</h3>
@@ -387,7 +447,7 @@ BS = GBS=MBS * GradAcc
 GBS=MBS * GradAcc * DP
 </d-math>
-<p>This means that we can reduce the number of gradient accumulation steps in favor of data parallel processes which speeds up training. In practice, people will tend to max out the number of data parallel nodes (the DP above) as much as possible as it’s inherently parallel versus the sequential Gradient Accumulation. Gradient accumulation is then added only to achieve a target batch size if DP alone is not sufficient. One exception to that is pipeline parallelism which we’ll discuss later.</p>
 <img src="assets/images/IMG_A95961668B3F-1.jpeg" alt="Data Parallelism">
@@ -403,17 +463,17 @@ GBS=MBS * GradAcc * DP
 <p>If the gradient accumulation ratio is lower than one, i.e. you have too many GPUs (!), you can either choose to not use all your GPUs or test if a lower MBS will speed up training. In these cases, you may want to prioritize throughput over the individual GPU utilization, you can then choose DP first and use a smaller MBS than possible in order to speed up training.</p>
-<p>Time to take a concrete example: We want to train a model with a GBS of 4M tokens and a sequence length of 4k. This means our batch size will be 1024 samples (we pick powers of two). We observe that a single of our GPU can fit MBS=2 in memory and we have 128 GPUs available for training. This means with 4 gradient accumulation steps we’ll achieve our goal of 1024 samples or 4M tokens per training step. Now what if we suddenly have 1024 GPUs available? We can achieve the same GBS and thus identical training by setting both MBS and gradient accumulation to 1 speeding up training significantly.</p>
 <p>[EXPERIMENTS WHERE WE INCREASE DP AND SHOW THROUGHPUT FOR SEVERAL MODELS]</p>
-<p>We’ve explored data parallelism, a simple strategy to scale training across more GPUs and gives consistent speed improvements. The keen reader might have noticed however that it rests on the assumption that we can fit at least one input sample forward pass (<em>MBS=1</em>) into our GPU memory. This is not always the case! In particular for larger models which often don’t fit into a single GPU anymore even with activation recomputations activated.</p>
-<p>In such case, we need to shard the model across devices! We’ll now study two complementary sharding methods, tensor and pipeline parallelism which are doing that. Let’s start by the simplest, tensor parallelism!</p>
 <h3>Tensor Parallelism</h3>
-<p>So you’ve exhausted all the previous textbook tricks to try to fit your model on a single GPU but it still doesn’t fit? Let’s try to distribute this model across several GPUs. Unlike DP we will not simply duplicate the model but various parts of the model instance will be living on various GPUs.</p>
 <p>If we take a look at a typical matrix multiplication (the core of a neural network), we can get an idea about how we could split the model:</p>
@@ -470,6 +530,8 @@ GeLU(XW1 + XW2) \neq GeLU(XW1) + GeLU(XW2)
 <p>If you rather like code, note that we can prove this with the following snippet as well:</p>
 <d-code block language="python">
 def example_gelu():
     from torch.nn.functional import gelu
@@ -495,9 +557,9 @@ def example_gelu():
     torch.testing.assert_close(y_row_1, y_row_2, rtol=1e-5, atol=1e-5)
 </d-code>
-<p>To avoid a synchronization step directly after the first MLP, we’ll thus start with Column Parallel and be able to directly perform parallel GELU.</p>
-<p>Now, what about the second MLP? Should it be column or row parallel? Let’s draft both options:</p>
 <ul>
     <li>Column Parallel followed by Column Parallel</li>
     <img src="assets/images/image%2013.png" alt="Column Parallel Schema 1">
@@ -505,9 +567,9 @@ def example_gelu():
     <img src="assets/images/image%2014.png" alt="Column Parallel Schema 2">
 </ul>
-<p>We see that the “Column Parallel followed by Row Parallel” schema only involves two communications instead of four. It’s thus the most efficient schema in terms of communications.</p>
-<p>Let’s take a quick look at the backward pass:</p>
 <img src="assets/images/image%2015.png" alt="Backward Pass 1">
 <img src="assets/images/image%2016.png" alt="Backward Pass 2">
@@ -586,7 +648,7 @@ if __name__ == "__main__":
     example_column_row_linear()
 </d-code>
-<p>Now that we’ve found the most efficient schema for the Feedforward part of the transformer, let’s take a look at the multi-head attention block (MHA).</p>
 <p>We can generally follow a similar approach where the Q, K, V will be split in a Column Parallel fashion and the output projection will be split along the Row dimension.</p>
@@ -604,7 +666,7 @@ if __name__ == "__main__":
 <p>Could we push this approach further?</p>
-<p>Sequence parallelism applies this same idea to other parts of our model. We’ve applied tensor parallelism to two main parts in our models where combination of MLP allowed to naturally split the weights along major axis.</p>
 <p>The rest of the model mostly comprises layer norms, dropout and various summation of residuals, these contribute little to the computation but come with rather large forward activations to store.</p>
@@ -640,14 +702,14 @@ if __name__ == "__main__":
 <p>In the transformer model, tokens have no inherent information about their positional information. For these reasons, we need to use a positional encoding function.</p>
-<p>Assuming that in the multi-head attention layer, <em>q_m</em> is the “position-aware” query vector corresponding to a token at position <em>m</em>, <em>k_n</em> the “position-aware” key vector corresponding to the token at position <em>n</em> and <em>f</em> is our position embedding function, we would like our position vector to be a function of the input vectors and absolute positions like this:</p>
 <d-math>
 q_m = f(q,m)
 k_n = f(k,n)
 </d-math>
-<p>We may also want the positional encoding to model relative positional information between two input tokens. Relative positions help the model to operate across longer context spans and even context lengths not seen during training. The attention operation is generally a dot product operation between “position-aware” vectors <em>q</em> and <em>k</em>, so for a positional encoding that contains relative positional information, we’ll want to have:</p>
 <d-math>
 <q_m, k_n> = g(q, k, m-n)
@@ -655,7 +717,7 @@ k_n = f(k,n)
 <p>In other words, we want the result of <em>⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩</em> to depend on the values of <em>q</em> and <em>k</em> themselves, as well as their relative position <em>m − n</em>, but not <em>m</em> and <em>n</em>. This way, the model can focus on the relative difference between two tokens rather than their absolute positions.</p>
-<p>Let’s show that the RoPE positional embedding formulation satisfies the above formula.</p>
 <p><strong>Rotation matrix</strong></p>
@@ -695,9 +757,9 @@ R(θ) =
 <p><strong>Implementation</strong></p>
-<p>In our case, our internal vectors (the activations in our model) have much more than two elements. Let’s pair elements to get 2D vectors and apply the 2D rotation operation on these pairs.</p>
-<p>There are combinatorially many ways we can pair elements but generally two options are the most popular for implementing RoPE: we call them the <em>interleaved</em> and <em>non-interleaved</em> versions. (It’s still rather unfortunate to have two popular options)</p>
 <ol>
     <li>In the interleaved version, we pair consecutive elements <em>(x<sub>0</sub>, x<sub>1</sub>),(x<sub>2</sub>,x<sub>3</sub>),…</em> before applying the rotation matrix:</li>
@@ -860,13 +922,13 @@ x_{d-1}\cos mθ_{d/2-1} + x_{d-1}\sin mθ_{d/2-1}  \\
 <h2>References</h2>
 <ul>
-    <li>Harm’s posts:
         <ul>
             <li><a href="https://www.harmdevries.com/post/context-length/">https://www.harmdevries.com/post/context-length/</a></li>
             <li><a href="https://www.harmdevries.com/post/model-size-vs-compute-overhead/">https://www.harmdevries.com/post/model-size-vs-compute-overhead/</a></li>
         </ul>
     </li>
-    <li>Stas’ guides:
         <ul>
             <li><a href="https://github.com/stas00/ml-engineering">https://github.com/stas00/ml-engineering</a></li>
             <li><a href="https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md">https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md</a></li>

     <title>FineWeb: decanting the web for the finest text data at scale</title>
     <link rel="stylesheet" href="style.css">
     <style>
+        #controls {
+            display: grid;
+            grid-template-columns: auto 1fr auto;
+            gap: 5px;
+            align-items: center;
+            max-width: 600px;
+            margin-bottom: 20px;
+        }
+        #controls label {
+            text-align: right;
+        }
+        #controls input[type="range"] {
+            width: 100%;
+        }
+        #controls input[type="number"] {
+            width: 60px;
+        }
+        #controls .row {
+            display: contents;
+        }
         #graph svg {
             font-family: sans-serif;
         }
         <aside>We are extremely thankful to the whole <a href="https://distill.pub/">distill.pub</a> team for creating the template on which we based this blog post.</aside>
+        <div id="graph" style="position: relative; width: 700px; height: 500px;"></div>
+        <div id="controls">
+            <div class="row">
+                <label for="a">Attention Heads (a):</label>
+                <input type="range" id="a" name="a" min="1" max="128" value="8">
+                <input type="number" id="a_input" value="8" min="1" max="128">
+            </div>
+            <div class="row">
+                <label for="b">Micro Batch Size (b):</label>
+                <input type="range" id="b" name="b" min="1" max="53248" value="32">
+                <input type="number" id="b_input" value="32" min="1" max="53248">
+            </div>
+            <div class="row">
+                <label for="h">Hidden Dimension Size (h):</label>
+                <input type="range" id="h" name="h" min="1" max="16384" value="512">
+                <input type="number" id="h_input" value="512" min="128" max="16384">
+            </div>
+            <div class="row">
+                <label for="h_ff">Feedforward Dimension Size (h_ff):</label>
+                <input type="range" id="h_ff" name="h_ff" min="1" max="65536" value="2048">
+                <input type="number" id="h_ff_input" value="2048" min="512" max="65536">
+            </div>
+            <div class="row">
+                <label for="L">Number of Layers (L):</label>
+                <input type="range" id="L" name="L" min="1" max="126" value="12">
+                <input type="number" id="L_input" value="12" min="1" max="126">
+            </div>
+            <div class="row">
+                <label for="s">Sequence Length (s):</label>
+                <input type="range" id="s" name="s" min="1" max="128000" value="128">
+                <input type="number" id="s_input" value="128" min="64" max="128000">
+            </div>
+            <div class="row">
+                <label for="v">Vocabulary Size (v):</label>
+                <input type="range" id="v" name="v" min="1000" max="100000" value="30522">
+                <input type="number" id="v_input" value="30522" min="1000" max="100000">
+            </div>
+            <div class="row">
+                <label for="k">Optimizer Parameters (k):</label>
+                <input type="range" id="k" name="k" min="1" max="16" value="8">
+                <input type="number" id="k_input" value="8" min="1" max="16">
+            </div>
+            <div class="row">
+                <label for="tp">Tensor Model Parallelism (t):</label>
+                <input type="range" id="tp" name="tp" min="1" max="16" value="8">
+                <input type="number" id="tp_input" value="8" min="1" max="16">
+            </div>
+            <div class="row">
+                <label for="dp">Data Model Parallelism (d):</label>
+                <input type="range" id="dp" name="dp" min="1" max="256" value="1">
+                <input type="number" id="dp_input" value="1" min="1" max="256">
+            </div>
+            <div class="row">
+                <label for="mixed">Mixed Precision:</label>
+                <input type="checkbox" id="mixed" name="mixed" checked>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="recomputation">Recomputation:</label>
+                <select id="recomputation" name="recomputation">
+                    <option value="none">None</option>
+                    <option value="selective">Selective</option>
+                    <option value="full">Full</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="zero">Zero:</label>
+                <select id="zero" name="zero">
+                    <option value="Optimizer">Optimizer</option>
+                    <option value="Gradients">Gradients</option>
+                    <option value="Parameters">Parameters</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="ff_activation">FF Activation:</label>
+                <select id="ff_activation" name="ff_activation">
+                    <option value="relu">ReLU</option>
+                    <option value="gelu">GELU</option>
+                    <option value="swiglu">SwiGLU</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="presets">Presets:</label>
+                <select id="presets" name="presets">
+                    <option value="Llama 3 Tiny">Llama 3 Tiny</option>
+                    <option value="Llama 3 8B">Llama 3 8B</option>
+                    <option value="Llama 3 70B">Llama 3 70B</option>
+                    <option value="Llama 3 405B">Llama 3 405B</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
         </div>
     <p><strong>TLDR:</strong> This blog covers a discussion on processing and evaluating data quality at scale, the 🍷 FineWeb
         recipe (listing and explaining all of our design choices), and the process followed to create its 📚 FineWeb-Edu subset.</p>
 <p>Most frameworks these days use FlashAttention (TODO: see later) which makes the attention computation less memory intensive through kernel fusion, thus most trainings use the <code>full</code> settings.</p>
+<p>We can save some GPU memory with activation recomputation but this only delays by a bit the next bottleneck: as hinted earlier for LLM training there is usually a sweet spot for the GBST and we need to work out the training configuration backward from there. However, you can't choose MBS to be an arbitrary large number on your GPU; at some point you will run out of GPU memory again since you need to store at least some of the activations in memory.</p>
 <p>There is a useful trick to compensate for that: <strong>gradient accumulation</strong> (<em>GradAcc</em>). With gradient accumulation we will split our batch in micro-batch, do forward and backward passes repeatedly on each micro-batch, compute the gradients, and, as the name suggests, sum the gradients step by step before doing a final optimizer step.</p>
+<p>We call the <code>micro batch size</code> (MBS) the batch size for each forward pass on a single node (the number of samples flowing through the model in one forward pass). We'll refer to the overall batch size between each optimizer step as the <code>global batch size</code> (GBS). If we do one optimizer step each 8 forward/backward pass, the <code>global batch size</code> will be 8 times the <code>micro batch size</code>.</p>
+<p>What we now call <code>global batch size</code> thus corresponds to what we've called up to now just <code>batch size</code> for simplicity (we now make the terms more precise to avoid ambiguity).</p>
 <p>With gradient accumulation the global batch size can be computed as follows:</p>
 <p>This is actually a bummer since the forward/backward passes for each micro-batch could actually totally be run in parallel. They are independent from each other and the only changing parameter are the input samples.</p>
+<p>Here comes data parallelism to solve exactly this problem! Let's take a look, you say? Okay sure!</p>
 <h3>Data Parallelism</h3>
 GBS=MBS * GradAcc * DP
 </d-math>
+<p>This means that we can reduce the number of gradient accumulation steps in favor of data parallel processes which speeds up training. In practice, people will tend to max out the number of data parallel nodes (the DP above) as much as possible as it's inherently parallel versus the sequential Gradient Accumulation. Gradient accumulation is then added only to achieve a target batch size if DP alone is not sufficient. One exception to that is pipeline parallelism which we'll discuss later.</p>
 <img src="assets/images/IMG_A95961668B3F-1.jpeg" alt="Data Parallelism">
 <p>If the gradient accumulation ratio is lower than one, i.e. you have too many GPUs (!), you can either choose to not use all your GPUs or test if a lower MBS will speed up training. In these cases, you may want to prioritize throughput over the individual GPU utilization, you can then choose DP first and use a smaller MBS than possible in order to speed up training.</p>
+<p>Time to take a concrete example: We want to train a model with a GBS of 4M tokens and a sequence length of 4k. This means our batch size will be 1024 samples (we pick powers of two). We observe that a single of our GPU can fit MBS=2 in memory and we have 128 GPUs available for training. This means with 4 gradient accumulation steps we'll achieve our goal of 1024 samples or 4M tokens per training step. Now what if we suddenly have 1024 GPUs available? We can achieve the same GBS and thus identical training by setting both MBS and gradient accumulation to 1 speeding up training significantly.</p>
 <p>[EXPERIMENTS WHERE WE INCREASE DP AND SHOW THROUGHPUT FOR SEVERAL MODELS]</p>
+<p>We've explored data parallelism, a simple strategy to scale training across more GPUs and gives consistent speed improvements. The keen reader might have noticed however that it rests on the assumption that we can fit at least one input sample forward pass (<em>MBS=1</em>) into our GPU memory. This is not always the case! In particular for larger models which often don't fit into a single GPU anymore even with activation recomputations activated.</p>
+<p>In such case, we need to shard the model across devices! We'll now study two complementary sharding methods, tensor and pipeline parallelism which are doing that. Let's start by the simplest, tensor parallelism!</p>
 <h3>Tensor Parallelism</h3>
+<p>So you've exhausted all the previous textbook tricks to try to fit your model on a single GPU but it still doesn't fit? Let's try to distribute this model across several GPUs. Unlike DP we will not simply duplicate the model but various parts of the model instance will be living on various GPUs.</p>
 <p>If we take a look at a typical matrix multiplication (the core of a neural network), we can get an idea about how we could split the model:</p>
 <p>If you rather like code, note that we can prove this with the following snippet as well:</p>
 <d-code block language="python">
+```
+</region_of_file_to_rewritten_file>
 def example_gelu():
     from torch.nn.functional import gelu
     torch.testing.assert_close(y_row_1, y_row_2, rtol=1e-5, atol=1e-5)
 </d-code>
+<p>To avoid a synchronization step directly after the first MLP, we'll thus start with Column Parallel and be able to directly perform parallel GELU.</p>
+<p>Now, what about the second MLP? Should it be column or row parallel? Let's draft both options:</p>
 <ul>
     <li>Column Parallel followed by Column Parallel</li>
     <img src="assets/images/image%2013.png" alt="Column Parallel Schema 1">
     <img src="assets/images/image%2014.png" alt="Column Parallel Schema 2">
 </ul>
+<p>We see that the "Column Parallel followed by Row Parallel" schema only involves two communications instead of four. It's thus the most efficient schema in terms of communications.</p>
+<p>Let's take a quick look at the backward pass:</p>
 <img src="assets/images/image%2015.png" alt="Backward Pass 1">
 <img src="assets/images/image%2016.png" alt="Backward Pass 2">
     example_column_row_linear()
 </d-code>
+<p>Now that we've found the most efficient schema for the Feedforward part of the transformer, let's take a look at the multi-head attention block (MHA).</p>
 <p>We can generally follow a similar approach where the Q, K, V will be split in a Column Parallel fashion and the output projection will be split along the Row dimension.</p>
 <p>Could we push this approach further?</p>
+<p>Sequence parallelism applies this same idea to other parts of our model. We've applied tensor parallelism to two main parts in our models where combination of MLP allowed to naturally split the weights along major axis.</p>
 <p>The rest of the model mostly comprises layer norms, dropout and various summation of residuals, these contribute little to the computation but come with rather large forward activations to store.</p>
 <p>In the transformer model, tokens have no inherent information about their positional information. For these reasons, we need to use a positional encoding function.</p>
+<p>Assuming that in the multi-head attention layer, <em>q_m</em> is the "position-aware" query vector corresponding to a token at position <em>m</em>, <em>k_n</em> the "position-aware" key vector corresponding to the token at position <em>n</em> and <em>f</em> is our position embedding function, we would like our position vector to be a function of the input vectors and absolute positions like this:</p>
 <d-math>
 q_m = f(q,m)
 k_n = f(k,n)
 </d-math>
+<p>We may also want the positional encoding to model relative positional information between two input tokens. Relative positions help the model to operate across longer context spans and even context lengths not seen during training. The attention operation is generally a dot product operation between "position-aware" vectors <em>q</em> and <em>k</em>, so for a positional encoding that contains relative positional information, we'll want to have:</p>
 <d-math>
 <q_m, k_n> = g(q, k, m-n)
 <p>In other words, we want the result of <em>⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩</em> to depend on the values of <em>q</em> and <em>k</em> themselves, as well as their relative position <em>m − n</em>, but not <em>m</em> and <em>n</em>. This way, the model can focus on the relative difference between two tokens rather than their absolute positions.</p>
+<p>Let's show that the RoPE positional embedding formulation satisfies the above formula.</p>
 <p><strong>Rotation matrix</strong></p>
 <p><strong>Implementation</strong></p>
+<p>In our case, our internal vectors (the activations in our model) have much more than two elements. Let's pair elements to get 2D vectors and apply the 2D rotation operation on these pairs.</p>
+<p>There are combinatorially many ways we can pair elements but generally two options are the most popular for implementing RoPE: we call them the <em>interleaved</em> and <em>non-interleaved</em> versions. (It's still rather unfortunate to have two popular options)</p>
 <ol>
     <li>In the interleaved version, we pair consecutive elements <em>(x<sub>0</sub>, x<sub>1</sub>),(x<sub>2</sub>,x<sub>3</sub>),…</em> before applying the rotation matrix:</li>
 <h2>References</h2>
 <ul>
+    <li>Harm's posts:
         <ul>
             <li><a href="https://www.harmdevries.com/post/context-length/">https://www.harmdevries.com/post/context-length/</a></li>
             <li><a href="https://www.harmdevries.com/post/model-size-vs-compute-overhead/">https://www.harmdevries.com/post/model-size-vs-compute-overhead/</a></li>
         </ul>
     </li>
+    <li>Stas' guides:
         <ul>
             <li><a href="https://github.com/stas00/ml-engineering">https://github.com/stas00/ml-engineering</a></li>
             <li><a href="https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md">https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md</a></li>

dist/main.bundle.js CHANGED Viewed

The diff for this file is too large to render. See raw diff

dist/main.bundle.js.map CHANGED Viewed

The diff for this file is too large to render. See raw diff

src/index.html CHANGED Viewed

@@ -9,6 +9,26 @@
     <title>FineWeb: decanting the web for the finest text data at scale</title>
     <link rel="stylesheet" href="style.css">
     <style>
         #graph svg {
             font-family: sans-serif;
         }
@@ -68,61 +88,101 @@
         <aside>We are extremely thankful to the whole <a href="https://distill.pub/">distill.pub</a> team for creating the template on which we based this blog post.</aside>
-        <div>
-            <label for="a">Attention Heads (a):</label>
-            <input type="range" id="a" name="a" min="1" max="128" value="8">
-            <input type="number" id="a_input" value="8" min="1" max="128">
-            <br>
-            <label for="b">Micro Batch Size (b):</label>
-            <input type="range" id="b" name="b" min="1" max="53248" value="32">
-            <input type="number" id="b_input" value="32" min="1" max="53248">
-            <br>
-            <label for="h">Hidden Dimension Size (h):</label>
-            <input type="range" id="h" name="h" min="1" max="16384" value="512">
-            <input type="number" id="h_input" value="512" min="128" max="16384">
-            <br>
-            <label for="h_ff">Feedforward Dimension Size (h_ff):</label>
-            <input type="range" id="h_ff" name="h_ff" min="1" max="65536" value="2048">
-            <input type="number" id="h_ff_input" value="2048" min="512" max="65536">
-            <br>
-            <label for="L">Number of Layers (L):</label>
-            <input type="range" id="L" name="L" min="1" max="126" value="12">
-            <input type="number" id="L_input" value="12" min="1" max="126">
-            <br>
-            <label for="s">Sequence Length (s):</label>
-            <input type="range" id="s" name="s" min="1" max="128000" value="128">
-            <input type="number" id="s_input" value="128" min="64" max="128000">
-            <br>
-            <label for="v">Vocabulary Size (v):</label>
-            <input type="range" id="v" name="v" min="1000" max="100000" value="30522">
-            <input type="number" id="v_input" value="30522" min="1000" max="100000">
-            <br>
-            <label for="mixed">Mixed Precision:</label>
-            <input type="checkbox" id="mixed" name="mixed" checked>
-            <br>
-            <label for="recomputation">Recomputation:</label>
-            <select id="recomputation" name="recomputation">
-                <option value="none">None</option>
-                <option value="selective">Selective</option>
-                <option value="full">Full</option>
-            </select>
-            <br>
-            <label for="ff_activation">FF Activation:</label>
-            <select id="ff_activation" name="ff_activation">
-                <option value="relu">ReLU</option>
-                <option value="gelu">GELU</option>
-                <option value="swiglu">SwiGLU</option>
-            </select>
-            <br>
-            <label for="presets">Presets:</label>
-            <select id="presets" name="presets">
-                <option value="Tiny">Tiny</option>
-                <option value="8B">8B</option>
-                <option value="70B">70B</option>
-                <option value="405B">405B</option>
-            </select>
         </div>
-        <div id="graph" style="position: relative; width: 960px; height: 500px;"></div>
     <p><strong>TLDR:</strong> This blog covers a discussion on processing and evaluating data quality at scale, the 🍷 FineWeb
         recipe (listing and explaining all of our design choices), and the process followed to create its 📚 FineWeb-Edu subset.</p>
@@ -357,13 +417,13 @@ m_{act} =  2 * L* seq * bs * h * (34 +   \frac{5*n_{heads}*seq}{h})
 <p>Most frameworks these days use FlashAttention (TODO: see later) which makes the attention computation less memory intensive through kernel fusion, thus most trainings use the <code>full</code> settings.</p>
-<p>We can save some GPU memory with activation recomputation but this only delays by a bit the next bottleneck: as hinted earlier for LLM training there is usually a sweet spot for the GBST and we need to work out the training configuration backward from there. However, you can’t choose MBS to be an arbitrary large number on your GPU; at some point you will run out of GPU memory again since you need to store at least some of the activations in memory.</p>
 <p>There is a useful trick to compensate for that: <strong>gradient accumulation</strong> (<em>GradAcc</em>). With gradient accumulation we will split our batch in micro-batch, do forward and backward passes repeatedly on each micro-batch, compute the gradients, and, as the name suggests, sum the gradients step by step before doing a final optimizer step.</p>
-<p>We call the <code>micro batch size</code> (MBS) the batch size for each forward pass on a single node (the number of samples flowing through the model in one forward pass). We’ll refer to the overall batch size between each optimizer step as the <code>global batch size</code> (GBS). If we do one optimizer step each 8 forward/backward pass, the <code>global batch size</code> will be 8 times the <code>micro batch size</code>.</p>
-<p>What we now call <code>global batch size</code> thus corresponds to what we’ve called up to now just <code>batch size</code> for simplicity (we now make the terms more precise to avoid ambiguity).</p>
 <p>With gradient accumulation the global batch size can be computed as follows:</p>
@@ -377,7 +437,7 @@ BS = GBS=MBS * GradAcc
 <p>This is actually a bummer since the forward/backward passes for each micro-batch could actually totally be run in parallel. They are independent from each other and the only changing parameter are the input samples.</p>
-<p>Here comes data parallelism to solve exactly this problem! Let’s take a look, you say? Okay sure!</p>
 <h3>Data Parallelism</h3>
@@ -387,7 +447,7 @@ BS = GBS=MBS * GradAcc
 GBS=MBS * GradAcc * DP
 </d-math>
-<p>This means that we can reduce the number of gradient accumulation steps in favor of data parallel processes which speeds up training. In practice, people will tend to max out the number of data parallel nodes (the DP above) as much as possible as it’s inherently parallel versus the sequential Gradient Accumulation. Gradient accumulation is then added only to achieve a target batch size if DP alone is not sufficient. One exception to that is pipeline parallelism which we’ll discuss later.</p>
 <img src="assets/images/IMG_A95961668B3F-1.jpeg" alt="Data Parallelism">
@@ -403,17 +463,17 @@ GBS=MBS * GradAcc * DP
 <p>If the gradient accumulation ratio is lower than one, i.e. you have too many GPUs (!), you can either choose to not use all your GPUs or test if a lower MBS will speed up training. In these cases, you may want to prioritize throughput over the individual GPU utilization, you can then choose DP first and use a smaller MBS than possible in order to speed up training.</p>
-<p>Time to take a concrete example: We want to train a model with a GBS of 4M tokens and a sequence length of 4k. This means our batch size will be 1024 samples (we pick powers of two). We observe that a single of our GPU can fit MBS=2 in memory and we have 128 GPUs available for training. This means with 4 gradient accumulation steps we’ll achieve our goal of 1024 samples or 4M tokens per training step. Now what if we suddenly have 1024 GPUs available? We can achieve the same GBS and thus identical training by setting both MBS and gradient accumulation to 1 speeding up training significantly.</p>
 <p>[EXPERIMENTS WHERE WE INCREASE DP AND SHOW THROUGHPUT FOR SEVERAL MODELS]</p>
-<p>We’ve explored data parallelism, a simple strategy to scale training across more GPUs and gives consistent speed improvements. The keen reader might have noticed however that it rests on the assumption that we can fit at least one input sample forward pass (<em>MBS=1</em>) into our GPU memory. This is not always the case! In particular for larger models which often don’t fit into a single GPU anymore even with activation recomputations activated.</p>
-<p>In such case, we need to shard the model across devices! We’ll now study two complementary sharding methods, tensor and pipeline parallelism which are doing that. Let’s start by the simplest, tensor parallelism!</p>
 <h3>Tensor Parallelism</h3>
-<p>So you’ve exhausted all the previous textbook tricks to try to fit your model on a single GPU but it still doesn’t fit? Let’s try to distribute this model across several GPUs. Unlike DP we will not simply duplicate the model but various parts of the model instance will be living on various GPUs.</p>
 <p>If we take a look at a typical matrix multiplication (the core of a neural network), we can get an idea about how we could split the model:</p>
@@ -470,6 +530,8 @@ GeLU(XW1 + XW2) \neq GeLU(XW1) + GeLU(XW2)
 <p>If you rather like code, note that we can prove this with the following snippet as well:</p>
 <d-code block language="python">
 def example_gelu():
     from torch.nn.functional import gelu
@@ -495,9 +557,9 @@ def example_gelu():
     torch.testing.assert_close(y_row_1, y_row_2, rtol=1e-5, atol=1e-5)
 </d-code>
-<p>To avoid a synchronization step directly after the first MLP, we’ll thus start with Column Parallel and be able to directly perform parallel GELU.</p>
-<p>Now, what about the second MLP? Should it be column or row parallel? Let’s draft both options:</p>
 <ul>
     <li>Column Parallel followed by Column Parallel</li>
     <img src="assets/images/image%2013.png" alt="Column Parallel Schema 1">
@@ -505,9 +567,9 @@ def example_gelu():
     <img src="assets/images/image%2014.png" alt="Column Parallel Schema 2">
 </ul>
-<p>We see that the “Column Parallel followed by Row Parallel” schema only involves two communications instead of four. It’s thus the most efficient schema in terms of communications.</p>
-<p>Let’s take a quick look at the backward pass:</p>
 <img src="assets/images/image%2015.png" alt="Backward Pass 1">
 <img src="assets/images/image%2016.png" alt="Backward Pass 2">
@@ -586,7 +648,7 @@ if __name__ == "__main__":
     example_column_row_linear()
 </d-code>
-<p>Now that we’ve found the most efficient schema for the Feedforward part of the transformer, let’s take a look at the multi-head attention block (MHA).</p>
 <p>We can generally follow a similar approach where the Q, K, V will be split in a Column Parallel fashion and the output projection will be split along the Row dimension.</p>
@@ -604,7 +666,7 @@ if __name__ == "__main__":
 <p>Could we push this approach further?</p>
-<p>Sequence parallelism applies this same idea to other parts of our model. We’ve applied tensor parallelism to two main parts in our models where combination of MLP allowed to naturally split the weights along major axis.</p>
 <p>The rest of the model mostly comprises layer norms, dropout and various summation of residuals, these contribute little to the computation but come with rather large forward activations to store.</p>
@@ -640,14 +702,14 @@ if __name__ == "__main__":
 <p>In the transformer model, tokens have no inherent information about their positional information. For these reasons, we need to use a positional encoding function.</p>
-<p>Assuming that in the multi-head attention layer, <em>q_m</em> is the “position-aware” query vector corresponding to a token at position <em>m</em>, <em>k_n</em> the “position-aware” key vector corresponding to the token at position <em>n</em> and <em>f</em> is our position embedding function, we would like our position vector to be a function of the input vectors and absolute positions like this:</p>
 <d-math>
 q_m = f(q,m)
 k_n = f(k,n)
 </d-math>
-<p>We may also want the positional encoding to model relative positional information between two input tokens. Relative positions help the model to operate across longer context spans and even context lengths not seen during training. The attention operation is generally a dot product operation between “position-aware” vectors <em>q</em> and <em>k</em>, so for a positional encoding that contains relative positional information, we’ll want to have:</p>
 <d-math>
 <q_m, k_n> = g(q, k, m-n)
@@ -655,7 +717,7 @@ k_n = f(k,n)
 <p>In other words, we want the result of <em>⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩</em> to depend on the values of <em>q</em> and <em>k</em> themselves, as well as their relative position <em>m − n</em>, but not <em>m</em> and <em>n</em>. This way, the model can focus on the relative difference between two tokens rather than their absolute positions.</p>
-<p>Let’s show that the RoPE positional embedding formulation satisfies the above formula.</p>
 <p><strong>Rotation matrix</strong></p>
@@ -695,9 +757,9 @@ R(θ) =
 <p><strong>Implementation</strong></p>
-<p>In our case, our internal vectors (the activations in our model) have much more than two elements. Let’s pair elements to get 2D vectors and apply the 2D rotation operation on these pairs.</p>
-<p>There are combinatorially many ways we can pair elements but generally two options are the most popular for implementing RoPE: we call them the <em>interleaved</em> and <em>non-interleaved</em> versions. (It’s still rather unfortunate to have two popular options)</p>
 <ol>
     <li>In the interleaved version, we pair consecutive elements <em>(x<sub>0</sub>, x<sub>1</sub>),(x<sub>2</sub>,x<sub>3</sub>),…</em> before applying the rotation matrix:</li>
@@ -860,13 +922,13 @@ x_{d-1}\cos mθ_{d/2-1} + x_{d-1}\sin mθ_{d/2-1}  \\
 <h2>References</h2>
 <ul>
-    <li>Harm’s posts:
         <ul>
             <li><a href="https://www.harmdevries.com/post/context-length/">https://www.harmdevries.com/post/context-length/</a></li>
             <li><a href="https://www.harmdevries.com/post/model-size-vs-compute-overhead/">https://www.harmdevries.com/post/model-size-vs-compute-overhead/</a></li>
         </ul>
     </li>
-    <li>Stas’ guides:
         <ul>
             <li><a href="https://github.com/stas00/ml-engineering">https://github.com/stas00/ml-engineering</a></li>
             <li><a href="https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md">https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md</a></li>

     <title>FineWeb: decanting the web for the finest text data at scale</title>
     <link rel="stylesheet" href="style.css">
     <style>
+        #controls {
+            display: grid;
+            grid-template-columns: auto 1fr auto;
+            gap: 5px;
+            align-items: center;
+            max-width: 600px;
+            margin-bottom: 20px;
+        }
+        #controls label {
+            text-align: right;
+        }
+        #controls input[type="range"] {
+            width: 100%;
+        }
+        #controls input[type="number"] {
+            width: 60px;
+        }
+        #controls .row {
+            display: contents;
+        }
         #graph svg {
             font-family: sans-serif;
         }
         <aside>We are extremely thankful to the whole <a href="https://distill.pub/">distill.pub</a> team for creating the template on which we based this blog post.</aside>
+        <div id="graph" style="position: relative; width: 700px; height: 500px;"></div>
+        <div id="controls">
+            <div class="row">
+                <label for="a">Attention Heads (a):</label>
+                <input type="range" id="a" name="a" min="1" max="128" value="8">
+                <input type="number" id="a_input" value="8" min="1" max="128">
+            </div>
+            <div class="row">
+                <label for="b">Micro Batch Size (b):</label>
+                <input type="range" id="b" name="b" min="1" max="53248" value="32">
+                <input type="number" id="b_input" value="32" min="1" max="53248">
+            </div>
+            <div class="row">
+                <label for="h">Hidden Dimension Size (h):</label>
+                <input type="range" id="h" name="h" min="1" max="16384" value="512">
+                <input type="number" id="h_input" value="512" min="128" max="16384">
+            </div>
+            <div class="row">
+                <label for="h_ff">Feedforward Dimension Size (h_ff):</label>
+                <input type="range" id="h_ff" name="h_ff" min="1" max="65536" value="2048">
+                <input type="number" id="h_ff_input" value="2048" min="512" max="65536">
+            </div>
+            <div class="row">
+                <label for="L">Number of Layers (L):</label>
+                <input type="range" id="L" name="L" min="1" max="126" value="12">
+                <input type="number" id="L_input" value="12" min="1" max="126">
+            </div>
+            <div class="row">
+                <label for="s">Sequence Length (s):</label>
+                <input type="range" id="s" name="s" min="1" max="128000" value="128">
+                <input type="number" id="s_input" value="128" min="64" max="128000">
+            </div>
+            <div class="row">
+                <label for="v">Vocabulary Size (v):</label>
+                <input type="range" id="v" name="v" min="1000" max="100000" value="30522">
+                <input type="number" id="v_input" value="30522" min="1000" max="100000">
+            </div>
+            <div class="row">
+                <label for="k">Optimizer Parameters (k):</label>
+                <input type="range" id="k" name="k" min="1" max="16" value="8">
+                <input type="number" id="k_input" value="8" min="1" max="16">
+            </div>
+            <div class="row">
+                <label for="tp">Tensor Model Parallelism (t):</label>
+                <input type="range" id="tp" name="tp" min="1" max="16" value="8">
+                <input type="number" id="tp_input" value="8" min="1" max="16">
+            </div>
+            <div class="row">
+                <label for="dp">Data Model Parallelism (d):</label>
+                <input type="range" id="dp" name="dp" min="1" max="256" value="1">
+                <input type="number" id="dp_input" value="1" min="1" max="256">
+            </div>
+            <div class="row">
+                <label for="mixed">Mixed Precision:</label>
+                <input type="checkbox" id="mixed" name="mixed" checked>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="recomputation">Recomputation:</label>
+                <select id="recomputation" name="recomputation">
+                    <option value="none">None</option>
+                    <option value="selective">Selective</option>
+                    <option value="full">Full</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="zero">Zero:</label>
+                <select id="zero" name="zero">
+                    <option value="Optimizer">Optimizer</option>
+                    <option value="Gradients">Gradients</option>
+                    <option value="Parameters">Parameters</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="ff_activation">FF Activation:</label>
+                <select id="ff_activation" name="ff_activation">
+                    <option value="relu">ReLU</option>
+                    <option value="gelu">GELU</option>
+                    <option value="swiglu">SwiGLU</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
+            <div class="row">
+                <label for="presets">Presets:</label>
+                <select id="presets" name="presets">
+                    <option value="Llama 3 Tiny">Llama 3 Tiny</option>
+                    <option value="Llama 3 8B">Llama 3 8B</option>
+                    <option value="Llama 3 70B">Llama 3 70B</option>
+                    <option value="Llama 3 405B">Llama 3 405B</option>
+                </select>
+                <span></span> <!-- Empty span to maintain grid alignment -->
+            </div>
         </div>
     <p><strong>TLDR:</strong> This blog covers a discussion on processing and evaluating data quality at scale, the 🍷 FineWeb
         recipe (listing and explaining all of our design choices), and the process followed to create its 📚 FineWeb-Edu subset.</p>
 <p>Most frameworks these days use FlashAttention (TODO: see later) which makes the attention computation less memory intensive through kernel fusion, thus most trainings use the <code>full</code> settings.</p>
+<p>We can save some GPU memory with activation recomputation but this only delays by a bit the next bottleneck: as hinted earlier for LLM training there is usually a sweet spot for the GBST and we need to work out the training configuration backward from there. However, you can't choose MBS to be an arbitrary large number on your GPU; at some point you will run out of GPU memory again since you need to store at least some of the activations in memory.</p>
 <p>There is a useful trick to compensate for that: <strong>gradient accumulation</strong> (<em>GradAcc</em>). With gradient accumulation we will split our batch in micro-batch, do forward and backward passes repeatedly on each micro-batch, compute the gradients, and, as the name suggests, sum the gradients step by step before doing a final optimizer step.</p>
+<p>We call the <code>micro batch size</code> (MBS) the batch size for each forward pass on a single node (the number of samples flowing through the model in one forward pass). We'll refer to the overall batch size between each optimizer step as the <code>global batch size</code> (GBS). If we do one optimizer step each 8 forward/backward pass, the <code>global batch size</code> will be 8 times the <code>micro batch size</code>.</p>
+<p>What we now call <code>global batch size</code> thus corresponds to what we've called up to now just <code>batch size</code> for simplicity (we now make the terms more precise to avoid ambiguity).</p>
 <p>With gradient accumulation the global batch size can be computed as follows:</p>
 <p>This is actually a bummer since the forward/backward passes for each micro-batch could actually totally be run in parallel. They are independent from each other and the only changing parameter are the input samples.</p>
+<p>Here comes data parallelism to solve exactly this problem! Let's take a look, you say? Okay sure!</p>
 <h3>Data Parallelism</h3>
 GBS=MBS * GradAcc * DP
 </d-math>
+<p>This means that we can reduce the number of gradient accumulation steps in favor of data parallel processes which speeds up training. In practice, people will tend to max out the number of data parallel nodes (the DP above) as much as possible as it's inherently parallel versus the sequential Gradient Accumulation. Gradient accumulation is then added only to achieve a target batch size if DP alone is not sufficient. One exception to that is pipeline parallelism which we'll discuss later.</p>
 <img src="assets/images/IMG_A95961668B3F-1.jpeg" alt="Data Parallelism">
 <p>If the gradient accumulation ratio is lower than one, i.e. you have too many GPUs (!), you can either choose to not use all your GPUs or test if a lower MBS will speed up training. In these cases, you may want to prioritize throughput over the individual GPU utilization, you can then choose DP first and use a smaller MBS than possible in order to speed up training.</p>
+<p>Time to take a concrete example: We want to train a model with a GBS of 4M tokens and a sequence length of 4k. This means our batch size will be 1024 samples (we pick powers of two). We observe that a single of our GPU can fit MBS=2 in memory and we have 128 GPUs available for training. This means with 4 gradient accumulation steps we'll achieve our goal of 1024 samples or 4M tokens per training step. Now what if we suddenly have 1024 GPUs available? We can achieve the same GBS and thus identical training by setting both MBS and gradient accumulation to 1 speeding up training significantly.</p>
 <p>[EXPERIMENTS WHERE WE INCREASE DP AND SHOW THROUGHPUT FOR SEVERAL MODELS]</p>
+<p>We've explored data parallelism, a simple strategy to scale training across more GPUs and gives consistent speed improvements. The keen reader might have noticed however that it rests on the assumption that we can fit at least one input sample forward pass (<em>MBS=1</em>) into our GPU memory. This is not always the case! In particular for larger models which often don't fit into a single GPU anymore even with activation recomputations activated.</p>
+<p>In such case, we need to shard the model across devices! We'll now study two complementary sharding methods, tensor and pipeline parallelism which are doing that. Let's start by the simplest, tensor parallelism!</p>
 <h3>Tensor Parallelism</h3>
+<p>So you've exhausted all the previous textbook tricks to try to fit your model on a single GPU but it still doesn't fit? Let's try to distribute this model across several GPUs. Unlike DP we will not simply duplicate the model but various parts of the model instance will be living on various GPUs.</p>
 <p>If we take a look at a typical matrix multiplication (the core of a neural network), we can get an idea about how we could split the model:</p>
 <p>If you rather like code, note that we can prove this with the following snippet as well:</p>
 <d-code block language="python">
+```
+</region_of_file_to_rewritten_file>
 def example_gelu():
     from torch.nn.functional import gelu
     torch.testing.assert_close(y_row_1, y_row_2, rtol=1e-5, atol=1e-5)
 </d-code>
+<p>To avoid a synchronization step directly after the first MLP, we'll thus start with Column Parallel and be able to directly perform parallel GELU.</p>
+<p>Now, what about the second MLP? Should it be column or row parallel? Let's draft both options:</p>
 <ul>
     <li>Column Parallel followed by Column Parallel</li>
     <img src="assets/images/image%2013.png" alt="Column Parallel Schema 1">
     <img src="assets/images/image%2014.png" alt="Column Parallel Schema 2">
 </ul>
+<p>We see that the "Column Parallel followed by Row Parallel" schema only involves two communications instead of four. It's thus the most efficient schema in terms of communications.</p>
+<p>Let's take a quick look at the backward pass:</p>
 <img src="assets/images/image%2015.png" alt="Backward Pass 1">
 <img src="assets/images/image%2016.png" alt="Backward Pass 2">
     example_column_row_linear()
 </d-code>
+<p>Now that we've found the most efficient schema for the Feedforward part of the transformer, let's take a look at the multi-head attention block (MHA).</p>
 <p>We can generally follow a similar approach where the Q, K, V will be split in a Column Parallel fashion and the output projection will be split along the Row dimension.</p>
 <p>Could we push this approach further?</p>
+<p>Sequence parallelism applies this same idea to other parts of our model. We've applied tensor parallelism to two main parts in our models where combination of MLP allowed to naturally split the weights along major axis.</p>
 <p>The rest of the model mostly comprises layer norms, dropout and various summation of residuals, these contribute little to the computation but come with rather large forward activations to store.</p>
 <p>In the transformer model, tokens have no inherent information about their positional information. For these reasons, we need to use a positional encoding function.</p>
+<p>Assuming that in the multi-head attention layer, <em>q_m</em> is the "position-aware" query vector corresponding to a token at position <em>m</em>, <em>k_n</em> the "position-aware" key vector corresponding to the token at position <em>n</em> and <em>f</em> is our position embedding function, we would like our position vector to be a function of the input vectors and absolute positions like this:</p>
 <d-math>
 q_m = f(q,m)
 k_n = f(k,n)
 </d-math>
+<p>We may also want the positional encoding to model relative positional information between two input tokens. Relative positions help the model to operate across longer context spans and even context lengths not seen during training. The attention operation is generally a dot product operation between "position-aware" vectors <em>q</em> and <em>k</em>, so for a positional encoding that contains relative positional information, we'll want to have:</p>
 <d-math>
 <q_m, k_n> = g(q, k, m-n)
 <p>In other words, we want the result of <em>⟨ 𝑞_𝑚 , 𝑘_𝑛 ⟩</em> to depend on the values of <em>q</em> and <em>k</em> themselves, as well as their relative position <em>m − n</em>, but not <em>m</em> and <em>n</em>. This way, the model can focus on the relative difference between two tokens rather than their absolute positions.</p>
+<p>Let's show that the RoPE positional embedding formulation satisfies the above formula.</p>
 <p><strong>Rotation matrix</strong></p>
 <p><strong>Implementation</strong></p>
+<p>In our case, our internal vectors (the activations in our model) have much more than two elements. Let's pair elements to get 2D vectors and apply the 2D rotation operation on these pairs.</p>
+<p>There are combinatorially many ways we can pair elements but generally two options are the most popular for implementing RoPE: we call them the <em>interleaved</em> and <em>non-interleaved</em> versions. (It's still rather unfortunate to have two popular options)</p>
 <ol>
     <li>In the interleaved version, we pair consecutive elements <em>(x<sub>0</sub>, x<sub>1</sub>),(x<sub>2</sub>,x<sub>3</sub>),…</em> before applying the rotation matrix:</li>
 <h2>References</h2>
 <ul>
+    <li>Harm's posts:
         <ul>
             <li><a href="https://www.harmdevries.com/post/context-length/">https://www.harmdevries.com/post/context-length/</a></li>
             <li><a href="https://www.harmdevries.com/post/model-size-vs-compute-overhead/">https://www.harmdevries.com/post/model-size-vs-compute-overhead/</a></li>
         </ul>
     </li>
+    <li>Stas' guides:
         <ul>
             <li><a href="https://github.com/stas00/ml-engineering">https://github.com/stas00/ml-engineering</a></li>
             <li><a href="https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md">https://github.com/bigscience-workshop/bigscience/blob/master/train/tr11-176B-ml/chronicles.md</a></li>

src/memory.js CHANGED Viewed

@@ -8,27 +8,35 @@ export function activationMemory(
     L, // number of layers
     s, // sequence length
     v, // vocab size
     mixed = true,
     recomputation = "none",
     ff_activation = "relu"
 ) {
-    console.log('activationMemory called with:', { a, b, h, h_ff, L, s, v, mixed, recomputation, ff_activation });
     // https://arxiv.org/pdf/2205.05198
     const bytesPerValue = mixed ? 2 : 4;
-    const oneLayerAttention = s * b * h * (bytesPerValue * 5 + 1) + ((2 * bytesPerValue + 1) * a * s * s * b); // eq (2)
     let oneLayerFeedforward;
     if (ff_activation === "relu") {
-        oneLayerFeedforward = (s * b * h * bytesPerValue + (s * b * h_ff * bytesPerValue) // inputs of 1st/2nd linear layers
             + s * b * h);  // dropout
     } else if (ff_activation === "gelu") {
-        oneLayerFeedforward = (s * b * h * bytesPerValue + (s * b * h_ff * bytesPerValue) // inputs of 1st/2nd linear layers
-            + s * b * h_ff * bytesPerValue // inputs of activation function (not really necessary for Relu)
             + s * b * h);  // dropout
     } else if (ff_activation === "swiglu") {
-        oneLayerFeedforward = (s * b * h * bytesPerValue + (s * b * h_ff * bytesPerValue) // inputs of input/output linear layers
-            + s * b * h_ff * bytesPerValue * 3 // inputs of activation function
             + s * b * h);  // dropout (note that dropout is lower-precision - boolean)
     }
@@ -41,41 +49,56 @@ export function activationMemory(
     let oneLayer;
-    if (recomputation === "none") {
-        oneLayer = oneLayerAttention + oneLayerFeedforward + 2 * layerNorm; // eq (2)
-    } else if (recomputation === "selective") {
-        oneLayer = s * b * h * 34; // eq (6)
     } else if (recomputation === "full") {
-        oneLayer = s * b * h * 2;
-    } else {
         throw new Error("Invalid recomputation value");
     }
-    const data = {
-        name: "activationMemory",
-        children: [
-            ...Array.from({ length: L }, (_, index) => ({
-                name: `Layer ${index + 1}`,
-                children: [
-                    { name: 'Attention', value: oneLayerAttention },
-                    { name: 'Feedforward', value: oneLayerFeedforward },
-                    { name: 'LayerNorm', value: 2 * layerNorm },
-                ]
-            })),
-            { name: 'Dropout', value: inputDropout },
-            { name: 'LayerNorm', value: outputLayerNorm },
-            { name: 'Projection', value: outputLayerProjection },
-            { name: 'Cross Entropy', value: outputCrossEntropy }
-        ]
-    };
-    const total = L * oneLayer + inputDropout + outputLayerNorm + outputLayerProjection + outputCrossEntropy;
     return data;
 }
-export function paramGradsOpt(h, L, s, v, k = 8, mixed = true) {
-    console.log('paramGradsOpt called with:', { h, L, s, v, k, mixed });
     const emb = h * (v + s);
     const oneLayer = 12 * h ** 2 + 13 * h;
     const other = 2 * h;
@@ -87,9 +110,16 @@ export function paramGradsOpt(h, L, s, v, k = 8, mixed = true) {
     }
     const bytesPerParameter = mixed ? 2 : 4;
-    const result = [bytesPerParameter * n, bytesPerParameter * n, k * n];
-    console.log('paramGradsOpt result:', result);
-    return result;
 }
 export function updateGraph() {
@@ -101,15 +131,18 @@ export function updateGraph() {
     const L = +document.getElementById('L').value;
     const s = +document.getElementById('s').value;
     const v = +document.getElementById('v').value;
     const mixed = document.getElementById('mixed').checked;
     const recomputation = document.getElementById('recomputation').value;
     const ff_activation = document.getElementById('ff_activation').value;
-    console.log('Slider values:', { a, b, h, h_ff, L, s, v, mixed, recomputation, ff_activation });
-    const fixedSize100GB = 100 * 1024 * 1024 * 1024; // 100GB in bytes
-    const activationMemoryData = activationMemory(a, b, h, h_ff, L, s, v, mixed, recomputation, ff_activation);
-    const paramGradsOptValue = paramGradsOpt(h, L, s, v)[0];
     const data = {
         name: "root",
@@ -119,7 +152,7 @@ export function updateGraph() {
                 value: 0,
                 children: [
                     activationMemoryData,
-                    { name: 'paramGradsOpt', value: paramGradsOptValue }
                 ]
             }
         ]
@@ -147,9 +180,11 @@ export function updateGraph() {
         .sum(d => d.value);
         // .sort((a, b) => b.value - a.value);
-    if (root.children[0].value < fixedSize100GB) {
-        root.children[0].value = fixedSize100GB;
-    }
     console.log('Treemap root:', root);
@@ -157,16 +192,20 @@ export function updateGraph() {
     const color = d => {
         switch(d.data.name) {
-            case 'paramGradsOpt': return '#4e79a7';  // Blue
             case 'activationMemory': return '#f28e2c';  // Orange
             case 'fixed100GB': return '#59a14f';  // Green
             case 'Attention': return '#e15759';  // Red
-            case 'Feedforward': return '#f28e2c';  // Orange
-            case 'LayerNorm': return '#9b59b6';  // Purple
-            case 'Dropout': return '#e15759';  // Red
-            case 'Projection': return '#f28e2c';  // Orange
-            case 'Cross Entropy': return '#e15759';  // Red
-            default: return '#59a14f';  // Red (for unexpected cases)
         }
     };
@@ -178,7 +217,7 @@ export function updateGraph() {
     cell.append("rect")
         .attr("width", d => d.x1 - d.x0)
         .attr("height", d => d.y1 - d.y0)
-        .attr("fill", d => d.depth === 1 ? "none" : color(d))
         .attr("stroke", d => d.depth === 1 ? color(d) : "none")
         .attr("stroke-width", 2);
@@ -196,8 +235,7 @@ export function updateGraph() {
             const name = d.data.name;
             const value = formatBytes(d.value);
-            if (d.depth === 1) {
-                // Parent node (fixed100GB)
                 node.attr("transform", `translate(${padding},${fontSize + padding})`)
                     .attr("font-weight", "bold")
                     .text(`${name}: ${value}`);
@@ -235,8 +273,8 @@ export function updateGraph() {
         .attr("x", 0)
         .attr("width", 19)
         .attr("height", 19)
-        .attr("fill", d => d.data.name === 'fixed100GB' ? 'none' : color(d))
-        .attr("stroke", d => d.data.name === 'fixed100GB' ? color(d) : 'none')
         .attr("stroke-width", 2);
     legend.append("text")
@@ -256,10 +294,10 @@ function formatBytes(bytes) {
 }
 const presets = {
-    "Tiny": { a: 16, b: 3, h: 1024, h_ff: 4096, L: 1, s: 7, v: 30522, mixed: true, recomputation: "none", ff_activation: "gelu" },
-    "8B": { a: 32, b: 32, h: 4096, h_ff: 16384, L: 32, s: 256, v: 30522, mixed: true, recomputation: "none", ff_activation: "swiglu" },
-    "70B": { a: 64, b: 32, h: 8192, h_ff: 32768, L: 80, s: 256, v: 30522, mixed: true, recomputation: "none", ff_activation: "swiglu" },
-    "405B": { a: 128, b: 32, h: 16384, h_ff: 65536, L: 126, s: 256, v: 30522, mixed: true, recomputation: "none", ff_activation: "swiglu" }
 };
 function setPresetValues(preset) {
@@ -306,41 +344,80 @@ function syncSliderAndInput(sliderId, inputId) {
 }
 export const init_memory_plot = function () {
-    console.log('DOM fully loaded and parsed');
-    const sliderIds = ['a', 'b', 'h', 'h_ff', 'L', 's', 'v'];  // Added 'v'
     sliderIds.forEach(id => {
-        syncSliderAndInput(id, `${id}_input`);
     });
     const recomputationSelect = document.getElementById('recomputation');
-    recomputationSelect.addEventListener('change', updateGraph);
     const ffActivationSelect = document.getElementById('ff_activation');
-    ffActivationSelect.addEventListener('change', updateGraph);
     const mixedCheckbox = document.getElementById('mixed');
-    mixedCheckbox.addEventListener('change', updateGraph);
     const presetSelect = document.getElementById('presets');
-    presetSelect.addEventListener('change', (event) => {
-        setPresetValues(event.target.value);
-    });
-    // Set max values for sliders based on the highest values in the presets
-    document.getElementById('a').max = 128;
-    document.getElementById('b').max = 53248;
-    document.getElementById('h').max = 16384;
-    document.getElementById('h_ff').max = 65536;
-    document.getElementById('L').max = 126;
-    document.getElementById('s').max = 128000;
-    document.getElementById('v').max = 100000;  // Set a reasonable max for vocabulary size
     console.log('Adding svg');
-    const svg = d3.select("#graph")
-        .append("svg")
-        .attr("width", 960)
-        .attr("height", 500);
     updateGraph();
 };

     L, // number of layers
     s, // sequence length
     v, // vocab size
+    tp = 1, // tensor model parallelism
     mixed = true,
     recomputation = "none",
     ff_activation = "relu"
 ) {
+    console.log('activationMemory called with:', { a, b, h, h_ff, L, s, v, tp, mixed, recomputation, ff_activation });
     // https://arxiv.org/pdf/2205.05198
     const bytesPerValue = mixed ? 2 : 4;
+    let oneLayerAttention;
+    if (recomputation === "none" || recomputation === "full") {
+        oneLayerAttention = s * b * h * (bytesPerValue * 4 / tp + bytesPerValue + 1) + ((2 * bytesPerValue + 1) * a * s * s * b); // eq (2)
+    } else if (recomputation === "selective") {
+        oneLayerAttention = s * b * h * (bytesPerValue * 4 / tp + bytesPerValue + 1); // table 2
+    } else {
+        throw new Error("Invalid recomputation value");
+    }
     let oneLayerFeedforward;
     if (ff_activation === "relu") {
+        oneLayerFeedforward = (s * b * h * bytesPerValue + (s * b * h_ff * bytesPerValue / tp) // inputs of 1st/2nd linear layers
             + s * b * h);  // dropout
     } else if (ff_activation === "gelu") {
+        oneLayerFeedforward = (s * b * h * bytesPerValue + (s * b * h_ff * bytesPerValue / tp) // inputs of 1st/2nd linear layers
+            + s * b * h_ff * bytesPerValue / tp // inputs of activation function (not really necessary for Relu)
             + s * b * h);  // dropout
     } else if (ff_activation === "swiglu") {
+        oneLayerFeedforward = (s * b * h * bytesPerValue + (s * b * h_ff * bytesPerValue / tp) // inputs of input/output linear layers
+            + s * b * h_ff * bytesPerValue * 3 / tp // inputs of activation function
             + s * b * h);  // dropout (note that dropout is lower-precision - boolean)
     }
     let oneLayer;
+    let data
+    if (recomputation === "none" || recomputation === "selective") {
+         data = {
+            name: "activationMemory",
+            children: [
+                ...Array.from({ length: L }, (_, index) => ({
+                    name: `Layer ${index + 1}`,
+                    children: [
+                        { name: 'Attention', value: oneLayerAttention },
+                        { name: 'Feedforward', value: oneLayerFeedforward },
+                        { name: 'LayerNorm', value: 2 * layerNorm },
+                    ]
+                })),
+                { name: 'Dropout', value: inputDropout },
+                { name: 'LayerNorm', value: outputLayerNorm },
+                { name: 'Projection', value: outputLayerProjection },
+                { name: 'Cross Entropy', value: outputCrossEntropy }
+            ]
+        };
     } else if (recomputation === "full") {
+         data = {
+            name: "activationMemory",
+            children: [
+                { name: 'LayerInput', value: s * b * h * bytesPerValue * L},
+                { name: 'Dropout', value: inputDropout },
+                { name: 'LayerNorm', value: outputLayerNorm },
+                { name: 'Projection', value: outputLayerProjection },
+                { name: 'Cross Entropy', value: outputCrossEntropy }
+            ]
+        };
+        } else {
         throw new Error("Invalid recomputation value");
     }
     return data;
 }
+export function paramGradsOpt(h, L, s, v, k = 8, dp = 1, zero = "Optimizer", mixed = true) {
+    // h, # hidden dimension size
+    // L, # number of layers
+    // s, # sequence length
+    // v, # vocab size
+    // k=8, # parameters for optimizer (Adam: 8 = 4 bytes moments + 4 bytes variance)
+    // dp=1, # data parallelism
+    // zero = "Optimizer", # zero data parallelism
+    // mixed=True # mixed precision training
+    console.log('paramGradsOpt called with:', { h, L, s, v, k, dp, zero, mixed });
     const emb = h * (v + s);
     const oneLayer = 12 * h ** 2 + 13 * h;
     const other = 2 * h;
     }
     const bytesPerParameter = mixed ? 2 : 4;
+    const data = {
+        name: "ParametersGradientOps",
+        children: [
+            { name: 'Parameters', value: zero === "Parameters" ? bytesPerParameter * n / dp : bytesPerParameter * n },
+            { name: 'Gradients', value: zero === "Gradients" ? bytesPerParameter * n / dp : bytesPerParameter * n },
+            { name: 'OptimizerAverages', value: zero === "Optimizer" ? k * n / dp : k * n }
+        ]
+    };
+    console.log('paramGradsOpt result:', data);
+    return data;
 }
 export function updateGraph() {
     const L = +document.getElementById('L').value;
     const s = +document.getElementById('s').value;
     const v = +document.getElementById('v').value;
+    const k = +document.getElementById('k').value;
+    const tp = +document.getElementById('tp').value;  // New: t parameter
+    const zero = document.getElementById('zero').value;
+    const dp = document.getElementById('dp').value;
     const mixed = document.getElementById('mixed').checked;
     const recomputation = document.getElementById('recomputation').value;
     const ff_activation = document.getElementById('ff_activation').value;
+    console.log('Slider values:', { a, b, h, h_ff, L, s, v, k, tp, zero, dp, mixed, recomputation, ff_activation });
+    const activationMemoryData = activationMemory(a, b, h, h_ff, L, s, v, tp, mixed, recomputation, ff_activation);
+    const paramGradsOptValue = paramGradsOpt(h, L, s, v, k, dp, zero, mixed);
     const data = {
         name: "root",
                 value: 0,
                 children: [
                     activationMemoryData,
+                    paramGradsOptValue
                 ]
             }
         ]
         .sum(d => d.value);
         // .sort((a, b) => b.value - a.value);
+        // const fixedSize100GB = 100 * 1024 * 1024 * 1024; // 100GB in bytes
+        // if (root.children[0].value < fixedSize100GB) {
+    //     root.value = fixedSize100GB;
+    //     root.children[0].value = fixedSize100GB;
+    // }
     console.log('Treemap root:', root);
     const color = d => {
         switch(d.data.name) {
+            case 'Parameters': return '#4e79a7';  // Blue
+            case 'Gradients': return '#f28e2c';  // Orange
+            case 'OptimizerAverages': return '#e15759';  // Green
             case 'activationMemory': return '#f28e2c';  // Orange
             case 'fixed100GB': return '#59a14f';  // Green
             case 'Attention': return '#e15759';  // Red
+            case 'Feedforward': return '#1f77b4';  // Light Blue
+            case 'LayerNorm': return '#ff7f0e';  // Dark Orange
+            case 'Dropout': return '#2ca02c';  // Dark Green
+            case 'Projection': return '#d62728';  // Dark Red
+            case 'Cross Entropy': return '#9467bd';  // Violet
+            case 'Total': return '#59a14f';  // Green
+            case 'root': return '#d3d3d3';  // Light Grey
+            default: return '#a0c4ff';  // Lighter Blue (for unexpected cases)
         }
     };
     cell.append("rect")
         .attr("width", d => d.x1 - d.x0)
         .attr("height", d => d.y1 - d.y0)
+        .attr("fill", d => color(d))
         .attr("stroke", d => d.depth === 1 ? color(d) : "none")
         .attr("stroke-width", 2);
             const name = d.data.name;
             const value = formatBytes(d.value);
+            if (d.depth === 1 || d.depth === 2) {
                 node.attr("transform", `translate(${padding},${fontSize + padding})`)
                     .attr("font-weight", "bold")
                     .text(`${name}: ${value}`);
         .attr("x", 0)
         .attr("width", 19)
         .attr("height", 19)
+        .attr("fill", d => color(d))
+        .attr("stroke", 'grey')
         .attr("stroke-width", 2);
     legend.append("text")
 }
 const presets = {
+    "Llama 3 Tiny": { a: 16, b: 3, h: 1024, h_ff: 4096, L: 1, s: 7, v: 30522, k: 8, tp: 1, zero: "Optimizer", dp: 1, mixed: true, recomputation: "none", ff_activation: "gelu" },
+    "Llama 3 8B": { a: 32, b: 32, h: 4096, h_ff: 16384, L: 32, s: 256, v: 30522, k: 8, tp: 1, zero: "Optimizer", dp: 1, mixed: true, recomputation: "none", ff_activation: "swiglu" },
+    "Llama 3 70B": { a: 64, b: 32, h: 8192, h_ff: 32768, L: 80, s: 256, v: 30522, k: 8, tp: 1, zero: "Optimizer", dp: 1, mixed: true, recomputation: "none", ff_activation: "swiglu" },
+    "Llama 3 405B": { a: 128, b: 32, h: 16384, h_ff: 65536, L: 126, s: 256, v: 30522, k: 8, t: 1, mixed: true, recomputation: "none", ff_activation: "swiglu" }
 };
 function setPresetValues(preset) {
 }
 export const init_memory_plot = function () {
+    console.log('Initializing memory plot');
+    const sliderIds = ['a', 'b', 'h', 'h_ff', 'L', 's', 'v', 'k', 'tp', 'dp'];
     sliderIds.forEach(id => {
+        const slider = document.getElementById(id);
+        const input = document.getElementById(`${id}_input`);
+        if (slider && input) {
+            syncSliderAndInput(id, `${id}_input`);
+        } else {
+            console.warn(`Elements for ${id} not found`);
+        }
     });
     const recomputationSelect = document.getElementById('recomputation');
+    if (recomputationSelect) {
+        recomputationSelect.addEventListener('change', updateGraph);
+    } else {
+        console.warn('Recomputation select not found');
+    }
     const ffActivationSelect = document.getElementById('ff_activation');
+    if (ffActivationSelect) {
+        ffActivationSelect.addEventListener('change', updateGraph);
+    } else {
+        console.warn('FF Activation select not found');
+    }
     const mixedCheckbox = document.getElementById('mixed');
+    if (mixedCheckbox) {
+        mixedCheckbox.addEventListener('change', updateGraph);
+    } else {
+        console.warn('Mixed checkbox not found');
+    }
     const presetSelect = document.getElementById('presets');
+    if (presetSelect) {
+        presetSelect.addEventListener('change', (event) => {
+            setPresetValues(event.target.value);
+        });
+    } else {
+        console.warn('Preset select not found');
+    }
+    // Set max values for sliders
+    sliderIds.forEach(id => {
+        const slider = document.getElementById(id);
+        if (slider) {
+            switch(id) {
+                case 'a': slider.max = '128'; break;
+                case 'b': slider.max = '53248'; break;
+                case 'h': slider.max = '16384'; break;
+                case 'h_ff': slider.max = '65536'; break;
+                case 'L': slider.max = '126'; break;
+                case 's': slider.max = '128000'; break;
+                case 'v': slider.max = '100000'; break;
+                case 'k': slider.max = '16'; break;
+                case 'tp': slider.max = '16'; break;
+                case 'dp': slider.max = '256'; break;
+            }
+        } else {
+            console.warn(`Slider ${id} not found`);
+        }
+    });
     console.log('Adding svg');
+    const graphContainer = document.getElementById('graph');
+    if (graphContainer) {
+        const svg = d3.select("#graph")
+            .append("svg")
+            .attr("width", 960)
+            .attr("height", 500);
+    } else {
+        console.warn('Graph container not found');
+    }
     updateGraph();
 };