Tinker API Implementation on TPUs, 30B models and up to 20k context

What it actually takes to run RL on a 27-billion-parameter language model on Google TPUs — the parallelism choices, the serving loop, and the memory engineering that stands between a working short-context run and a long-context one.

Sachin Konan · Princeton University

The reinforcement-learning algorithm here is simple: sample several answers to a problem, score each one, and raise the model's probability of producing the answers that scored well. Running it on TPUs is less simple. Sampling and training have different computational profiles — generation is bound by memory bandwidth and runs best on a dedicated inference server, while training needs gradients and optimizer state resident — and a TPU slice gives a fixed amount of memory per chip that cannot be extended.

This post documents building that system. It begins with SkyRL, an open RL framework written for GPUs, and ends with a 27-billion-parameter model training at a 22,528-token context on a TPU v5p slice — 44× the context I started from. The last third covers memory engineering, which is where most of the difficulty turned out to be.

1. The machine

A TPU pod slice is sold by its number of TensorCores, which is neither the number of chips nor the number of physical machines. The distinction matters because every layout decision later in this post follows from it.

Three properties matter for what follows. A host always has four chips, so a larger slice means more machines rather than denser ones. Each chip carries 95 GiB of high-bandwidth memory (HBM), which is roughly an H100's memory and still, as Section 10 shows, insufficient at long context. And each chip's two TensorCores appear to software as a single device: a v5p-16 ("16 cores") is 8 chips across 2 hosts, and JAX reports 4 devices per host.

2. Why I needed a new backend

SkyRL ships a JAX training backend that implements the model forward pass, the optimizer, and the sampling loop directly. It is readable and self-contained, and it was the right starting point. But it is a reimplementation, and on TPUs a reimplementation competes with a great deal of work that Google has already done and tuned:

• MaxText — reference LLM implementations written for TPUs, with the sharding annotations, attention kernels, and rematerialization policies already chosen and benchmarked.
• tunix — a post-training library that wraps those models with LoRA (via qwix), mesh construction, and checkpoint resharding.
• tokamax — a kernel library underneath, supplying the fused attention primitives (MaxText can route its splash-attention scheduler through it).

The kernels I would have had to write

The concrete thing MaxText supplies is TPU kernels. Attention on a TPU is not a matrix multiply you write in JAX and hope the compiler handles; the fast implementations are hand-written in Pallas, JAX's low-level kernel language, and they are what separates a working model from a fast one.

Attention runs through a hand-written Pallas kernel on both halves of the stack, and neither is code I wrote: splash attention on the trainer through MaxText, ragged paged attention on the samplers through vLLM's TPU backend — where it is not a flag but the attention path itself. That pair is most of the reason this stack is built on MaxText and vLLM rather than on anything I would have had to write.

The two inactive rows are training-side, and the ragged one is not a flag away. The backend does build a padding-aware mask, but the MaxText path discards it — MaxText applies causal masking internally, and segment_ids is accepted by the adapter and never forwarded. The kernel therefore sees a causal mask over the full length and has no way to know which positions are filler; padded positions are dropped from the loss afterwards, once the compute has already happened. Section 10 explains what it would take to reclaim that.

Sampling kernel: tpu_inference/kernels/ragged_paged_attention/v3/, dispatched from layers/common/attention_interface.py.

The device mesh

The other thing I inherited is tunix's way of describing where a tensor lives. A device mesh is a named grid of chips: instead of writing "shard this matrix across four devices," you declare axes — fsdp, tensor, data — give the mesh a shape over the physical chips, and annotate each tensor with which of its dimensions map to which axis. The compiler derives every collective from those annotations.

This is why Section 3's choice is a one-line change rather than a rewrite, and why the "rows must be a multiple of the chip count" rule exists at all: the batch dimension is mapped to a mesh axis of size four, so it has to divide by four. tunix also carries the idea further than I use it — a mesh per role, so a trainer and a sampler can occupy different chips with weights resharded between them, which is the principled version of the trainer/sampler split I built by hand.

Writing my own forward pass would mean re-earning all of that by hand, per model family. So I kept SkyRL's interface — the Tinker API server, which speaks forward_backward, optim_step, and sample — and replaced the engine underneath it with one that drives MaxText models through tunix. The RL client never knows the difference.

Code: skyrl/backends/tunix_backend.py (~2,500 lines), plugged in via --backend tunix.

3. How the model is split across chips

Three ways to spread a model over four chips, and the vocabulary matters because the choice determines the shape of every tensor downstream.

I use FSDP, and no tensor parallelism. Sharding the weights keeps 40 GB per chip free for activations. Skipping TP avoids a communication step inside every layer, and TP earns its keep only when a model cannot fit otherwise — a 27B model at bf16 (54 GB) fits in 95 GiB with room to spare.

I verified this rather than assuming it — an instrumented model load reported 53.8 GB of parameters globally and 13.5 GB resident per chip, exactly one quarter. It is worth confirming: a configuration slip that silently replicates weights costs 40 GB per chip and shows up much later as an unexplained failure.

4. The Tunix backend

Why not write against Tunix directly?

Tunix is a library: the natural way to use it is to import it and build a trainer inside its abstractions, in JAX. I wanted an API boundary in between instead, for three reasons.

LoRA is enough for this. An RL update rides a very thin signal. Supervised learning supplies information on the order of the number of tokens in an episode; a policy gradient supplies O(1) — an entire rollout collapsing to a single advantage number. Thinking Machines' analysis makes that argument and finds LoRA matching full fine-tuning on policy-gradient RL at ranks as low as 1, on two conditions: that the adapter is not capacity-constrained, and that it covers the MLP layers rather than attention alone. Every run here trains a rank-32 adapter over attention and MLP and never touches the base weights — far above what the signal can carry, and what makes a 27B model trainable on a single host at all.

The internals stay switchable. When the algorithm talks to an API rather than to a trainer, what sits underneath becomes configuration. I can develop and debug against my own TPU server and run the same code against the hosted Tinker service when a run needs to be larger than my slice — a different base URL, not a different program. Section 9 is that property used as a test.

Many adapters over one base. The direction I want to explore next is multi-agent, and LoRA makes that cheap in the right way: a server holds many adapters against a single set of base weights (eight by default here), so several policies can be served and trained at once without several copies of a 27B model. Multiple bases would not fit; multiple adapters already do. That is an extension of the interface rather than a rewrite of it.

All of it was reachable because SkyRL had already reimplemented the Tinker API — the HTTP server, the futures, the request batching and barrier scheduling, the weight-sync semantics. That is a narrow contract:

Writing a backend beneath that contract, rather than a trainer beside Tunix, meant inheriting all of the above and staying compatible with every client already written against it. Tunix is still here — qwix LoRA, mesh construction, the MaxText plumbing all sit inside the backend. It moved from being the framework to being a dependency behind an interface.

What the contract buys is visible in how little an algorithm has to say. This is a complete GRPO step:

sampling_client = await training_client.save_weights_and_get_sampling_client_async()

results = await asyncio.gather(*[
    sampling_client.sample_async(prompt=p, num_samples=group_size, sampling_params=sp)
    for p in prompts
])

datums = []
for result, prompt in zip(results, prompts):
    rewards    = [grade(seq) for seq in result.sequences]
    baseline   = sum(rewards) / len(rewards)
    advantages = [r - baseline for r in rewards]      # <- GRPO, in full
    if all(a == 0.0 for a in advantages):
        continue                                      # no signal in this group
    for seq, adv in zip(result.sequences, advantages):
        datums.append(tinker.Datum(
            model_input   = prompt.append(EncodedTextChunk(tokens=seq.tokens[:-1])),
            loss_fn_inputs = {"target_tokens": ..., "logprobs": seq.logprobs,
                              "advantages": ...},
        ))

await training_client.forward_backward_async(datums, loss_fn="importance_sampling")
await training_client.optim_step_async(adam_params)

Two lines are the algorithm. Nothing mentions a device, a mesh, a shard, or a compiler: advantages is a list of Python floats, and switching to PPO is loss_fn="ppo". Everything this post is about — sharding, packing, tiled cross-entropy, compiled shapes — happens on the far side of forward_backward, where the client cannot see it and does not have to.

That separation is what made the check in Section 9 possible at all: an unmodified published recipe, pointed at a different URL.

What the backend has to do

The backend's job is to make MaxText models honor those verbs. A few decisions shaped it:

• Load MaxText directly through pyconfig, bypassing tunix's model-name registry, so any forked or unreleased model config can be loaded by path. Converted checkpoints are cached (Section 6).
• Respect TPU shape rules. Sequence lengths round up to the 512-token attention block; batches pad to a multiple of the chip count. Both are handled inside the backend so the client can send whatever it has.
• Train adapters, not the model. LoRA (rank 32) via qwix, targeting attention and MLP projections. Only adapter gradients exist, which is what makes a 27B model trainable on one host at all.

5. Sampling: vLLM, and inflight LoRA

An RL step spends most of its wall clock generating tokens, and generation is a different discipline from training — continuous batching, paged attention, careful KV-cache management. Rather than write that, I run vLLM on the sampling hosts.

The subtlety is that in RL the policy changes every step, so the sampler must be updated constantly. Restarting a server per step would dominate the budget. Instead each save_weights_for_sampler exports a standard LoRA adapter directory and hot-swaps it into the running vLLM under a versioned name — load the new adapter, unload the previous, keep the base weights and KV cache warm. The sampler tracks the trainer within one step, and that is measurable: the divergence between the policy that generated a token and the policy being trained stayed between 1×10^-4 and 5.3×10^-4 across a 180-step run.

How the sampler splits the model

The trainer uses FSDP. The sampler does not: vLLM runs tensor parallel across the four chips of its host, and the launcher derives that from the topology rather than being told:

The trainer and the sampler make opposite choices on identical hardware, and for the sampler the reason is the KV cache. Tensor parallelism splits the weights, and what that buys is room:

A full replica on a single chip would spend 54 of its 102 GB on weights before one sequence arrived. Splitting them four ways leaves nearly the whole host for cached attention state, which is the constraint that actually binds when a step is dozens of concurrent 20k-token rollouts. Decode is bandwidth-bound rather than arithmetic-bound, so the same split also means each chip reads a quarter of the weights per token; the activations exchanged per layer are one token wide.

That is the shape of the sampler numbers in Section 7: about 97 GB of HBM reserved against 1.7% of the TensorCores busy — memory reserved and bytes moved, not arithmetic performed.

Adapter path: tpu/vllm_tpu_server.py and a fork of tpu-inference pinned at skyrl/v0.23.0-lora, so the runtime-LoRA hooks are source-installed rather than patched at deploy time.

6. Caching, because the machine will be taken away

I run on preemptible (spot) slices: cheap, and reclaimed without warning. In practice mine lasted between one hour and two days. Any run longer than that has to survive losing its machine, which turns each expensive artifact into the same question: if this slice disappears, what should not have to be computed again?

The last row is the one people forget. An RL step is sample, then train, and sampling is the long half. Preemption during the training half would otherwise throw away an hour of finished generation. Caching the graded rollouts mid-step means a replacement machine reloads them and goes straight to the gradient. With all five caches warm, a fresh slice goes from nothing to training in about 25 minutes, and the run resumes at the step it lost rather than the step it last checkpointed.

7. The training loop: synchronous, then pipelined

The first working version was the obvious one. Sample the whole batch, wait, train on it, wait, repeat. It worked, and it wasted half the machine: while the sampling hosts generate, the training host idles, and vice versa.

The pipelined version streams finished minibatches to the trainer as they are graded. Training begins as soon as the first wave lands, and the sampling of later waves hides inside the gradient math of earlier ones. Measured on a v5p-16 with Qwen3-8B, 1024 rollouts per step:

Inside a 210-second step: gradient waves take ~130 s and pace everything (sampling's ~95 s hides entirely inside that window), adapter export and load take ~45 s of pure I/O, and the optimizer step itself takes 0.1 s. The proportions are the useful part: the gradient update is essentially free, and nearly all of a step is spent moving data and waiting on generation.

Utilization: nothing here is compute-bound

A sidecar samples each chip's counters during training and logs them alongside the run. The numbers say where the remaining performance is.

Two different bottlenecks. Decode is memory-bound by nature — generating one token reads the whole model and does almost no arithmetic with it — so 1.7% TensorCore is close to expected, and vLLM's job is to keep enough sequences in flight to hide it. The trainer's 15–21% is the more interesting figure: it is not FLOP-limited but pipeline-limited, waiting on requests, results, and adapter I/O. That is what the pipelined loop attacks, and why the next unclaimed gains are also structural — one-step-off-policy sampling, and an in-place adapter reload to cut the 45-second sync.

Parameters that actually matter

8. Making Qwen3.5 exist in MaxText

I wanted two strong open models: Gemma-4-31B and Qwen3.5-27B. Gemma was straightforward — MaxText already ships a reference implementation, so it was a matter of adapter plumbing. Qwen3.5 was not implemented at all, and its architecture is unusual enough to be interesting: a hybrid stack where three of every four layers use linear attention (a gated DeltaNet) and every fourth uses full attention, with a 248,320-token vocabulary.

I added it in a MaxText fork: the dense-MLP branch in the decoder, a model config, and a HuggingFace→orbax converter mapping. Most of the effort then went to a subtler problem — proving the port was numerically correct. If the trainer's model disagrees with the sampler's, RL still runs; it just quietly optimizes the wrong objective.

A three-way log-probability comparison (the MaxText port, the HuggingFace reference, and the vLLM sampler), together with a per-layer bisection, found three bugs. They failed in different ways, which is why all three are listed here.

The payoff is measurable: at step 0 the sampler and trainer disagreed by 1.7×10^-4 nats — the cleanest parity of any model I brought up, and the evidence that the port was right before a single gradient was taken.

9. Does it learn?

I validated on mathematical reasoning with verifiable rewards: the model answers competition-style problems, correct answers are rewarded, and accuracy is tracked on a held-out set.

First Qwen3-8B from a cold start, 180 steps on a v5p-16 — held-out accuracy went 1.2% → 77.4%. That run also survived a preemption at step 60 and resumed on a fresh machine with optimizer state intact, which was the real test of Section 6.

Then Qwen3.5-27B, same recipe, same hardware:

Against a published reference

The client in all of this is not my code. It is the Tinker Cookbook math-RL recipe run unmodified — the same module, the same defaults — with two arguments changed: model_name, and a base_url pointing at my server instead of the hosted API. The training code never learns it is talking to my TPUs.

Thinking Machines publishes a result for that recipe, which gives an outside number to check against:

python -m tinker_cookbook.recipes.math_rl.train env=math \
  model_name="Qwen/Qwen3.5-9B" group_size=16 groups_per_batch=64 \
  learning_rate=2e-5 max_tokens=512

# after 180 steps:  "test/env/all/correct": 0.838

Learning rate (2e-5), token budget (512), LoRA rank (32), loss function and step count are identical across all three rows. The rest is not controlled, and I would not read it as a head-to-head: the middle row is a different model generation, and the bottom row is a larger model trained on a quarter of the batch — two differences that push in opposite directions. What the table does establish is that the same recipe pointed at my backend lands in the same band as the published number rather than somewhere anomalous.

10. The 22,528-token wall

Everything above happens at a 512-token generation budget. That is enough for math answers and nowhere near enough for the actual target: a model that reasons at length — thinking for thousands of tokens before writing a program that gets executed and scored.

Raising the context from 512 to 22,528 tokens broke the system in five distinct ways. This is the part worth reading, because the failures were not "out of memory, use less" — they were structural, and each fix is a small piece of engineering with a clean idea in it.

Where the memory actually goes

Start with the arithmetic. A training step must, for every token, compare the model's predicted distribution against the token that was actually sampled. The naive way materializes the full distribution:

Qwen3.5's vocabulary of 248,320 is what makes this severe. The logits tensor grows linearly in context length and is multiplied by a quarter of a million — at 20k tokens it alone exceeds what the chip has left, before counting the activations of 64 transformer layers or the copy that computing a log-softmax would produce.

1 · Score log-probabilities on the training chips

RL needs the log-probability of tokens that were already generated. The natural place to get them is the inference server — but vLLM's TPU backend crashes when asked to return them for long prompts. So I moved that computation onto the training host, where the model already lives, jitted it, and gathered only the target token's value inside the compiled program rather than returning a full distribution. The scoring path also has to pad its batch to the chip count, exactly like training.

2 · Pack micro-batches by token budget, not sequence count

"Eight sequences per batch" is meaningless when sequences range from 1,000 to 20,000 tokens. I pack by total tokens instead:

Memory then depends on a quantity I control, rather than on whatever lengths the model happened to produce.

3 · Never materialize a log-softmax

The standard way to get a token's log-probability is log_softmax(logits)[target], which allocates a second tensor the size of the first. The identity

saves the entire tensor. The two sides are equal by definition, but they cost different amounts: the left-hand side produces a vector with one entry per vocabulary token — 248,320 numbers, of which exactly one is ever read — while the right-hand side is a gather (x[i]) minus a scalar reduction. The vocabulary axis is summed away rather than written out, and XLA fuses the reduction into the preceding operation, so the intermediate never lands in memory at all.

The numerics are unchanged. A stable logsumexp is computed as m + log Σ exp(x − m) with m = max(x) — the identical shift that log_softmax applies internally. Same arithmetic, same floating-point result, one fewer tensor. This one line is what made a KL-regularized run fit at all.

4 · Fused linear cross-entropy: the logits never exist

The identity in (3) removes the copy, but the logits themselves are still built. The real fix is to never form them: split the tokens into tiles, and for each tile project hidden states to vocabulary, take the one number needed, and throw the tile away.

This is where a subtlety cost me three attempts. Tiling the forward pass is easy; the problem is the backward pass. If you write the loop with a standard scan — even wrapped in gradient checkpointing — automatic differentiation saves every tile's logits so it can reuse them, stacks them back together, and reconstructs the exact 81 GB tensor you were avoiding. The loop looks tiled and is not.

The fix is a custom gradient: it tells the compiler that the forward pass keeps only the hidden states, and that the backward pass should recompute one tile's logits at a time. Peak memory is then one tile in both directions, enforced by the gradient definition itself rather than left to the compiler's discretion.

Code: _flce_target_logprobs in skyrl/backends/tunix_backend.py — a jax.custom_vjp whose forward residual is deliberately the hidden tiles, never the logits.

5 · Keep the compiler to one shape

XLA compiles a separate program for every tensor shape it is given, and each program reserves its own block of memory to work in. Feed it many shapes and those blocks pile up — tens of gigabytes each — until a new shape has nowhere left to load. Measured on this workload, the reserved block grows at roughly

I tried the clever approach first — bucketing sequences into a ladder of lengths to waste less padding — and it backfired: padding dropped from 60% to 22%, but six live memory blocks replaced one, and the run failed the moment a seventh shape appeared. What worked was the simpler approach, and the one MaxText and tunix follow by convention: pad everything to a single shape. One program, one reserved block, for the whole run. Roughly 60% of the tokens in a tile are then padding — real compute, spent on nothing — and it is still the better trade, because the alternative does not run at all.

6 · The leak that hid behind all of it

With one shape and fused cross-entropy, a training step still died — deterministically, around the 55th tile of 92, always with the same "cannot allocate" message. Memory grew ~170 MB per tile, with nothing in my code to explain it.

Instrumentation settled it. Logging every live array on the chip every ten tiles showed watched one shape's count climb: [4, 16384, 5120], one more copy per tile, forever. Those are hidden states — and MaxText's decoder sows them into the model object as a debugging convenience. Under JAX's functional-object wrapper, changes a model makes to itself during a compiled call persist after it. Ninety-two calls, ninety-two retained copies. The fix is one line, executed after each tile:

nnx.pop(model, nnx.Intermediate)   # discard sown intermediates

Memory went flat at 29.1 GB across all 92 tiles, and the first end-to-end training step completed. Two related fixes landed with it: gradient accumulation moved inside the compiled function with its buffer donated (the previous version kept three copies of the gradients alive), and the out-of-memory handler learned to retry gently before evicting the compiled program — the old handler turned a single transient failure into an unrecoverable spiral by discarding the very program it needed to reload.

11. Cookbook

Everything in this post runs from one repository. The recipes below create and tear down their own slices, so the only prerequisites are credentials and capacity.

git clone --recurse-submodules https://github.com/SachinKonan/SkyRLTpu.git
cd SkyRLTpu

gcloud auth login
gcloud config set project <your-gcp-project>

export ZONE=<zone-with-v5p-capacity>

export TUNIX_MAXTEXT_CKPT_CACHE_GCS=gs://<your-bucket>/skyrl-maxtext-ckpts
export HF_CACHE_GCS=gs://<your-bucket>/hf-cache
export VLLM_XLA_CACHE_GCS=gs://<your-bucket>/vllm-xla-cache

./tpu/autoresume_run.sh tpu/runs/qwen35-27b.env

MODEL_NAME=Qwen/Qwen3.5-27B
TUNIX_MAXTEXT_MODEL_NAME=qwen3.5-27b
TUNIX_MAXTEXT_PIP_SPEC="maxtext @ git+https://github.com/SachinKonan/maxtext.git@skyrl/qwen35-dense"
TRAIN_WORKERS=0            # host 0 trains
VLLM_WORKERS=1             # host 1 samples
TRAIN_MICRO_BATCH_SIZE=8
GROUP_SIZE=8  GROUPS_PER_BATCH=32   # 256 rollouts / step
MAX_TOKENS=512  LORA_RANK=32  LEARNING_RATE=2e-5
STREAM_NUM_MINIBATCHES=8   # pipelined overlap
RENDERER_NAME=qwen3_5_disable_thinking
MAX_STEPS=180  SAVE_EVERY=20  EVAL_EVERY=20

TINKER_BACKEND=tunix  MODEL_NAME=Qwen/Qwen3.5-27B \
TRAIN_WORKERS=0  VLLM_WORKERS=1,2,3 \
TUNIX_UNIFORM_SEQ_LEN=20480 \    # one program shape, ever  (fix 5)
TUNIX_TRAIN_TOKEN_BUDGET=81920 \ # 4 sequences x 20480      (fix 2)
TUNIX_FLCE_TILE_SIZE=2048 \      # fused linear CE tile     (fix 4)
TUNIX_MAXTEXT_KWARGS='{"num_vocab_tiling": 8}' \
  ./tpu/start_colocated_vllm_tinker.sh

EXPERIMENT_NAME=erdos-qwen35  TTD_ENV=erdos_min_overlap \
PHASE1_MAX_TOKENS=15872  CONTEXT_WINDOW=20480  TTD_TRAIN_MAX_SEQ=20480 \
GROUP_SIZE=16  GROUPS_PER_BATCH=32  KL_PENALTY_COEF=0 \
  bash tpu/run_ttd_on_tpu_host.sh

12. What next

Two directions are open, and they are hard for different reasons. One is a serving problem I partly solved and do not fully trust; the other is a capacity problem I have not attempted.

Mixture-of-experts, and where adapters attach

Everything above trains dense models. The obvious way to get more capability per unit of compute is a mixture-of-experts model, where each token is routed to a few of many expert networks. I built the path for gpt-oss-20b and it works, but it required breaking an assumption that holds everywhere else in this stack.

The assumption is that a LoRA adapter can be attached to a layer. In a dense model, vLLM wraps each target layer so the adapter is applied at runtime — which is what makes hot-swapping cheap, and is the mechanism Section 5 depends on. In an MoE model on TPU, vLLM fuses all experts into one large tensor and runs them as a single batched operation, because doing otherwise would be far slower. There is no per-expert layer left to wrap: the abstraction the adapter needs has been optimized away.

It is worth the trouble because adapting attention alone is not equivalent. The LoRA result in Section 4 holds only when the adapter covers the MLP layers — and in a mixture-of-experts model, the MLP is the experts. Skipping them is precisely the capacity-constrained case the analysis warns about, so an MoE model trained through attention only is not the same experiment.

The resolution is to split the exported adapter in two. The attention projections ship as a standard PEFT directory and take vLLM's native path. The expert factors ship as a sidecar and are merged into the fused tensor at load time, through a custom collective RPC — apply_moe_lora_deltas — that hands each worker its factors and adds them in place. Getting there also meant disabling vLLM's own MoE-LoRA wrapper on TPU, which fails to boot against bf16 experts.

Code: tpu/vllm_tpu_server.py plus the tpu-inference fork (merge extraction, donated in-place adds, equivalence tests against incremental merges).

The instability I have not solved

This path runs, and it is not yet trustworthy. Sampling occasionally returns NaN log-probabilities for individual tokens. Downstream that is catastrophic rather than merely noisy, because the RL loss exponentiates a difference of log-probabilities:

The current defense is three layers deep, and every layer is containment rather than a cure. NaNs are clamped to a sentinel value at the source in the fork; any token carrying that sentinel has its loss weight zeroed and its value neutralized (a zero weight is not enough — inf × 0 is NaN inside a compiled function); and non-finite optimizer updates are skipped outright. The gradient is protected. What I still cannot say is why the fused MoE path produces NaN log-probabilities in the first place — whether it is an overflow in the bf16 expert matmuls, the router, or the merge itself. Until that is answered I would not run a long MoE job unattended.

Larger models: is tensor parallelism the answer?

Nothing here has been run on a model larger than 27B on my own hardware. (A 120B model appears in my results, but those runs used hosted infrastructure, not this stack.) The arithmetic for attempting it is simple enough to state:

So tensor parallelism is probably not what I need. TP earns its place when a single tensor is too large for one chip; here no individual tensor is the problem, total capacity is — and capacity is what sharding across more chips buys. The natural move for 120B is multi-host FSDP over two or four hosts, not a different parallelism strategy within one.

That is also the honest description of the blocker: my backend trains on a single host by construction, and refuses more. Multi-host is not a missing flag but a missing mechanism — every process must execute the same compiled program in lockstep, which means a coordinator broadcasting each operation to workers. SkyRL's original JAX backend already implements exactly that pattern, so the work is porting it rather than inventing it.

I did try tensor parallelism, out of curiosity and because it would have shrunk the memory footprint of a training tile. Two configurations, two different failures: splitting four chips as 2×2 left 0.7 GB free after loading the model, against 27–29 GB under plain FSDP, for reasons I never isolated; and a 4-way tensor split failed inside the attention kernel, which requires its mask blocks to divide evenly across devices — the same divisibility rule from Section 3, reappearing one level down. Both are recorded as measurements rather than conclusions.

Built on SkyRL, MaxText, tunix, and vLLM. Measurements come from runs on v5p-16 and v5p-32 slices; the model-parallelism, memory, and timing figures quoted here were instrumented on those runs rather than estimated.