diff --git a/apps/Ollama_Gemma4_WebSearch/Modelfile b/apps/Ollama_Gemma4_WebSearch/Modelfile new file mode 100644 index 00000000..04da276f --- /dev/null +++ b/apps/Ollama_Gemma4_WebSearch/Modelfile @@ -0,0 +1,20 @@ +FROM gemma3:latest + +# Gemma 4 requires a very strict template for Tool Calling / Web Search to function. +# Default Ollama templates may strip the <|tool_call|> and <|tool_response|> tokens, +# which causes Gemma to "see nothing" when the search results are returned. + +TEMPLATE """{{ if .System }}system +{{ .System }} +{{ end }}{{ if .Prompt }}user +{{ if .Tools }}You have access to the following functions: + +{{ range .Tools }}{{ .Function }} +{{ end }} + +Use the `<|tool_call|>` and `<|tool_response|>` tags for web search operations.{{ end }} +{{ .Prompt }} +{{ end }}model +{{ if .Response }}{{ .Response }}{{ end }}""" + +PARAMETER temperature 0.1 diff --git a/apps/Ollama_Gemma4_WebSearch/README.md b/apps/Ollama_Gemma4_WebSearch/README.md new file mode 100644 index 00000000..7e8cb65a --- /dev/null +++ b/apps/Ollama_Gemma4_WebSearch/README.md @@ -0,0 +1,23 @@ +# Gemma 4: Fixing Ollama Web Search + +When using **Gemma 4** with the Ollama Web Search tool (or within wrappers like OpenWebUI), you might encounter an issue where the model attempts a search but then "sees nothing" or acts as if no query was submitted. + +This happens because Gemma 4 uses a strict structured token format (`<|tool_call|>` and `<|tool_response|>`) for tools. If Ollama's default prompt template doesn't explicitly instruct the model to use these tags or if it strips the returned results from the template, Gemma 4 will drop the context. + +## Solution + +We have provided a custom `Modelfile` that injects the required structured tags into the system template, ensuring web search payloads are correctly returned to Gemma. + +### Usage + +1. Build the custom model locally: +```bash +ollama create gemma4-search-fixed -f Modelfile +``` + +2. Run it with the web search tool: +```bash +ollama run gemma4-search-fixed +``` + +3. (If using OpenWebUI): Select `gemma4-search-fixed` as your default model and enable the Web Search toggle. The tool results will now be correctly interpreted by Gemma. diff --git a/apps/concurrent_spark/README.md b/apps/concurrent_spark/README.md new file mode 100644 index 00000000..2abd9d69 --- /dev/null +++ b/apps/concurrent_spark/README.md @@ -0,0 +1,29 @@ +# Gemma 4 Concurrent on NVIDIA DGX Spark Clusters + +This guide demonstrates how to adapt the macOS/local `apps/concurrent` orchestrator to run distributed concurrent Gemma 4 agents across a multi-node NVIDIA DGX Spark cluster. + +## Architectural Overview + +DGX clusters typically feature high-speed ring topologies (e.g., connected via 2x 200 Gbps CX-7 cables) and NVLink for ultra-fast inter-GPU and inter-node communication. + +To scale the "Specialist Agent" pattern: +1. **The Orchestrator (Spark Driver)**: Uses Gemma to plan the tasks. +2. **The Specialists (Spark Executors)**: We replace local `subprocess` terminals with Spark partitions using `mapInPandas`. +3. **The LLM Backend (vLLM / TensorRT-LLM)**: We run an inference engine on each DGX node that binds to the local GPUs via NCCL/NVLink. + +### Network Topology Consideration +Because DGX nodes are connected via CX-7 rings, we can maximize throughput by: +- **Data Parallelism**: Spawning one vLLM server per node, with PySpark mapping rows to `localhost:8000` on the executor. +- **Pipeline/Tensor Parallelism**: Utilizing Ray on Spark or Spark DL to shard a single massive Gemma 4 model across the DGX nodes using the CX-7 interconnects. + +## Quick Start +Run the PySpark adaptation provided in `spark_orchestrator.py`. + +```bash +# Submit to Spark cluster +spark-submit \ + --master spark://:7077 \ + --executor-memory 128G \ + --executor-cores 8 \ + spark_orchestrator.py --scenario translate --topic "Scaling Gemma 4 with Spark" +``` diff --git a/apps/concurrent_spark/spark_orchestrator.py b/apps/concurrent_spark/spark_orchestrator.py new file mode 100644 index 00000000..6c58619d --- /dev/null +++ b/apps/concurrent_spark/spark_orchestrator.py @@ -0,0 +1,100 @@ +""" +spark_orchestrator.py — PySpark implementation of the concurrent multi-agent demo. + +Instead of spawning macOS terminal windows, this script distributes +specialist tasks to PySpark executors, which query LLM inference engines +running on the DGX nodes. +""" + +import argparse +import pandas as pd +from pyspark.sql import SparkSession +from pyspark.sql.types import StructType, StructField, StringType +from typing import Iterator + +# Assuming demo modules are available in Python Path or deployed via Spark +# from demo.scenarios import get_scenario +# from demo.utils import stream_llm + +def process_agent_partition(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: + """ + Spark mapInPandas function. + Runs on the DGX Spark Executors. + """ + import requests # Standard REST calls to node-local LLM + + # In a real DGX setup, each node might run its own vLLM instance bound to the local GPUs. + # We query the local inference server. + LOCAL_LLM_URL = "http://localhost:8000/v1/chat/completions" + + for pdf in iterator: + results = [] + for _, row in pdf.iterrows(): + system_prompt = row['system_prompt'] + instruction = row['instruction'] + + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + messages.append({"role": "user", "content": instruction}) + + try: + response = requests.post(LOCAL_LLM_URL, json={ + "model": "gemma-4", + "messages": messages, + "max_tokens": 1024 + }) + # Fallback extraction + result_text = response.json()['choices'][0]['message']['content'] + except Exception as e: + result_text = f"Error: {e}" + + results.append(result_text) + + pdf['result'] = results + yield pdf + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--scenario", default="translate") + parser.add_argument("--topic", default="Gemma 4 distributed on DGX Spark") + args = parser.parse_args() + + spark = SparkSession.builder.appName("Gemma4-Concurrent-Spark").getOrCreate() + + # In a real execution, we would call plan_tasks(api_url, scenario, topic) + # For this example, we mock the planned tasks + planned_tasks = [ + {"name": f"Agent_{i}", "instruction": f"Process {args.topic} segment {i}", "system_prompt": "You are a specialist."} + for i in range(16) # 16 Concurrent agents + ] + + print(f"Distributed Orchestrator: Planning complete. {len(planned_tasks)} tasks generated.") + + # ─── Dispatch via Spark ─────────────────────────────────────── + df = spark.createDataFrame(planned_tasks) + + # Define output schema + schema = StructType([ + StructField("name", StringType(), True), + StructField("instruction", StringType(), True), + StructField("system_prompt", StringType(), True), + StructField("result", StringType(), True) + ]) + + # Run concurrent inference across DGX cluster + result_df = df.mapInPandas(process_agent_partition, schema=schema) + + # ─── Collect and Assemble ────────────────────────────────────── + results_list = result_df.collect() + + results = {row['name']: row['result'] for row in results_list} + print("All distributed tasks completed!") + + for name, res in results.items(): + print(f"{name}: {res[:50]}...") + + spark.stop() + +if __name__ == "__main__": + main() diff --git a/tutorials/Gemma_3_Keras_TPU_Parallelism.ipynb b/tutorials/Gemma_3_Keras_TPU_Parallelism.ipynb new file mode 100644 index 00000000..ecd75f57 --- /dev/null +++ b/tutorials/Gemma_3_Keras_TPU_Parallelism.ipynb @@ -0,0 +1,164 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "a34edc13da26" + }, + "source": [ + "# Distributed Inference with Gemma 3 on Kaggle TPU v5e-8\n", + "\n", + "This notebook demonstrates how to modernize your JAX/TPU workflows using **Gemma 3** and the **Keras 3 Distribution API**.\n", + "\n", + "By leveraging a Kaggle TPU v5e-8 (8-core mesh), we can perform true Data Parallelism for high-throughput batch inference across a 32k context window without relying on legacy Flax abstractions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fa6d0de2f9aa" + }, + "outputs": [], + "source": [ + "!pip install -U keras keras-nlp\n", + "!pip install -U jax jaxlib" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d79d4b93b350" + }, + "source": [ + "## 1. Environment Setup\n", + "\n", + "Set the backend to `jax` and configure Keras to allocate memory across the TPU mesh." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5240b523544d" + }, + "outputs": [], + "source": [ + "import os\n", + "\n", + "# Must be set before importing Keras\n", + "os.environ[\"KERAS_BACKEND\"] = \"jax\"\n", + "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"0.9\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f31870325f66" + }, + "source": [ + "## 2. Initialize Keras 3 Distribution API\n", + "\n", + "We use `keras.distribution.DataParallel` to automatically shard our input data across all 8 TPU cores. Keras will replicate the model weights on each device and independently process micro-batches." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "998fc6e4b92c" + }, + "outputs": [], + "source": [ + "import keras\n", + "import keras_nlp\n", + "import jax\n", + "\n", + "print(f\"Devices available: {jax.devices()}\")\n", + "\n", + "# Create an 8-core data parallel distribution mesh\n", + "devices = keras.distribution.list_devices()\n", + "data_parallel = keras.distribution.DataParallel(devices=devices)\n", + "\n", + "# Set the global distribution config\n", + "keras.distribution.set_distribution(data_parallel)\n", + "print(\"Distribution API configured!\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fe2c6ca1c15b" + }, + "source": [ + "## 3. Load Gemma 3 Model\n", + "\n", + "Because the global distribution is set to `DataParallel`, when we instantiate `Gemma3CausalLM`, Keras automatically replicates the weights across all 8 devices." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "a7e636f6b008" + }, + "outputs": [], + "source": [ + "# Load Gemma 3 (e.g., 4b parameter version)\n", + "model_id = \"gemma3_4b_en\" # Replace with standard Kaggle preset\n", + "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(model_id)\n", + "\n", + "gemma_lm.summary()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "f9fe8e903323" + }, + "source": [ + "## 4. Run Data-Parallel Batch Inference\n", + "\n", + "When passing a list of prompts, the workload is automatically sharded across the TPU cores." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5624ff57f34d" + }, + "outputs": [], + "source": [ + "batch_prompts = [\n", + " \"Explain the significance of true data parallelism in deep learning.\",\n", + " \"What are the key architectural improvements in Gemma 3?\",\n", + " \"Write a Python script to calculate Fibonacci using dynamic programming.\",\n", + " \"Describe the benefits of using JAX over PyTorch for TPU hardware.\",\n", + " \"How does the Keras 3 Distribution API simplify multi-core scaling?\",\n", + " \"Generate a summary of global warming mitigation strategies.\",\n", + " \"Write a creative short story about a sentient robot exploring Mars.\",\n", + " \"What is the maximum context window supported by Gemma 3, and how is it achieved?\"\n", + "]\n", + "\n", + "# Inference executes in parallel across the v5e-8 mesh\n", + "responses = gemma_lm.generate(batch_prompts, max_length=512)\n", + "\n", + "for i, response in enumerate(responses):\n", + " print(f\"\\n--- Prompt {i+1} ---\\n{response}\")" + ] + } + ], + "metadata": { + "colab": { + "name": "Gemma_3_Keras_TPU_Parallelism.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/tutorials/README.md b/tutorials/README.md index 1f55ecca..cc24e160 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -22,6 +22,7 @@ Explore Gemma's capabilities across different modalities: | [Vision - Video](../docs/capabilities/vision/video.ipynb) | Video understanding and analysis with Gemma 4. | | [Audio](../docs/capabilities/audio.ipynb) | Explore audio processing and understanding. | | [Thinking](../docs/capabilities/thinking.ipynb) | Reasoning capabilities. | +| [JAX/TPU Parallelism with Keras 3](Gemma_3_Keras_TPU_Parallelism.ipynb) | True data parallelism on Kaggle TPU v5e-8 mesh using Gemma 3. | ## Fine-tuning