KerasHub + JAX + TPU: Exploring a New Stack for LLM Post-Training Experiments

A Technical Signal Behind One Tweet

A developer recently shared their experience on social media: running extensive LLM post-training experiments using a combination of KerasHub, JAX, and TPU — and hinting that some of the work might be open-sourced. While brief, this post reflects a noteworthy trend in the toolchain choices being made in modern large model engineering.

Post-training is the critical bridge between a "completed pre-trained model" and one that's actually usable. It encompasses techniques like supervised fine-tuning (SFT), instruction alignment, RLHF, and DPO. Compared to pre-training — which consumes enormous resources — post-training is far more accessible to individual researchers and smaller engineering teams, making it one of the most actively experimented-on areas in recent years.

Why KerasHub + JAX + TPU

This isn't a random combination. Each component plays a distinct and deliberate role.

KerasHub: The Entry Point for Models and Workflows

KerasHub is the Keras ecosystem's library for pretrained models, offering a wide range of ready-to-use model architectures and weights. For post-training experiments, it dramatically lowers the barrier to "grab a base model and start modifying it." Developers don't need to build network architectures from scratch, freeing them to focus on training strategies and data engineering.

JAX: The Foundation for High-Performance Numerical Computing

JAX is known for its functional programming paradigm, automatic differentiation, and just-in-time (JIT) compilation. It excels in large-scale parallel computing scenarios, pairs naturally with TPU, and lets developers express complex training logic in concise code while achieving near-hardware-limit performance. For post-training experiments that require rapid iteration and trying different training recipes, JAX's flexibility is especially valuable.

JAX's core design philosophy combines a NumPy-style API with the XLA (Accelerated Linear Algebra) compiler. Its three core transformations — jit (just-in-time compilation), vmap (vectorized mapping), and pmap (parallel mapping) — allow developers to express computation in a style close to mathematical notation, with the compiler handling low-level optimization and device scheduling automatically. Unlike PyTorch's imperative execution model, JAX uses a functional pure-transformation model with constrained side effects, making cross-device and cross-node parallelism easier to reason about and optimize. In post-training scenarios, this means researchers can rapidly experiment with different loss function formulations (e.g., DPO loss, SimPO loss) by swapping out a few lines of pure functional code, without needing to refactor the entire training loop.

TPU: The Compute Backbone for Large-Scale Training

TPUs (Tensor Processing Units) are accelerators purpose-built for deep learning, offering strong cost-efficiency for the matrix-heavy workloads of large model training. The deep integration between JAX and TPU makes scaling distributed training relatively smooth — one of the key reasons some research teams opt for this path over the mainstream GPU + PyTorch approach.

Why Post-Training Experiments Matter

An increasing share of capability differentiation among today's large models comes from the post-training stage. As base models converge and become more homogeneous, the real competitive battleground is how well a model can be fine-tuned, aligned, and preference-optimized for specific tasks.

The developer's mention of "a wide range of experiments" implies that the ability to rapidly switch between training strategies and compare results within a single toolchain is critical. The KerasHub + JAX + TPU combination strikes a balance between experimental iteration speed and scalability — well suited for this kind of broad trial-and-error research.

Post-training techniques currently fall into three broad categories. First, supervised fine-tuning (SFT), which directly optimizes model output style using high-quality instruction-response pairs. Second, reinforcement learning from human feedback (RLHF), which trains a reward model and then optimizes the policy model using algorithms like PPO — effective, but engineering-intensive. Third, direct preference optimization (DPO) and its variants (SimPO, ORPO, etc.), which bypass the explicit reward model and optimize directly from preference data pairs, offering a cleaner implementation with competitive results. Because each approach suits different use cases, running "a wide range of experiments" typically means rapidly switching between these strategies on the same base model and comparing outcomes — placing high demands on toolchain flexibility.

The Potential Value of Open-Sourcing

The most intriguing part of the post is "might open-source some of that." Post-training toolchains built around PyTorch — such as TRL and Axolotl — are already quite mature, while practical resources for post-training with JAX/Keras remain comparatively scarce.

If these experimental codebases are open-sourced, they would provide invaluable reference templates for developers working within the JAX + TPU stack and fill a meaningful gap in the ecosystem. For teams looking to reduce dependence on a single framework and explore alternative technical paths, this kind of hands-on work would be especially valuable.

The PyTorch post-training ecosystem is already well-structured: TRL (Transformer Reinforcement Learning), maintained by Hugging Face, covers mainstream workflows including SFT, DPO, and PPO; Axolotl simplifies fine-tuning experiment reproducibility through a configuration-driven approach; and LLaMA-Factory offers an all-in-one solution widely used in China. By contrast, the JAX/Keras ecosystem lacks a post-training framework at a comparable level of abstraction — developers typically have to assemble their own training loops, data pipelines, and evaluation logic. This is both an ecosystem gap and an opportunity for meaningful contribution. A battle-tested JAX post-training template would have direct practical value for teams doing research on Google Cloud TPU or through the TPU Research Cloud.

A Measured View: Still an Early Signal

It's worth noting that this post itself is light on specifics — no concrete experimental results, no performance benchmarks, and no clear open-source timeline. It's more of a directional signal than a reproducible artifact.

For readers following large model engineering practices, it's worth keeping an eye on post-training explorations like this within the JAX ecosystem. They represent a viable alternative to the dominant PyTorch path, and serve as a reminder that when it comes to large model toolchain choices, there's more than one road worth taking.