Jacoguzo/tutorials editable code

docs/_static/css/custom.css

@@ -209,3 +209,35 @@ html[data-theme="dark"] .button-primary:visited:hover {
     padding-inline: 50px;
     padding-bottom: 24px;
 }
+
+/* Inline inputs for editable commands */
+.inline-input {
+    background-color: rgba(255, 255, 255, 0.1);
+    color: inherit;
+    border: 1px solid #555;
+    border-radius: 3px;
+    padding: 2px 6px;
+    font-family: inherit;
+    font-size: inherit;
+    display: inline-block;
+    vertical-align: middle;
+    margin: 0 2px;
+    box-sizing: content-box;
+}
+
+.inline-input:focus {
+    outline: none;
+    border-color: #1A73E8;
+    background-color: rgba(255, 255, 255, 0.2);
+}
+
+html[data-theme="light"] .inline-input {
+    background-color: rgba(0, 0, 0, 0.05);
+    border-color: #ccc;
+    color: #333;
+}
+
+html[data-theme="light"] .inline-input:focus {
+    background-color: rgba(0, 0, 0, 0.1);
+    border-color: #1A73E8;
+}

docs/_static/js/editable_commands.js

@@ -0,0 +1,129 @@
+/**
+ * Handles inline editable commands in documentation.
+ * Replaces placeholders in code blocks with inline input fields.
+ */
+document.addEventListener('DOMContentLoaded', () => {
+  const codeBlocks = document.querySelectorAll('div.highlight-sh pre, div.highlight-bash pre, div.highlight-default pre');
+
+  codeBlocks.forEach(block => {
+
+    const originalHTML = block.innerHTML;
+
+    const placeholders = [
+      "<your virtual env name>",
+      "<model name>",
+      "<tokenizer path>",
+      "<Hugging Face access token>",
+      "<output directory to store run logs>",
+      "<name for this run>",
+      "<number of fine-tuning steps to run>",
+      "<batch size per device>",
+      "<Hugging Face dataset name>",
+      "<data split for train>",
+      "<data columns to train on>",
+      "<gcs path for MaxText checkpoint>",
+      "<Google Cloud Project ID>",
+      "<Name of GKE Cluster>",
+      "<GKE Cluster Zone>",
+      "<Name of Workload>",
+      "<TPU Type>",
+      "<GCS Path for Output/Logs>",
+      "<Fine-Tuning Steps>",
+      "<Hugging Face Access Token>",
+      "<Model Name>",
+      "<Model Tokenizer>",
+      "<Hugging Face Dataset Name>",
+      "<Data Split for Train>",
+      "<Data Columns to Train on>",
+      "<cluster name>",
+      "<GCP project ID>",
+      "<zone name>",
+      "<path/to/gcr.io>",
+      "<number of slices>",
+      "<Flag to use zarr3>",
+      "<Flag to use ocdbt>",
+      "<Hugging Face Model>",
+      "<MaxText Model>",
+      "<Tokenizer>",
+      "<Name for this run>"
+    ];
+
+    let newHTML = originalHTML;
+
+    placeholders.forEach(placeholder => {
+      // 1. create robust regex for this placeholder
+      // escape chars
+      const escapeRegex = (string) => string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+
+      const htmlEscapedKey = placeholder
+        .replace(/&/g, '&amp;')
+        .replace(/</g, '&lt;')
+        .replace(/>/g, '&gt;');
+
+      let pattern = '';
+      for (let i = 0; i < htmlEscapedKey.length; i++) {
+        const char = htmlEscapedKey[i];
+        pattern += escapeRegex(char) + '(?:<[^>]+>)*';
+      }
+
+      const regex = new RegExp(pattern, 'g');
+
+      // Replace with an input element
+      // We use the original placeholder text as placeholder for the input
+      const inputHTML = `<input class="inline-input" placeholder="${placeholder}" style="width: ${placeholder.length + 2}ch;" />`;
+
+      newHTML = newHTML.replace(regex, inputHTML);
+    });
+
+    if (newHTML !== originalHTML) {
+      block.innerHTML = newHTML;
+    }
+  });
+
+  // Add event listeners to newly created inputs to auto-resize
+  document.querySelectorAll('.inline-input').forEach(input => {
+    input.addEventListener('input', function () {
+      this.style.width = Math.max(this.value.length, this.placeholder.length) + 2 + 'ch';
+    });
+  });
+
+  /**
+   * Intercept copy button clicks to include user input values.
+   * Runs in capture phase to precede sphinx-copybutton's listener.
+   */
+  document.addEventListener('click', (event) => {
+    // Check if the clicked element is a copy button or inside one
+    const button = event.target.closest('.copybtn');
+    if (!button) return;
+
+    // Find the associated code block
+    // Sphinx-copybutton places the button inside .highlight usually
+    const highlightDiv = button.closest('.highlight');
+    if (!highlightDiv) return;
+
+    const inputs = highlightDiv.querySelectorAll('input.inline-input');
+    if (inputs.length === 0) return;
+
+    const swaps = [];
+    inputs.forEach(input => {
+      // Create a temporary span with the input's current value
+      const span = document.createElement('span');
+      // If value is empty, fallback to placeholder to match original text behavior
+      const val = input.value;
+      span.textContent = val ? val : input.placeholder;
+
+      // Mimic input appearance slightly if needed, but plain text is what we want copied
+      span.style.color = val ? 'inherit' : 'gray';
+
+      input.replaceWith(span);
+      swaps.push({ input, span });
+    });
+
+    // Revert immediately after the current event loop
+    setTimeout(() => {
+      swaps.forEach(({ input, span }) => {
+        span.replaceWith(input);
+      });
+    }, 0);
+  }, true);
+});

docs/conf.py

@@ -48,6 +48,7 @@
 html_theme = "sphinx_book_theme"
 html_static_path = ["_static"]
 html_css_files = ["css/custom.css"]
+html_js_files = ["js/editable_commands.js"]
 html_logo = "_static/maxtext.png"
 
 # -- Options for myst ----------------------------------------------

docs/tutorials/posttraining/rl.md

@@ -69,9 +69,9 @@ Setup following environment variables before running GRPO/GSPO:
 
 ```bash
 # -- Model configuration --
-export HF_MODEL='llama3.1-8b-Instruct'
-export MODEL='llama3.1-8b'
-export TOKENIZER='meta-llama/Llama-3.1-8B-Instruct'
+export HF_MODEL=<Hugging Face Model> # e.g. 'llama3.1-8b-Instruct'
+export MODEL=<MaxText Model> # e.g. 'llama3.1-8b'
+export TOKENIZER=<Tokenizer> # e.g. 'meta-llama/Llama-3.1-8B-Instruct'
 export HF_TOKEN=<Hugging Face access token>
 
 # -- MaxText configuration --

docs/tutorials/posttraining/rl_on_multi_host.md

@@ -39,26 +39,36 @@ Setup following environment variables:
 
 ```bash
 # -- Model configuration --
-export HF_MODEL='llama3.1-70b-Instruct'
-export MODEL='llama3.1-70b'
-export TOKENIZER='meta-llama/Llama-3.1-70B-Instruct'
+export HF_MODEL=<Hugging Face Model> # e.g. 'llama3.1-70b-Instruct'
+export MODEL=<MaxText Model> # e.g. 'llama3.1-70b'
+export TOKENIZER=<Tokenizer> # e.g. 'meta-llama/Llama-3.1-70B-Instruct'
 export HF_TOKEN=<Hugging Face access token>
 
 # -- MaxText configuration --
 export BASE_OUTPUT_DIRECTORY=<output directory to store run logs> # e.g., gs://my-bucket/my-output-directory
-export RUN_NAME=llama-3-70b-grpo
+export RUN_NAME=<Name for this run> # e.g., llama-3-70b-grpo
 export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/0/items
 
 # -- Workload configuration --
 export WORKLOAD=${RUN_NAME}
-export TPU_TYPE='v5p-128'
+export TPU_TYPE=<TPU Type> # e.g., 'v5p-128'
 export TPU_CLUSTER=<cluster name>
 export PROJECT_ID=<GCP project ID>
 export ZONE=<zone name>
 ```
 
 ## Get your model checkpoint
 
+### Option 1: Using an existing MaxText checkpoint
+
+If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section.
+
+```bash
+export MAXTEXT_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
+```
+
+### Option 2: Converting from a Hugging Face checkpoint
+
 You can convert a Hugging Face checkpoint to MaxText format using the `src/MaxText/utils/ckpt_conversion/to_maxtext.py` script. This is useful if you have a pre-trained model from Hugging Face that you want to use with MaxText.
 
 First, ensure you have the necessary dependencies installed. Then, run the conversion script on a CPU machine. For large models, it is recommended to use the `--lazy_load_tensors` flag to reduce memory usage during conversion. \

docs/tutorials/posttraining/sft.md

@@ -24,6 +24,7 @@ We use [Tunix](https://github.com/google/tunix), a JAX-based library designed fo
 In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get started!
 
 ## Install dependencies
+
 ```sh
 # 1. Clone the repository
 git clone https://github.com/AI-Hypercomputer/maxtext.git

docs/tutorials/posttraining/sft_on_multi_host.md

@@ -73,7 +73,7 @@ export ZONE=<GKE Cluster Zone>
 # -- Workload Configuration --
 export WORKLOAD_NAME=<Name of Workload> # e.g., sft-$(date +%s)
 export TPU_TYPE=<TPU Type> # e.g., v6e-256
-export TPU_SLICE=1
+export TPU_SLICE=<number of slices>
 export DOCKER_IMAGE="gcr.io/${PROJECT}/${DOCKER_IMAGE_NAME}"
 
 # -- MaxText Configuration --

placeholders.txt

@@ -0,0 +1,31 @@
+<batch size per device>
+<cluster name>
+<data columns to train on>
+<Data Columns to Train on>
+<data split for train>
+<Data Split for Train>
+<Fine-Tuning Steps>
+<Flag to use ocdbt>
+<Flag to use zarr3>
+<GCP project ID>
+<gcs path for MaxText checkpoint>
+<GCS Path for Output/Logs>
+<GKE Cluster Zone>
+<Google Cloud Project ID>
+<Hugging Face access token>
+<Hugging Face Access Token>
+<Hugging Face dataset name>
+<Hugging Face Dataset Name>
+<model name>
+<Model Name>
+<Model Tokenizer>
+<name for this run>
+<Name of GKE Cluster>
+<Name of Workload>
+<number of fine-tuning steps to run>
+<output directory to store run logs>
+<path/to/gcr.io>
+<tokenizer path>
+<TPU Type>
+<your virtual env name>
+<zone name>