Training Marin 32B: What an open lab can build with TPUs, JAX, and a little persistence

by David Hall, Open Athena

Last summer, we partnered with Google to share how Marin trained a fully open 8B foundation model using JAX and TPUs. Since then, our process hasn't changed much, but the scale has. Over the summer, we trained a 32B model entirely in the open, and most days there was just one person keeping the run moving.

Large-scale training is usually associated with big teams and bigger infrastructure. Large model releases typically have hundreds of authors. Marin tests a different hypothesis: using open source software and data, small teams can train serious foundation models if the tooling is good, the platform is stable, and the process is transparent. The Marin 32B run was our strongest validation yet.

---

A model built with one hand on the helm

Marin was started at Stanford University's Center for Research on Foundation Models with the goal of building radically open foundation models. In May, we released Marin 8B Base, which bested the popular Llama 3.1 8B Base on 14 of 19 benchmarks. Marin 8B was trained using Google TPU v4 and TPU v5e from the TPU Research Cloud.

Building on that success, we set out to build a 32B model starting in June. Our 32B training run followed Marin's usual "Tootsie Roll" style: start with a solid recipe, instrument heavily, and adapt mid-flight when necessary. That flexibility matters, because the first time you train at a larger scale, issues inevitably arise.

The timing, however, was less than ideal, as universities tend to empty out over the summer. Students graduate, get internships, go home, or travel the world. Marin was no different. By June, our team was down to one full time research engineer, with a few PhD students providing guidance when they weren't busy with their dissertations. Nevertheless, we pushed forward.

To spoil the ending, the model turned out quite well. On release, Marin 32B Base was the best open source base model, and it outperformed comparable open-weights models like Google's Gemma 3 27B PT on 24 of 42 base-model evaluations.

There were many bumps along the way, resulting in multiple mid-run corrections, but through it all Google's TPU infrastructure stayed rock-solid, and JAX's predictable performance let us iterate quickly. This meant that even with a tiny team, we could diagnose, patch, and continue training without losing momentum.

To be blunt: one researcher kept the 32B run alive all summer, juggling preemptible slices, rebuilding optimizer state, switching architectures, and generally shepherding ~6.4 trillion tokens across v5p and v4 pods—while mostly working on other Marin projects. The fact that this was possible speaks to the stability of the TPU platform and the maturity of the JAX/Marin stack.

The short version of a long summer

Our retrospective goes into much more detail about every spike, switch and cooldown. Here's the condensed version.

We began with a Llama-3-style 32B backbone and our best 8B data mix, running on preemptible TPU v5p pods. Preemptions were predictable, and recovery was nearly automatic. As availability tightened, however, we moved to dedicated TPU v4 capacity. After a slight tweak to gradient checkpointing to accommodate the older hardware (made easy by JAX's built-in support), we were back up and running and performance stayed excellent.

Around 70k steps, persistent loss spikes appeared. We tried clipping, update-norm guards, skip-step heuristics, "necromancy" (rebuilding optimizer state), and swapping in optimizers like Muon. Nothing helped. The model needed architectural support.

So, we warm-started the run onto a Qwen3-style architecture, which is the same as the Llama 3 architecture, except that it adds QK-Norm to attention. After a brief loss bump, the spikes vanished. The model recovered to its expected trajectory within ~10 billion tokens and remained stable.

Towards the end of training, it was time for a cool down. When training LLMs, one "cools down" the model by lowering the learning rate and changing the data mix to higher quality data. Our first cooldown surfaced two issues: contamination from a cached math dataset, and a training-loss phase shift caused by our linear-congruential shuffle. Switching to a Feistel-based shuffle fixed the latter completely. After cleaning the data and re-running the cooldown, the second cooldown was smooth and produced the final model.

The result: a strong, open 32B base model

Marin 32B Base is a competitive open-source base model. It outperformed Olmo 2 32B Base—the previous best fully open-source base model—on 32 of 42 tasks, and it performs especially well on knowledge-heavy evaluations like ARC, BoolQ, and PIQA.

Head-to-head, Marin 32B Base also beat Gemma 3 27B PT on 24 of 42 tasks, and its overall average rank places it alongside Qwen 2.5 32B and the newer Olmo 3 32B models. On our evaluation suite, Marin 32B Base actually ties Olmo 3 32B Base in win rate, despite Olmo 3 being trained by a much larger team and arriving a month later.

Mean rank across our evaluation suite (lower is better). Marin 32B Base lands in the top cluster of open(-weight) models, alongside Qwen 2.5 and Olmo 3, and ahead of Gemma 3 27B PT and Olmo 2 32B. Gray bars indicate open weight models, while blue bars indicate open source models.

While Olmo 3 32B Base now comfortably leads on math and coding benchmarks, Marin 32B Base holds its own and still leads on many knowledge QA evaluations. For a model trained with a fraction of the team size typically expected for a 30B-scale run, we're proud of where it landed.

Because Marin 32B Base (like Olmo 3 32B) is open source, the weights, code, data recipes, and every experimental detour are public. Anyone can reproduce, audit, or build on the work.

---

The stack that made it possible

TPU stability across large slices

During the run, we moved across preemptible v5p-512 slices coordinated with Cloud TPU Multislice, a v4-2048 slice for the long middle, and several mid-run architectural transitions. Throughout, TPUs were completely reliable for us: no mysterious hangs, no collective-op debugging. Preemptions were predictable and easy to recover from.

JAX + Levanter = predictable performance

Levanter builds on JAX's XLA compilation. In practice, what mattered for us was deterministic restarts, stable MFU at scale without custom kernels, and JAX's activation checkpointing, which made the v5p to v4 migration easy.

Marin's experiment system

Marin logs every step of the experimental pipeline: hyperparameters, code versions, datasets, metrics, and artifacts. Even with architectural switches and restarts, the run never devolved into a tangle of scripts. And because it's all open, anyone can retrace or reproduce the training.

What's next

Marin 32B Base is a strong base model, but we're not done. Here's what's coming next:

• A reasoning-optimized Marin 32B
• Hardened multislice TPU support for smoother preemptible training
• Exploring MoE variants for the next scale
• Continuing to release everything, including successes and failures, openly

Closing thought

Training a 32B model with a small team isn't about heroics but about using the right tools and infrastructure. TPUs' reliability, JAX's clarity and performance, and Marin's open, reproducible process provided the leverage we needed. If the 8B run showed that open labs can build credible models, the 32B run showed they can do it at scale: quietly, steadily, and with far fewer people than you might expect.

SpatialReasoner: Teaching VLMs to "see" structure — Accelerated with Tunix on TPUs

by Yifan Shen & Ismini Lourentzou, University of Illinois Urbana-Champaign, and Srikanth Kilaru, Google ML Frameworks

Introduction

We are seeing an increasing interest in Tunix among researchers focusing on the post-training phase of model development. As a native JAX library, Tunix offers the flexibility needed to refine foundation models—including Vision-Language Models (VLMs) and not just LLMs—helping them significantly improve their spatial reasoning capabilities.

Today, we are highlighting the work of the PLAN Lab (Perception and LANguage Lab) at the University of Illinois Urbana-Champaign (UIUC). To address the critical lack of spatial awareness in VLMs, they built SpatialReasoner-R1, a model capable of fine-grained spatial logic. They utilized Tunix and leveraged the Google TPU Research Cloud (TRC) to scale their experiments.

In this blog, Professor Ismini Lourentzou and her team explain how they used Tunix's modular design to implement novel alignment algorithms and improve spatial reasoning in VLMs.

The "Where" Problem in VLMs

Modern Vision-Language Models (VLMs) can describe images and answer basic visual questions with impressive fluency. However, they often struggle with fine-grained spatial understanding. If you ask a VLM to estimate distances, directions, or the precise relative positions of objects, it frequently "hallucinates" coordinates or produces inconsistent reasoning with vague answers.

These capabilities are critical for real-world applications, such as robotics, where precise spatial reasoning enables safe and intelligent interaction with physical environments.

To bridge this gap, we developed the SpatialReasoner-R1 (4B and 8B versions), a model trained to perform step-by-step visually grounded spatial reasoning. It achieves 95.59 on Qualitative Accuracy and 77.3 on Quantitative Accuracy for our 8B fDPO model, outperforming the strongest baseline by ~9% in average accuracy on the SPATIALRGPT-Bench while preserving strong general vision-language abilities.

The Method: Fine-Grained Direct Preference Optimization (fDPO)

The secret sauce behind SpatialReasoner-R1 is a new technique called Fine-Grained Direct Preference Optimization (fDPO).

Standard alignment methods (like DPO) usually give a model a simple "thumbs up" or "thumbs down" for an entire response. But spatial reasoning is complex— for example, a model might correctly identify an object yet make a flawed logical inference about its location.

fDPO introduces segment-specific preference granularity. We optimize separate loss components for:

1. Descriptive Grounding: Does the model correctly perceive and describe the objects in the image?
2. Logical Reasoning: Is the step-by-step deduction sound and follows coherent spatial logic?

To generate high-quality training signals, we built a Multi-Model Monte Carlo Tree Search (M3CTS) data generation pipeline, which constructs diverse reasoning trajectories that guide the model toward reliable spatial understanding.

Tunix: Modularity for Novel Research

Implementing a custom objective like fDPO can be difficult in rigid frameworks. Tunix addresses this by providing a well-structured and extensible DPOTrainer that makes it possible to introduce new alignment objectives without reengineering the training pipeline.

This modularity meant we could reuse the entire underlying training stack—sharding, data loading, and loop management—while injecting our novel research logic with just a small amount of well-contained code.

While our backbone model (Sa2VA) required specific architectural handling, the core fDPO algorithm is model-agnostic. We found the Tunix experience smooth and well-documented, making it easy to prototype and iterate on fine-tuning workflows without reinventing the wheel.

Google TRC & TPUs: Reliability at Scale

Training a model to reason over long horizons requires significant compute. The Google TPU Research Cloud (TRC) provided the infrastructure we needed to make large-scale training practical.

• Scalability: Tunix's integration with TPUs allowed us to scale our experiments seamlessly.
• Reliability: The system performed reliably across multiple TPU runs, which was essential for conducting large-scale spatial reasoning benchmarks.
• Support: The Google Tunix and TRC teams assisted with infrastructure setup and experiment design, helping us refine our multi-model exploration strategy.

Looking Ahead: Open Source Contributions

We believe that open-source, extensible tools like Tunix are vital for fostering innovation. They lower the barrier for researchers to experiment with new training objectives without rebuilding core infrastructure.

In that spirit, we contributed our fDPO implementation back to the Tunix ecosystem. We open-source the core fDPO components, enabling the community to apply segment-specific preference optimization to their own models.

Get Started

You can explore our research and the tools we used below:

• SpatialReasoner Project Page
• Tunix Documentation
• Tunix GitHub Repository
• Google TRC
• JAX documentation
• JAX AI Stack documentation

GRL: Turning verifiable games into a post-training suite for LLM agents with Tunix on TPUs

by The GRL Team, UC San Diego and Lin Chai & Srikanth Kilaru, Google ML Frameworks

Introduction

JAX is widely recognized for its power in training large-scale AI models. However, a primary bottleneck in the next phase of AI development—LLM post-training with Reinforcement Learning (RL)—is the scarcity of environments with verifiable rewards.

Today, we are highlighting the work of the GRL (Game Reinforcement Learning) team at UC San Diego. To solve the data bottleneck, they have built a pipeline to turn video games into rigorous reasoning benchmarks. They utilized Tunix, a JAX-native research-friendly RL framework that supports multi-host, multi-turn capabilities, and leveraged the Google TPU Research Cloud (TRC) to scale their experiments. The results are promising: this approach has yielded significant improvements in model quality, particularly in planning and reasoning tasks, proving that games can be a viable substrate for serious AI capability training.

In this blog the GRL team explains how they are combining game environments, modular Tunix library for RL post-training, and TPU compute to train the next generation of agents.

---

Why Verifiable Games for LLM Post-Training?

Current RL post-training has shown strong gains in domains like math and coding because success can be auto-checked. However, these settings are often narrow and short-term. We are effectively overfitting RL to clean problems, while the next generation of agents must operate in messy, multi-step worlds.

To unlock RL as a systematic method for reasoning, we need a diverse pool of environments where rewards are grounded in explicit, machine-checkable rules. Games are this missing, underused substrate.

1. The Performance Gap: LLMs still perform surprisingly poorly on many strategy games, revealing a clear gap between model behavior and human-level interactive competence.
2. Verifiable Signals: Games come with built-in verifiable signals—wins, scores, puzzle completion—meaning outcomes are automatically and unambiguously graded without human labeling.
3. Long-Horizon Reasoning: Unlike short QA tasks, games force models to plan, explore, and reason over many steps.
4. Abundance: Decades of RL research has produced a standardized ecosystem of diverse environments ready to be recycled.

Game Reinforcement Learning (GRL): A Unified Game-to-Post-Training Pipeline

To harness this ecosystem, we built GRL, a comprehensive suite designed to recycle diverse game environments into a reusable post-training resource. Our mission is to prioritize environments with executable success checks—ranging from text-based puzzles to embodied 3D worlds and web/GUI workflows. Our code and ecosystem live under the LM Games organization (lmgame.org).

GRL provides three key capabilities:

• A Unified Pipeline: We standardize the conversion of games into RL-ready environments with structured states and consistent metrics. This makes results comparable across models and research groups.
• Versatile Configuration: Researchers can tailor interaction styles (e.g., max_turns, natural language feedback) while mixing training data from different tasks seamlessly. This allows for training on puzzles, math, and web tasks within a single run.
• Algorithm-Agnostic Interface: GRL works with any agentic training algorithm. While we frequently use PPO, the system serves as a robust testbed for developing new RL techniques.

The Engine: Plugging into the Tunix RL Framework

Designed for Research Flexibility and Multi-Turn Agents

In practice, plugging a GRL game agent into Tunix is seamless thanks to its modular design. Tunix is built specifically to support multi-turn agentic tasks, allowing researchers to leverage native one-turn inference APIs to achieve complex multi-turn rollouts, then batch those outputs directly back into the training flow. This research flexibility is key; the framework is lightweight enough for quick iteration and benchmarking, yet modular enough to allow fine-grained adjustments to reward functions, algorithms, and hardware-aware settings like mesh sizes.

We first define an agent_cfg (see picture above) that tells the system which game to play (eg. Sokoban or Tetris), how the LLM should talk (chat template + reasoning style), and its budgets (max turns, tokens per turn, action format). On the Tunix side, we then load a pre-trained model into three roles: actor, critic, and reference and build ClusterConfig to specify rollout and training configs and PpoConfig to specify RL hyperparameters. The glue is minimal and the layout is clear and research friendly: once agent_cfg, ppo_cfg, and cluster_cfg are defined, we construct an RLCluster and pass everything into PpoLearner, which gives us a complete multi-turn PPO trainer in JAX.

Our multi-turn RL workflow is equally lightweight from the user's point of view. For example, with a 5-turn budget, the trainer repeatedly lets the LLM "play" the game for up to five conversational turns: at each turn it sees the current grid or state, reasons in language using the chat template, outputs a series of actions, and receives the next state and a verifiable reward signal (win/loss/score/step penalty). GRL's agent + env configs handle all the orchestration: they log observations, actions, and rewards into structured trajectories, which Tunix then turns into token-level advantages and returns for PPO updates. You don't manually build datasets or rollouts; the trainer owns the loop - interact -> log -> compute rewards -> update policy -> repeat.

In our preliminary experiments using this setup, training Qwen2.5-7B-Instruct on Sokoban and Tetris yielded strong in-domain gains (+2-56% across game variants). We also observed modest generalization to out-of-domain tasks, with consistent improvements in planning tasks (Blocksworld: +3-7%) and positive but unstable signals in computer use (Webshop: ~+6%). All scripts and configs are available in the GRL repo: https://github.com/lmgame-org/GRL/tree/main. To reproduce the end-to-end Tunix + GRL training example (including our Sokoban/Tetris runs), you can simply clone the repo and run one line: bash tunix_quick_training_example.sh.

Google TRC & TPUs: Accelerating Game-Based RL at Scale

A critical component of our research was the Google TPU Research Cloud (TRC) program. Access to Cloud TPUs allowed us to move from small-scale prototypes to production-grade training runs with minimal friction.

TPUs and JAX directly attacked our two biggest bottlenecks:

1. Rollout Throughput: Using the vLLM-TPU path via tpu-inference, we could serve multiple model families on the same TPU v5p backend. This boosted sampling throughput, making the data-collection loop tighter and multi-environment concurrency cheaper.
2. Multi-Host Scale for 7B Models: Tunix's lightweight design combined with JAX's mesh-based sharding allowed us to scale the same code from a single host to multi-host setups declaratively. This capability was essential for our experiments with 7B parameter models (such as Qwen2.5-7B), where we leveraged 2 v5p-8 hosts with minimal code change (in fact, only an env var config). The scale up is seamless, proving that the infrastructure can handle the heavy computational lifting required for modern LLM post-training without requiring complex engineering overhauls.
3. Hardware Advantage: At the hardware level, the performance gains were significant. Each TPU v5p chip delivers around 459 BF16 TFLOPs, compared to roughly 312 on an NVIDIA A100. This raw power, combined with the TRC program's support, meant that large-N studies—involving more seeds, longer horizons, and more environments—became routine experiments rather than "special ops" engineering challenges.

This combination of Tunix's flexible abstraction and TRC's massive compute resources allowed us to iterate quickly on ideas while benefiting from production-grade infrastructure.

Get Started

GRL and Tunix are open for the community to explore. You can reproduce our end-to-end training example (including the Sokoban/Tetris runs) by cloning the repo, following the installation instructions, and then running a single command:

bash tunix_quick_training_example.sh

• GRL Repository
• Tunix Documentation
• Tunix GitHub Repository
• Google TRC
• JAX documentation
• JAX AI Stack documentation

ESCA: Grounding embodied agents with scene graphs — Accelerated by JAX

by The ESCA Team, University of Pennsylvania & Srikanth Kilaru, Google ML Frameworks

Introduction

Multi-Modal Language Models (MLLMs) are increasingly forming the core of the brain for general-purpose embodied agents — AI that can navigate and act in the physical world as robots. While MLLMs are making rapid progress, they often stumble on a critical hurdle: precise visual perception. They struggle to reliably capture the fine-grained links between low-level visual features and high-level textual semantics.

Today, we are highlighting the work of Prof. Mayur Naik's research team at the University of Pennsylvania. To bridge the gap between high-level language and low-level visual features, they developed ESCA (Embodied and Scene-Graph Contextualized Agent). By porting their neurosymbolic pipeline to JAX, they achieved the real-time performance necessary for high-throughput decision-making. This work also demonstrates that JAX drives performance gains across a wide range of hardware, including standard CPUs and NVIDIA GPUs, and not just on Google TPUs.

In this blog, the UPenn team explains how they combined structured scene graphs with JAX's functional design to reduce perception errors by over 50% and achieve a 25% speedup in inference.

---

The "Grounding" Problem in Embodied AI

Existing MLLMs are powerful, but they can be surprisingly "blind" when tasked with interacting with the physical world. In our empirical analysis of 60 navigation tasks from EmbodiedBench, we found that 69% of agent failures stemmed from perception errors. See the figure below.

The three top-level error types are Perception, Reasoning, and Planning. The second-level errors are Hallucination, Wrong Recognition, Spatial Understanding, Spatial Reasoning, Reflection Error, Inaccurate Action, and Collision. For clarity, the figure uses these acronyms to label the different error types.

The models struggle to capture fine-grained links between visual features and textual semantics. They might recognize a "kitchen," but fail to identify the specific spatial relationship between a knife and a cutting board required to complete a task.

Enter ESCA: The Anglerfish of AI

To solve this, we introduced ESCA, a framework designed to contextualize MLLMs through open-domain scene graph generation.

Think of ESCA like the bioluminescent lure of a deep-sea anglerfish. Just as the fish illuminates its dark surroundings to reveal prey, ESCA "illuminates" the agent's environment by generating a structured Scene Graph—a map of objects, attributes, and relationships (e.g., Cup [Red] ON Table).

A key innovation here is Selective Grounding. Injecting a massive scene graph of everything in the room can overwhelm the model. Instead, ESCA identifies only the subset of objects and relations pertinent to the current instruction. It performs probabilistic reasoning to construct prompts enriched with exactly the contextual details the agent needs to act.

The Engine: LASER and Scallop

At the core of ESCA is LASER, a CLIP-based foundation model trained on 87k video-caption pairs. LASER uses Scallop—our neurosymbolic programming language that supports JAX backends—to align predicted scene graphs with logical specifications. This pipeline allows us to train low-level perception models to produce detailed graphs without needing tedious frame-level annotations.

JAX User Experience

1. The Power of Statelessness

JAX's design encouraged a fully functional, stateless architecture. Every component, from feature extraction to similarity computation, was made into a pure modular function. This structure enabled effective use of jit (Just-In-Time) compilation. The XLA compiler could fuse sequences—like normalization, matrix multiplication, and softmax—into fewer kernels, reducing intermediate buffers and lowering GPU overhead.

2. Handling Complex Control Flow

Our pipeline requires selecting the "top-k" most relevant objects from a probabilistic scene graph. This introduces complex control flow. JAX provided the primitives we needed to handle this efficiently:

• We used jax.lax.cond to manage control flow inside the probabilistic graph.
• We leveraged jax.nn and jax.numpy for all activation functions and batched math in a JIT-friendly way.

3. Debugging and Transparency

Migrating to JAX was also a learning experience. Tools like jax.debug.print/callback() allowed us to inspect values inside jit-compiled functions, while jax.disable_jit() let us easily switch to eager execution to step through the program seeing intermediate values.

Furthermore, the transparency of the open-source system was impressive. Being able to read the annotated source code and see how Python functions trace into jaxpr (JAX expression) gave us deep insight into how to design inference logic that scales.

4. Seamless Integration with Flax

NNX fits into our workflow perfectly. We used nnx.Module to structure the model and FrozenDict to keep parameters organized and immutable. The TrainState object made managing model parameters and optimizer states straightforward, without adding the complexity often found in other frameworks.

JAX Performance: A 25% Speedup

Embodied agents operate in a continuous loop: planning, acting, and updating their understanding of a dynamic world. High latency here is a dealbreaker. We ported LASER from PyTorch to JAX to improve real-time performance, and the benefits were significant. By rewriting our core similarity computations and feature pipelines as pure functions wrapped in jax.jit, we achieved significant gains.

On an NVIDIA H100 GPU, JAX reduced the average time per frame from 18.15 ms (PyTorch) to 14.55 ms (JAX)—a roughly 25% speedup.

Framework	Hardware	Avg Time Per Frame (ms) ↓	FPS ↑
PyTorch	H100 GPU	18.15 ± 0.73	55.15 ± 2.31
JAX	H100 GPU	14.55 ± 0.64	68.82 ± 3.13

Conclusion

ESCA demonstrates that better data—structured, grounded scene graphs—can solve the perception bottleneck in Embodied AI. But it also demonstrates that better infrastructure is required to run these systems in the real world. JAX provided the speed, transparency, and modularity needed to turn our research into a real-time agent capable of reliable reasoning.

Acknowledgements

This research was made possible through support from a Google Research Award to the University of Pennsylvania and from the ARPA-H program on Safe and Explainable AI under award D24AC00253-00.

Get Started

You can explore the LASER code, the ESCA framework and documentation for JAX and Flax at:

• ESCA Repository
• LASER/SGClip Repository
• JAX Documentation
• Flax Documentation

Empowering app developers: Fine-tuning Gemma 3 for mobile with Tunix in Google Colab

by Henry Ndubuaku & Noah Cyclich, Cactus Compute and Lance Wang & Srikanth Kilaru, Google ML Frameworks

In the rapidly evolving world of AI models for mobile devices, a persistent challenge is how to bring SOTA LLMs to smartphones without compromising on privacy or requiring App developers to be Machine Learning engineers.

Today, we are excited to talk about how Cactus, a startup building a next-gen inference engine for mobile devices, fine-tunes the open-source Gemma 3 model. By leveraging Tunix, the LLM post-training library in the JAX ML ecosystem, they achieved this entirely on Google Colab's Free Tier.

The Challenge: Making Small Models "Expert"

For app developers, running Large Language Models (LLMs) in the cloud isn't always an option due to privacy concerns (like GDPR) and latency requirements. The solution lies in running models locally on the device. However, most smartphones globally lack specialized MPUs (Micro Processing Units), meaning developers need highly efficient, smaller models.

While compact models like Gemma (270M or 1B parameters) are incredibly efficient, they are often "generalists." To be useful for specific mobile applications—such as a medical imaging assistant or a legal document analyzer—they need to be fine-tuned to become domain experts.

The problem? Most app developers are not ML infrastructure experts. Setting up complex training pipelines, managing dependencies, and navigating steep learning curves creates too much friction.

The Solution: SFT via Tunix on Google Colab

To solve this, Cactus created a simplified "Low-Friction" workflow by implementing a Python script using Supervised Fine Tuning (SFT) APIs of Tunix in a Colab.

1. The Engine: Tunix

Cactus utilized Tunix, Google's lightweight and modular LLM post-training library, which supports both SFT and leading RL algorithms, and executes natively on TPUs. Tunix strips away the complexity of heavy frameworks, offering a simplified path to Supervised Fine-Tuning (SFT).

2. The Access: Google Colab Free Tier

Accessibility was a key requirement. Instead of requiring developers to set up complex cloud billing and project IDs immediately, the workflow operates entirely within a Google Colab Notebook. By utilizing the free tier of Colab, developers can:

• Load the Gemma 3 model.
• Upload their specific dataset (e.g., medical data or customer service logs).
• Run an SFT (Supervised Fine-Tuning) job using Tunix.
• Export the weights for conversion.

3. The Deployment: Cactus

Once tuned, the model is converted into the Cactus graph format. This allows the now-specialized Gemma 3 model to be deployed directly into a Flutter or native mobile app with just a few lines of code, running efficiently on a wide range of smartphone hardware.

Why This Matters

"Our users are app developers, not ML engineers," explains Henry Ndubuaku, co-founder of Cactus. "They want to pick a model, upload data, and click 'tune.' By using Tunix and Colab, we can give them a 'clone-and-run' experience that removes the intimidation factor from fine-tuning."

This workflow represents the "lowest hanging fruit" in democratizing AI:

• No complex local environment setup.
• No upfront infrastructure costs.
• High-performance JAX native Tunix library to tune a leading OSS model (Gemma).

What's Next?

While the Colab notebook provides an immediate, accessible solution, Cactus is exploring a future plan to build a full GUI-based portal for fine-tuning and quantization of LLMs with the back end compute as Google Cloud TPUs, allowing for scalable training of larger models and even more seamless integration into the mobile development lifecycle.

Get Started

Ready to turn your mobile app into an AI powerhouse? Check out the Tunix SFT Notebook for Cactus and start fine-tuning Gemma 3 for your device today:

• GRPO Notebook
• QLoRA Notebook
• DPO

You can explore Tunix sample scripts, documentation and repo at:

• Tunix sample scripts
• Tunix Documentation
• Tunix GitHub Repository

Shape the future with Google Summer of Code 2026!

By Stephanie Taylor, Mary Radomile & Lucila Ortiz

Are you a passionate beginner ready to make your mark in open source? Now is your chance to make an impact in the 2026 Google Summer of Code (GSoC) program!

For over two decades, GSoC has been a critical launchpad, consistently introducing fresh, enthusiastic talent into the open source ecosystem. The numbers speak for themselves: GSoC has connected more than 22,000 contributors with over 20,000 mentors from 1,000+ open source organizations, all collaborating to keep this vital community thriving.

---

Google, like everyone else, is built on open source. We depend on a healthy, vibrant open source ecosystem and want to lower the barrier to entry for people who want to contribute to the open community. Join in, learn great skills, and make an impact on people around the world.

Richard Seroter, Chief Evangelist, Google Cloud

---

Over 3+ months, contributors spend their time collaborating on real-world projects right alongside experienced mentors and their communities. This isn't just about coding; this deep immersion in open source does more than build valuable technical skills. It cultivates a strong understanding of community dynamics, best practices, and the soft skills needed to become a truly impactful open source contributor.

Check out the trailer below to learn more about the massive impact GSoC has made over the last two-plus decades.

Google Summer of Code trailer

Join GSoC as a Mentoring Organization

Application Period: January 19 – February 3

If you're interested in your organization participating, now is the time to start! We welcome around 30 new organizations to GSoC annually—and yours could be one of them.

• Visit our website where you'll find supportive materials to get started.
• The Mentor Guide is a must-read, offering an introduction to GSoC, community engagement tips, project idea suggestions, and guidance on applying.
• Pro Tip: For 2026, we'll have an expanded focus on the AI, Security, and Machine Learning domains.

Want to be a GSoC Contributor?

Application Period: March 16 – March 31

Ready to get started? The official website is your first stop!

• Watch the Introduction to GSoC video for a quick overview.
• For more in-depth preparation, read the Contributor Guide—it's packed with insightful ideas from past participants.
• Be sure to review the program rules and FAQ to understand eligibility requirements and how the program works.

---

Spread the Word!

Please help us amplify the message about GSoC 2026! Share this post with your peers, family members, colleagues, universities, and anyone interested in making a real difference in the open source community.

Join us and help shape the future of open source!

AI Conformant Clusters in GKE

AI Conformant Clusters in GKE

by Janet Kuo & Federico Bongiovanni, GKE

We are excited to announce that Google Kubernetes Engine (GKE) is now a CNCF-certified Kubernetes AI conformant platform, designed to provide a stable and optimized environment for your AI/ML applications. This initiative, culminating in a major announcement of the Kubernetes AI Conformance program by CNCF's CTO Chris Aniszczyk, at KubeCon NA 2025, is set to simplify AI/ML on Kubernetes for everyone. You can check out the Opening Keynote here.

During the keynote, Janet Kuo, author of this blog and Staff Software Engineer at Google, performed a live Demo, demonstrating the practical power of an AI Conformant cluster. If you are interested in the technical specifics, you can learn more about the demo here.

Why AI Conformance Matters

The primary goal of the Kubernetes AI Conformance program is to simplify AI/ML on Kubernetes, guarantee interoperability and portability for AI workloads, and enable a growing ecosystem of AI tools on a standard foundation.

Setting up a Kubernetes cluster for AI/ML can be a complex undertaking. An AI-conformant platform like GKE handles these underlying complexities for you, ensuring that your environment is optimized for scalability, performance, portability, and interoperability.

For a detailed look at all the requirements and step-by-step instructions on how to create an AI-conformant GKE cluster, we encourage you to read the GKE AI Conformance user guide.

What Makes GKE an AI-Conformant Platform?

A Kubernetes AI-conformant platform like GKE handles the underlying complexities for you, providing a verified set of capabilities to run AI/ML workloads reliably and efficiently. Here are some of the key requirements that GKE manages for you:

• Dynamic Resource Allocation (DRA): GKE enables more flexible and fine-grained resource requests for accelerators, going beyond simple counts. This is crucial for workloads that need specific hardware configurations.
• Intelligent Autoscaling for Accelerators: GKE implements autoscaling at both the cluster and pod level to ensure your AI workloads are both cost-effective and performant.
  • Cluster Autoscaling works at the infrastructure level. It automatically resizes node pools with accelerators, adding nodes when it detects pending Pods that require them and removing nodes to save costs when they are underutilized.
  • Horizontal Pod Autoscaling (HPA) works at the workload level. HPA can automatically scale the number of your pods up or down based on real-time demand. For AI workloads, this is especially powerful, as you can configure it to make scaling decisions based on custom metrics like GPU/TPU utilization.
• Rich Accelerator Performance Metrics: GKE exposes detailed, fine-grained performance metrics for accelerators. This allows for deep insights into workload performance and is essential for effective monitoring and autoscaling.
• Robust AI Operator Support: GKE ensures that complex AI operators, such as Kubeflow or Ray, can be installed and function reliably, enabling you to build and manage sophisticated ML platforms with CRDs.
• All-or-Nothing Scheduling for Distributed Workloads: GKE supports gang scheduling solutions like Kueue, which ensure that distributed AI jobs only start when all of their required resources are available, preventing deadlocks and resource wastage.

A Unified and Evolving Standard

The Kubernetes AI Conformance program is designed as a single, unified standard for a platform to support all AI/ML workloads. This reflects the reality that modern AI processes, from training to inference, increasingly rely on the same underlying high-performance infrastructure.

What's Next?

We invite you to explore the benefits of running your AI/ML workloads on an AI-conformant GKE cluster.

• Learn more about the program: check out the CNCF Kubernetes AI Conformance repository.
• Dive deeper into AI/ML on GKE: read the Introduction to AI/ML workloads on GKE.

The launch of the AI Conformance program is a significant milestone, but it is only the first step. We are eager to continue this conversation and work alongside the community to evolve and improve this industry standard as we head into 2026.

Secure-by-design firmware development with Wasefire

Improving firmware development

by Julien Cretin, Security, Privacy, and Abuse Research

Building firmware for embedded devices—like microcontrollers and IoT hardware—is hard. It's often complex, it requires deep expertise, and most importantly it is prone to security bugs. One of the key challenges is the limited resources available on these devices, such as constrained processing power, memory, and storage capacity. These constraints make implementing robust security measures at odds with performance and functionality. Unsafe IoT devices are then recruited by cyber criminals into botnets, to perform DDoS attacks, steal information, and act as proxies to evade detection (e.g. the Mirai botnet).

Today, we introduce a new framework that makes it easier to build and maintain safer embedded systems: Wasefire.

Wasefire simplifies the development process and incorporates security best practices by default. This enables developers to create secure firmware without requiring extensive security expertise and only focusing on the business logic they want to implement. To this end, Wasefire provides for each supported device, a platform on which device-agnostic sandboxed applets can run. Wasefire currently supports nRF52840 DK, nRF52840 Dongle, nRF52840 MDK Dongle, and OpenTitan Earlgrey. There is also a Host platform for testing without embedded devices.

The platform is written in Rust for its performance and built-in memory safety. Embedded devices are one of the four target domains of the Rust 2018 roadmap. So today, it is quite simple to write embedded code in Rust, or even integrate Rust in existing embedded code.

The platform expects the applets to be written in—or more realistically, compiled to—WebAssembly for its simplicity, portability, and security. WebAssembly is a binary instruction format for a stack-based virtual machine. It is designed for high-performance applications on the web (hence its name) but it also supports non-web environments. Fun fact: Wasefire uses WebAssembly in both environments: the main usage is non-web for the virtual machine to run applets, but the web interface of the Host platform also relies on WebAssembly.

Incidentally, WebAssembly is another one of the four target domains of the Rust 2018 roadmap. This means that writing applets in Rust and compiling them to WebAssembly is very simple. For this reason, Rust is the primary language to write applets for Wasefire. Starting a new project is as simple as the following steps:

WebAssembly on microcontrollers

Running WebAssembly on microcontrollers may seem like overkill if it were only for sandboxing. But using a virtual machine also provides binary-level portability like Java Cards. In particular, the same WebAssembly applet can be distributed in binary form and run on multiple platforms.

On a microcontroller, every byte matters. To cater to a variety of needs, Wasefire provides multiple alternatives to balance between security, performance, footprint, and portability:

• WebAssembly applets: Platforms may embed the Wasefire interpreter. This is a custom in-place interpreter for WebAssembly in the style of "A fast in-place interpreter for WebAssembly" with a very small footprint. The main drawback is that it doesn't support computation heavy applets.
• Pulley applets: Platforms may embed Wasmtime and its Pulley interpreter. WebAssembly was not designed for interpretation, but for compilation. So WebAssembly interpreters will experience some form of challenge (either performance or footprint). On the contrary, Pulley was designed for fast interpretation and can be compiled from WebAssembly. The main drawback is the larger footprint of this solution and the need for applets to be signed (which is not yet implemented) since Pulley cannot be validated like WebAssembly.
• Native applets: Platforms may link with an applet compiled as a static library for the target architecture. This solution is only provided as a last resort when no other existing alternative works. The main drawback is that almost all security benefits are nullified and binary-level portability is lost.
• CHERI applets: This alternative is planned (but not yet started) and would provide the performance and footprint advantage of Native applets while retaining the sandboxing advantage of WebAssembly and Pulley applets. The main drawback is that the target device needs to support CHERI and binary-level portability is lost.

To illustrate this tradeoff, let's look at a few examples from the Wasefire repository:

• The first example is a button-controlled blinking LED. This applet can run as a WebAssembly applet without problem.
• The second example is a FIDO2 security key implemented using the OpenSK library. This applet reaches the limits of the WebAssembly in-place interpreter in terms of performance at the moment. By using a Pulley applet instead, performance can be improved by degrading applet size and memory footprint.
• The third example is a BLE sniffer. Performance is a critical aspect of this applet. The in-place interpreter is too slow and many packets are dropped. Compiling this applet to Pulley doesn't drop any packet in a noisy BLE environment.

We can summarize the tradeoff in the table below. The platform size differs between examples because the second and third examples need optional drivers disabled by default. The platform is the nRF52840 DK. For the security key, applet performance is measured as the time between a FIDO2 GetInfo request and the last packet of its response. For the BLE sniffer, applet performance is measured as the number of processed packets per second. This metric saturates for Pulley and Native applets, so we only get a lower bound of performance in those cases.

Blinking LED	WebAssembly	Pulley	Native
Platform size (KiB)	98	299	49
Applet size (KiB)	3.3	12	5.6
Platform memory (KiB)	10	80	5
Security key	WebAssembly	Pulley	Native
Platform size (KiB)	133	334	80
Applet size (KiB)	125	247	73
Platform memory (KiB)	20	104	9
Applet performance (ms)	1191	60	23
BLE sniffer	WebAssembly	Pulley	Native
Platform size (KiB)	102	303	53
Applet size (KiB)	7.2	18	7.6
Platform memory (KiB)	16	82	8.8
Applet performance (packet/s)	= 55 (dropping packets)	> 195 (not dropping)	> 195 (not dropping)

Looking forward

Wasefire is still an experimental project. Many features are missing (including security features) and many improvements are planned. For example, the platform currently runs a single applet and provides all the resources this applet asks for. Ultimately, applets would come with a manifest describing which resources they are permitted to use, and those resources would be isolated to that single applet. It would also be possible to run multiple applets concurrently.

The project is open source, so bug reports, feature requests, and pull requests are welcome. The project is licensed under Apache-2.0, so commercial use is permitted.

Feel free to give it a try (no hardware needed) and share the word!

How JAX makes high-performance economics accessible

How JAX makes high-performance economics accessible

John Stachurski, QuantEcon & Srikanth Kilaru, Google ML Frameworks

JAX is widely recognized for its power in training large-scale AI models, but its core design as a system for composable function transformations unlocks its potential in a much broader scientific landscape. We're seeing adoption for applications as disparate as AI-driven protein engineering to solving high-order Partial Differential Equations (PDEs). Today, we're excited to highlight another frontier where JAX is making a significant impact: enabling economists to model complex, real-world scenarios that shape national policy—computational economics.
I recently spoke with economist John Stachurski, a co-founder of QuantEcon and an early advocate for open-source scientific computing. His story of collaborating with the Central Bank of Chile demonstrates how JAX makes achieving performance easy and accessible. John's journey shows how JAX's intuitive design and abstractions allow domain experts to solve scientific problems without needing to become parallel programming specialists. John shares the story in his own words.

---

A Tale of Two Implementations: The Central Bank of Chile's Challenge
Due to my work with QuantEcon, I was contacted by the Central Bank of Chile (CBC), which was facing a computational bottleneck with one of their core models. The bank's work is high-stakes; their role is to set monetary policy and act as the lender of last resort during financial crises. Such crises are inherently non-linear in nature, involving self-reinforcing cycles and feedback loops that make them challenging to model and assess.
To better prepare themselves for such crises, the CBC began working on a model originally developed by Jarvier Bianchi, in which an economic shock worsens the balance sheets of domestic economic agents, reducing collateral and tightening credit constraints. This leads to further deterioration in balance sheets, which again tightens credit constraints, and so on. The result is a downward spiral. The ramifications can be large in a country such as Chile, where economic and political instability are historically linked.

The Problem:

The task of implementing this model was led by talented CBC economist Carlos Rondon. Carlos wrote the first version using a well-known proprietary package for mathematical modeling that has been used extensively by economists over the past few decades. The completed model took 12 hours to run -- that is, to generate prices and quantities implied by a fixed set of parameters -- on a $10,000 mainframe with 356 CPUs and a terabyte of RAM. A 12 hour run-time made it almost impossible to calibrate the model and run useful scenarios. A better solution had to be found.

Carlos and I agreed that the problem was rooted in the underlying software package. The issue was that, to avoid using slow loops, all operations needed to be vectorized, so that they could be passed to precompiled binaries generated from Fortran libraries such as LAPACK. However, as users of these traditional vectorization-based environments will know, it is often necessary to generate many intermediate arrays in order to obtain a given output array. When these arrays are high-dimensional, this process is slow and extremely memory intensive. Moreover, while some manual parallelization is possible, truly efficient parallelization is difficult to achieve.

The JAX Solution:

I flew to Santiago and we began a complete rewrite in JAX. Working side-by-side, we soon found that JAX was exactly the right tool for our task. In only two days we were able to reimplement the model and — running on a consumer-grade GPU — we observed a dramatic improvement in wall-clock time . The algorithm was unchanged, but even a cheap GPU outperformed the industrial server by a factor of a thousand. Now the model was fully operational: fast, clean, and ready for calibration.
There were several factors behind the project's success. First, JAX's elegant functional style allowed us to express the economic model's logic in a way that closely mirrored the underlying mathematics. Second, we fully exploited JAX's vmap by layering it to represent nested for loops. This allowed us to work with functions that operate on scalar values (think of a function that performs the calculations on the inside of a nested for loop), rather than attempting to operate directly on high dimensional arrays — a process that is inherently error-prone and difficult to visualize.

Third, JAX automates parallelization and does it extremely efficiently. We both had experience with manual parallelization prior to using JAX. I even fancied I was good at this task. But, at the end of the day, the majority of our expertise is in economics and mathematics, not computer science. Once we handed over parallelization to JAX's compiler OpenXLA we saw a massive speed up. Of course, the fact that XLA generates specialised GPU kernels on the fly was a key part of our success.
I have to stress how much I enjoyed completing this project with JAX. First, we could write code on a laptop and then run exactly the same code on any GPU, without changing a single line. Second, for scientific computing, the pairing of an interpreted language like Python with a powerful JIT compiler provides the ideal combination of interactivity and speed. To my mind, everything about the JAX framework and compilers is just right. A functional programming style makes perfect sense in a world where functions are individually JIT-compiled. Once we adopt this paradigm, everything becomes cleaner. Throw in automatic differentiation and NumPy API compatibility and you have a close-to-perfect environment for writing high performance code for economic modeling.

---

Unlocking the Next Generation of Economic Models

John's story captures the essence of JAX's power. By making high performance accessible to researchers, JAX is not just accelerating existing workloads; it's democratizing access to performance and enabling entirely new avenues of research.
As economists build models that incorporate more realistic heterogeneity—such as varying wealth levels, firm sizes, ages, and education—JAX enables them to take full advantage of modern accelerators like GPUs and Google TPUs. JAX's strengths in both scientific computing and deep learning make it the ideal foundation to bridge this gap.

Explore the JAX Scientific Computing Ecosystem

Stories like John's highlight a growing trend: JAX is much more than a framework for building the largest machine learning models on the planet. It is a powerful, general-purpose framework for array-based computing across all sciences which, together with accelerators such as Google TPUs and GPUs, is empowering a new generation of scientific discovery. The JAX team at Google is committed to supporting and growing this vibrant ecosystem, and that starts with hearing directly from you.

• Share your story: Are you using JAX to tackle a challenging scientific problem? We would love to learn how JAX is accelerating your research.
• Help guide our roadmap: Are there new features or capabilities that would unlock your next breakthrough? Your feature requests are essential for guiding the evolution of JAX.

Please reach out to the team via GitHub to share your work or discuss what you need from JAX. You can also find documentation, examples, news, events, and more at jaxstack.ai and jax.dev.

Sincere thanks to John Stachurski for sharing his insightful journey with us. We're excited to see how he and other researchers continue to leverage JAX to solve the world's most complex scientific problems.

Unleashing autonomous AI agents: Why Kubernetes needs a new standard for agent execution

Unleashing autonomous AI agents: Why Kubernetes needs a new standard for agent execution

Joan Kallogjeri & Tomer Glottmann, GCP Google Kubernetes Engine

The arrival of autonomous AI Agents capable of reasoning, planning, and executing actions by generating their own code and interacting with the runtime environment marks a paradigm shift in how applications are built and operated. However, these new capabilities also introduce a fundamental security gap: how to safely allow agents to run untrusted, unverified generated code, perform actions and get access to data in runtime environments, especially pertaining to mission-critical infrastructure and environments that may have proprietary data.

We are excited to announce a major initiative within the Kubernetes community to address this exact challenge: we are launching Agent Sandbox as a formal subproject of Kubernetes SIG Apps, hosted under kubernetes-sigs/agent-sandbox.

This is more than just a tool; it is designed to standardize and evolve Kubernetes into the most secure and scalable platform for the agentic workloads.

The Latency Crisis for Interactive AI

Agent behavior often involves quick, iterative tool calls — checking a file, running a calculation, or querying an API. For security reasons, each of these calls requires its own isolated sandbox.

The challenge is that these sandboxes must be created from scratch, extremely quickly, to ensure isolated environments between executions. Because security and isolation are non-negotiable, the "spin-up" time becomes the critical bottleneck. If the secure execution environment takes too long to spin up, the entire agent application stalls, killing the interactive experience.

The Bottleneck of Massive Throughput

Enterprise platforms require infrastructure that can handle overwhelming scale. Users engaged in complex AI agent workloads demand support for up to tens of thousands of parallel sandboxes, processing thousands of queries per second. To meet this challenge, we are extending Kubernetes' proven capabilities of managing high-capacity, low latency applications, models and infrastructure to fit a growing class of single-instance workloads, like AI agent runtimes or dev environments, that require a lightweight, VM-like abstraction. A standardized, controller-based Sandbox API provides Kubernetes-native solution for these use cases, avoiding the workarounds required today, paving the way for the next generation of cloud-native AI applications.

The Agent Sandbox: A new Agent Standard for Kubernetes

To solve these problems, we are introducing a new, declarative resource focused strictly on the Sandbox primitive, designed from the ground up to be backend-agnostic.
The goal is to provide a persistent, isolated instance for single-container, stateful, singleton workloads, managed entirely through familiar Kubernetes constructs. The core APIs include:

• Sandbox : The core resource defining the agent sandbox workload for running an isolated instance of the agent's environment
• SandboxTemplate : Defines the secure blueprint of a sandbox archetype, including resource limits, base image, and initial security policies
• SandboxClaim : A transactional resource allowing users or higher-level frameworks (like ADK or LangChain) to request an execution environment, abstracting away the complex provisioning logic.

In addition to the Sandbox primitive we are also launching with additional features that make the experience as a whole better for the user:

• WarmPools — In order to support fast instance startup time, which is an important part of the usability of agenting sandboxes, we introduced the Warm Pool extension. The Sandbox Warm Pool Orchestrator utilizes a dedicated CRD to maintain a pool of pre-warmed pods, allowing the Sandbox Controller to claim a ready instance upon creation and reduce cold startup latency to less than one second.
• Shutdown Time — Since agentic behaviour can be unpredictable, this feature supports clean termination and cleanup of sandboxes, it automates the deletion by providing an absolute time for the sandbox to terminate.
• Python API/SDK — For better usability and a developer-friendly interface to programmatically interact with these CRDs, we provide an example SDK that abstracts away Kubernetes complexities with simple Pythonic functions.

The standard is designed to seamlessly support multiple isolation backends like gVisor and Kata Containers allowing developers to choose the technology that best fits their security and performance trade-offs.

The new Agent Sandbox features and implementations are available now in the github repo kubernetes-sigs/agent-sandbox and on our website agent-sandbox.sigs.k8s.io. We invite all developers, partners, and experts to join this critical community effort to define the secure scalable future of autonomous AI on Kubernetes.

We will be presenting a technical deep dive and officially launching the project at Kubecon Atlanta, November 2025. We hope to see you there!

Announcing Magika 1.0: now faster, smarter, and rebuilt in Rust

Announcing Magika 1.0: now faster, smarter, and rebuilt in Rust

Elie Bursztein, Julien Cretin, and Yanick Fratantonio, Applied Cybersecurity Research

Early last year, we open sourced Magika, Google's AI-powered file type detection system. Magika has seen great adoption by open source communities since that alpha release, with over one million monthly downloads. Today, we are happy to announce the release of Magika 1.0, a first stable version that introduces new features and a host of major improvements since last announcement. Here are the highlights:

• Expanded file type support for more than 200 types (up from ~100).
• A brand-new, high-performance engine rewritten from the ground up in Rust.
• A native Rust command-line client for maximum speed and security.
• Improved accuracy for challenging text-based formats like code and configuration files.
• A revamped Magika Python and TypeScript module for even easier integrations.

Smarter Detection: Doubling Down on File Types

Magika 1.0 now identifies more than 200 content types, doubling the number of file-types supported from the initial release. This isn't just about a bigger number; it unlocks far more granular and useful identification, especially for specialized, modern file types.

Some of the notable new file types detected include:

• Data Science & ML: We've added support for formats such as Jupyter Notebooks (ipynb), Numpy arrays (npy, npz), PyTorch models (pytorch), ONNX (onnx) files, Apache Parquet (parquet), and HDF5 (h5).
• Modern Programming & Web: The model now recognizes dozens of languages and frameworks. Key additions include Swift (swift), Kotlin (kotlin), TypeScript (typescript), Dart (dart), Solidity (solidity), Web Assembly (wasm), and Zig (zig).
• DevOps & Configuration: We've expanded detection for critical infrastructure and build files, such as Dockerfiles (dockerfile), TOML (toml), HashiCorp HCL (hcl), Bazel (bazel) build files, and YARA (yara) rules.
• Databases & Graphics: We also added support for common formats like SQLite (sqlite) databases, AutoCAD (dwg, dxf) drawings, Adobe Photoshop (psd) files, and modern web fonts (woff, woff2).
• Enhanced Granularity: Magika is now smarter at differentiating similar formats that might have been grouped together. For example, it can now distinguish:
  • JSONL (jsonl) vs. generic JSON (json)
  • TSV (tsv) vs. CSV (csv)
  • Apple binary plists (applebplist) from regular XML plists (appleplist)
  • C++ (cpp) vs. C (c)
  • JavaScript (javascript) vs. TypeScript(typescript)

Expanding Magika's detection capabilities introduced two significant technical hurdles: data volume and data scarcity.

First, the scale of the data required for training was a key consideration. Our training dataset grew to over 3TB when uncompressed, which required an efficient processing pipeline. To handle this, we leveraged our recently released SedPack dataset library. This tool allows us to stream and decompress this large dataset directly to memory during training, bypassing potential I/O bottlenecks and making the process feasible.

Second, while common file types are plentiful, many of the new, specialized, or legacy formats presented a data scarcity challenge. It is often not feasible to find thousands of real-world samples for every file type. To overcome this, we turned to generative AI. We leveraged Gemini to create a high-quality, synthetic training set by translating existing code and other structured files from one format to another. This technique, combined with advanced data augmentation, allowed us to build a robust training set, ensuring Magika performs reliably even on file types for which public samples are not readily available.

The complete list of all 200+ supported file types is available in our revamped documentation.

Under the Hood: A High-Performance Rust Engine

We completely rewrote Magika's core in Rust to provide native, fast, and memory-safe content identification. This engine is at the heart of the new Magika native command line tool that can safely scan hundreds of files per second.

Output of the new Magika Rust based command line tool

Magika is able to identify hundreds of files per second on a single core and easily scale to thousands per second on modern multi-core CPUs thanks to the use of the high-performance ONNX Runtime for model inference and Tokio for asynchronous parallel processing, For example, as visible in the chart below, on a MacBook Pro (M4), Magika processes nearly 1,000 files per second.

Getting Started

Ready to try it out? Getting started with the native command-line client is as simple as typing a single command line:

• On Linux and MacOS: curl -LsSf https://securityresearch.google/magika/install.sh | sh
• On Windows (PowerShell): powershell -ExecutionPolicy ByPass -c "irm https://securityresearch.google/magika/install.ps1 | iex"

Alternatively, the new Rust command-line client is also included in the magika python package, which you can install with: pipx install magika.

For developers looking to integrate Magika as a library into their own applications in Python, JavaScript/TypeScript, Rust, or other languages, head over to our comprehensive developer documentation to get started.

What's next

We're incredibly excited to see what you will build using Magika's enhanced file detection capabilities.

We invite you to join the community:

• Try Magika: Install it and run it on your files, or try it out in our web demo.
• Integrate Magika into your software: Visit our documentation to get started.
• Give us a star on GitHub to show your support.
• Report issues or suggest new file types you'd like to see by opening a feature request.
• Contribute new features and bindings by opening a pull request.

Thank you to everyone who has contributed, provided feedback, and used Magika over the past year. We can't wait to see what the future holds.

Acknowledgements

Magika's continued success was made possible by the help and support of many people, including: Ange Albertini, Loua Farah, Francois Galilee, Giancarlo Metitieri, Alex Petit-Bianco, Kurt Thomas, Luca Invernizzi, Lenin Simicich, and Amanda Walker.