Scaling LLM Fine-Tuning Optimization: A Multi-Agent Parallel Orchestration Pattern using Google Antigravity & Keras 3

How we built a real-time, parallel agentic web application to reduce developer time-to-value when configuring LLM fine-tuning on JAX & Keras 3.

Introduction

Fine-tuning a Large Language Model (LLM) like Google Gemma 2 is a highly non-trivial task. Developers often spend hours wrestling with three distinct configuration vectors:
1. Hyperparameters & LoRA Adapters: Finding the optimal learning rate, LoRA rank, alpha, and adapter targets for a given dataset and task.
2. Accelerated Infrastructure Sharding: Setting up DeviceMesh, LayoutMap, and FSDP (Fully Sharded Data Parallel) structures under Jax and Keras 3 to fit model weights and activations in limited GPU/TPU VRAM.
3. Data Pipeline Optimization: Building a high-throughput, non-blocking tf.data pipeline that tokenizes inputs, formats prompts matching Gemma’s template, and pre-fetches batches to prevent accelerator starvation.

To solve this, we built the Agentic Fine-Tuning Assistant. By leveraging the Google Antigravity SDK, the assistant orchestrates multiple specialist agents running in parallel to solve these vectors concurrently, presenting a unified, runnable fine-tuning script (fine_tune_gemma.py) to the developer in seconds. The generated code is available at github.

1. Multi-Agent Orchestration Architecture

To keep the reasoning pathways clean and minimize inference latency, we avoided a monolithic agent. Instead, we designed a parallel coordinator-specialist agent topology:

The Specialist Agents

• TunerAgent: Configures LoRA settings, batch size, learning rates, epochs, and optimizer choices based on the dataset size and task type.-
• HardwareAgent: Analyzes the hardware configuration (e.g., Single L4 GPU, TPU v5e-4) and creates JAX sharding strategies and LayoutMap bindings.
• DataAgent: Generates clean Python code utilizing tf.data and keras_nlp preprocessors to tokenize and batch the dataset matching Gemma’s instruction format.

The Master Orchestrator

• OrchestratorAgent: Takes the outputs of all three parallel subagents, cross-checks for configurations (e.g. ensuring batch sizes align between the sharding strategy and learning rates), and fuses them into a single, fully executable script.

Specialist subagent terminals reasoning in parallel

2. Backend Design Decisions

The backend is built using FastAPI and the Google Antigravity Python SDK.

Real-Time Streaming via Server-Sent Events (SSE)

Instead of waiting for the subagents to finish and presenting a slow page load, we implemented a real-time streaming endpoint /api/optimize returning a StreamingResponse (MIME type text/event-stream).

Get Fernanda Wanderley’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The SDK separates the model’s inner reasoning stream (thoughts) from its actual text generation chunk. We pipe both streams from all parallel agents into an asynchronous queue:

async def run_agent_and_push(agent_id, sys_instruction, prompt):
    config = LocalAgentConfig(
        system_instructions=sys_instruction,
        api_key=os.environ.get("GEMINI_API_KEY"),
        policies=[policy.deny_all()], # Deny all built-in filesystem/shell tools for safety
    )
    async with Agent(config) as agent:
        response = await agent.chat(prompt)
        
        # 1. Stream agent reasoning process (thoughts)
        async for thought in response.thoughts:
            await queue.put({"agent": agent_id, "type": "thought", "content": thought})
        
        # 2. Stream actual generation content
        async for chunk in response:
            await queue.put({"agent": agent_id, "type": "content", "content": chunk})
            
        await queue.put({"agent": agent_id, "type": "done"})

Safety and Sandboxing

Because the agents are running within our web application context, we enforce a strict Safety Policy by specifying policies=[policy.deny_all()] in LocalAgentConfig. This blocks the agents from triggering any built-in tools (such as local command execution or file creation) while permitting high-speed text reasoning.

Master Orchestrator agent streaming its thought process in real-time

3. Frontend Design Decisions & UX

Our frontend is built with a clean, high-fidelity Light Theme that prioritizes scanning and legibility:

• Typography: Outfit for clean, modern interface elements and JetBrains Mono for code blocks.
• Layout: A two-column structure. The left column houses the config form (model size, hardware, dataset). The right column shows the three subagent terminal cards side-by-side.
• Real-Time Progress UI: When the user clicks optimize, all three specialist terminals activate simultaneously. The UI shows custom progress spinners and streams the agents’ thoughts in distinct purple italic blocks (.thought-log) and response code in dark slate (.content-log).
• Code Console Mockup: Once the specialists finish, the Master Orchestrator activates. Once the final script is generated, the log collapses, revealing a clean, syntax-represented code-editor container with a “Copy to Clipboard” button.

4. Development Velocity: Reducing Time-to-Value

One of the most impressive metrics of the Google Antigravity platform is the developer velocity it unlocks. Under a traditional setup, building a multi-agent asynchronous orchestration dashboard with streaming thoughts, safety policies, and an interactive UI would take a senior developer days or weeks of work (requiring Celery, WebSockets, custom lock patterns, and extensive state machine wrappers).

Using Google Antigravity, we designed, coded, tested, and polished the entire application in under 20 minutes:

This represents a 95%+ reduction in development time, proving that Antigravity changes the time-to-value equation for agentic architectures.

5. Anatomy of the Generated Fine-Tuning Code

The orchestrator combines the specialist outputs into a complete fine_tune_gemma.py file. Let’s dissect the components of the generated code:

A. JAX/Keras 3 Distributed Sharding

For large models like Gemma 2, model parallelism is essential. Keras 3 maps cleanly to JAX’s native sharding APIs:

import keras
import numpy as np
# 1. Establish the Device Mesh matching target accelerator topology
# Example: 1x4 mesh for data-parallel and model-parallel sharding
mesh = keras.distribution.DeviceMesh(
 shape=(1, 4), 
 axis_names=("data", "model")
)
# 2. Define the sharding LayoutMap mapping weights to mesh axes
layout_map = keras.distribution.LayoutMap(mesh)
layout_map["token_embedding/embeddings"] = (None, "model")
layout_map["decoder_block.*attention.*query.*kernel"] = ("data", "model")
layout_map["decoder_block.*attention.*output.*kernel"] = ("model", "data")
# 3. Apply the layout map globally
keras.distribution.set_distribution(layout_map)

B. Gemma Chat Prompt Formatting & Tokenization

The script uses the official Gemma preprocessor to format datasets for instruction tuning:

import tensorflow as tf
import keras_nlp
# Instantiating the model preset downloads tokenizer and metadata
gemma_preprocessor = keras_nlp.models.GemmaCausalLM.from_preset(
 "gemma2_instruct_9b_en", 
 preprocessor_only=True
)
def format_prompt(instruction, response):
 # Formats inputs matching Gemma's official chat template:
 # <start_of_turn>user\nInstruction…<end_of_turn>
 return f"<start_of_turn>user\n{instruction}<end_of_turn>\n<start_of_turn>model\n{response}<end_of_turn>"
def preprocess_fn(features):
 prompts = format_prompt(features["instruction"], features["response"])
 return gemma_preprocessor(prompts)
# Build tf.data pipeline
dataset = tf.data.Dataset.from_tensor_slices(data)
dataset = dataset.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.cache().shuffle(1000).batch(4).prefetch(tf.data.AUTOTUNE)

C. LoRA Adapter Configuration

Using KerasNLP, enabling parameter-efficient tuning is a single method call:

# 1. Load the model weights under the defined distribution strategy
gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset("gemma2_instruct_9b_en")
# 2. Enable LoRA on the backbone with specialized adapter rank
gemma_lm.backbone.enable_lora(rank=8)
# 3. Compile and fit using mixed bfloat16 precision
gemma_lm.compile(
 loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
 optimizer=keras.optimizers.AdamW(learning_rate=2e-5, weight_decay=0.01),
 weighted_metrics=["accuracy"]
)
gemma_lm.fit(dataset, epochs=3)

Conclusion

By orchestrating parallel specialist subagents, we reduced the time-to-value of setting up Gemma 2 fine-tuning from hours of manual research to a 10-second interactive browser flow. Streaming the subagents’ thoughts directly to the user keeps the UX interactive and engaging, demonstrating the power of Google Antigravity to build developer-centric agentic applications.

Google Cloud credits are provided for this project. #AgenticArchitect #GoogleAntigravity