The JAX trainer backend

The trainer fleet has two interchangeable backends. The default is PyTorch DTensor/FSDP2; the alternative is a JAX backend built on Flax NNX, jax.sharding, and optax. Both sit below the same models/policy/interfaces.py ABCs, so the GRPO loop, the rewards, the sandbox, and the inference fleet are unchanged regardless of which trainer runs. Only the trainer fleet differs; the inference fleet stays on vLLM and the sandbox on the executor.

The backend is selected per config (policy.jax_cfg.enabled); see examples/configs/grpo_swe_jax.yaml. The implementation lives in models/jax/.

Why a second backend

JAX/XLA brings GSPMD sharding (mesh + PartitionSpec) and a different compilation and numerics story to the same training objective. Keeping it below the policy interface means it is a drop-in: the refit seam emits the same HF weight names the inference engine expects, so a JAX-trained policy serves through the identical vLLM path.

What it implements

The backend reproduces the trainer’s responsibilities in JAX:

  • Models — Flax NNX modules with HF-weight loaders and an exact name map: Qwen3 dense, the MoE path (mirroring models/dtensor/moe/), the hybrid linear-attention path (Gated-DeltaNet) for the Qwen3-Next lineage, and Gemma 4 (gemma4_text).

  • Loss and log-probs — a pure-JAX mirror of the clipped-PG loss and the log-prob gather, validated for value- and grad-parity against the torch loss.

  • Train step — nnx.split + jax.value_and_grad, microbatch grad accumulation, an optax optimizer chain (AdamW + scheduler + global-norm clip), and the aux-loss-free load-balance bias stepped via a pre-optimizer hook.

  • Refit / checkpoint — JAX→torch weight conversion (dlpack on GPU, numpy on CPU) through the shared name map, including the MoE per-expert re-expansion; Orbax checkpoints plus an HF-safetensors export so a JAX run’s weights reload through the torch loader.

Build phases

The re-platform was built and validated in phases (J0–J11), recorded in handoff/jax-trainer-replatform-plan.md and handoff/jax-build-progress.md: scaffolding and the config flag (J0); dense Qwen3 with HF-logit parity (J1); mesh/PartitionSpec sharding (J2); loss, train step, refit, log-prob and checkpoint seams (J3–J7); the Ray/driver glue (J9); MoE (J8); and hybrid linear/hybrid attention for Qwen3-Next (J11).

Gemma 4

Gemma 4 (gemma4_text) is supported natively, selected automatically from the HF model_type. Its compute path diverges from Qwen3 in several ways the module reproduces exactly (verified against the HF reference logits):

  • per-layer-type attention — a 5:1 sliding/full interleave where sliding layers use the base head_dim with full-rotary RoPE and global layers a larger global_head_dim with proportional partial-rotary RoPE (the non-rotated frequencies are left as identity / NoPE);

  • q/k/v RMSNorm with a scale-free v_norm and attention scaling fixed at 1.0, plus attention_k_eq_v (global layers reuse the K projection as V) and KV-layer sharing (the shared tail layers carry no K/V weights and consume the K/V published by the last non-shared layer of their type);

  • per-layer input embeddings (PLE) — a packed table gathered and projected into every decoder layer;

  • additive MoE — a routed branch (a Gemma-specific router plus fused gelu-tanh experts) summed with the dense MLP rather than replacing it; and

  • sandwich norms, final_logit_softcapping, and sqrt(hidden) embedding scaling.

Because several of its modules hold more than one parameter (the router’s scale vectors, the experts’ fused gate/up and down stacks) and its per-layer embedding table nests one level deeper than the HF checkpoint, Gemma 4 carries its own refit name-map rather than the generic dense one — see weight synchronization. Unlike the Qwen3-MoE refit, its experts stay fused, matching the layout the Gemma 4 checkpoint and the inference engine use; a load→refit round-trip reproduces the HF state dict. Expert-parallel Gemma 4 MoE is hardware-gated alongside the other multi-device paths.

Validation posture

The entire CPU-validatable surface — dense → MoE → hybrid, plus loss/grad, train-step, refit, checkpoint, and load-balance — is proven on CPU against HF reference logits, with a clean type-check. Anything that genuinely needs a GPU or multiple devices (the sharded multi-GPU train step, the live NCCL refit into vLLM, the expert-parallel ragged_all_to_all collective, bf16 numerics) is GPU-gated and tracked, with its exact bring-up check, in handoff/hardware-deferred-validation.md. GPU validation is currently parked by project decision; the ledger means resuming is “run the harness,” not re-derive.