Restore production sharding and warm-start support

README.md

@@ -229,7 +229,11 @@ The example uses `datasets/train/foundation_5000.jsonl`, writes
 parameters through JSON. The command writes the trained model at the end of
 the run.
 
-To skip burn-in using compatible post-burn-in sampler states:
+Set the optional top-level `checkpoint` field to start from a full router-model
+checkpoint. This loads model parameters only; KFAC state, sampler state, the
+step counter, and the learning-rate schedule start fresh.
+
+To load compatible sampler states:
 
 ```bash
 hamiltonzero train examples/train.json --reuse-mcmc path/to/mcmc-states
@@ -237,7 +241,9 @@ hamiltonzero train examples/train.json --reuse-mcmc path/to/mcmc-states
 
 For multisystem training, the path is a directory containing
 `<system-index>.eqx` files. For a one-system training panel it may be a single
-file.
+file. Training runs `mcmc.burn_in` iterations after either fresh initialization
+or loading reused states. Each iteration uses `mcmc.burn_in_replica_steps`
+MCMC moves.
 
 ## Fine-tune
 

src/hamiltonzero/compiled/trunk.py

@@ -127,3 +127,31 @@ def compile_shared_trunk(model, ctx) -> SharedTrunk:
         real_mask=ctx.mask,
         balanced_mask=ctx.bmask,
     )
+
+
+def compile_shared_trunk_from_kernel(kernel: TrunkCompilerKernel, ctx) -> SharedTrunk:
+    edge_feat, local_feat, global_feat = kernel.featurizer(
+        ctx.J_double_prime,
+        ctx.mask,
+        ctx.h_prime,
+    )
+    g_seed = tree_sphere(global_feat.astype(local_feat.dtype))
+    node_raw, edge_raw, g_seed = kernel.trunk(
+        ctx,
+        edge_feat,
+        local_feat,
+        g_seed,
+    )
+    global_stream = kernel.shared_global(
+        g_seed.astype(edge_raw.dtype),
+        edge_raw,
+        ctx.mask,
+    )
+    return SharedTrunk(
+        node_raw=node_raw,
+        edge_raw=edge_raw,
+        global_raw=global_feat,
+        global_stream=global_stream,
+        real_mask=ctx.mask,
+        balanced_mask=ctx.bmask,
+    )

src/hamiltonzero/config.py

@@ -151,6 +151,7 @@ class TrainConfig:
     steps: int
     seed: int = 777
     n_max: int = 64
+    checkpoint: Path | None = None
     model: ModelConfig = field(default_factory=ModelConfig)
     router: RouterConfig = field(default_factory=RouterConfig)
     mcmc: MCMCConfig = field(default_factory=MCMCConfig)

src/hamiltonzero/evaluation/runtime.py

@@ -40,6 +40,7 @@ from hamiltonzero.data.systems import (
     load_system as load_spin_hamiltonian,
 )
 from hamiltonzero.energy import vmc_energy_custom_lap_compiled
+from hamiltonzero.energy.custom_lap import build_W_levels
 from hamiltonzero.energy.frame import compile_energy_frame, route_energy_inputs
 from hamiltonzero.mcmc.runtime import (
     adapt_batched,
@@ -212,7 +213,6 @@ def _step_single(
     )
 
 
-@partial(jax.jit, static_argnames=("replica_steps", "walker_chunk_size"))
 def _step_compiled_rows(
     state,
     model,
@@ -249,7 +249,6 @@ def _adapt_single(
     )
 
 
-@jax.jit
 def _adapt_rows(
     state,
     beta_history_weight,
@@ -269,27 +268,110 @@ def _adapt_rows(
 
 
 def _compiled_energy_single(kernel, tree, frame, q, *, chunk_size: int):
-    values = vmc_energy_custom_lap_compiled(
+    return vmc_energy_custom_lap_compiled(
         kernel,
         tree,
         frame,
         q,
         chunk_size=chunk_size,
     )
-    return tuple(jnp.expand_dims(value, axis=0) for value in values)
 
 
-@partial(jax.jit, static_argnames=("chunk_size",))
+def _build_compiled_energy_single(mesh: Mesh, chunk_size: int):
+    batch = P("batch", None, None)
+    batch_output = P("batch")
+
+    def local_energy(kernel, tree, frame, q):
+        outputs = _compiled_energy_single(
+            kernel,
+            tree,
+            frame,
+            q,
+            chunk_size=int(chunk_size),
+        )
+        local_count = jnp.asarray(q.shape[0], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "batch")
+        guard = global_count.astype(jnp.float32) * jnp.asarray(0.0, jnp.float32)
+        return tuple(value + guard.astype(value.dtype) for value in outputs)
+
+    mapped = jax.shard_map(
+        local_energy,
+        mesh=mesh,
+        in_specs=(P(), P(), P(), batch),
+        out_specs=(batch_output,) * 4,
+        check_vma=False,
+    )
+    replicated_sharding = NamedSharding(mesh, P())
+    batch_sharding = NamedSharding(mesh, batch)
+    batch_output_sharding = NamedSharding(mesh, batch_output)
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            replicated_sharding,
+            replicated_sharding,
+            replicated_sharding,
+            batch_sharding,
+        ),
+        out_shardings=(batch_output_sharding,) * 4,
+    )
+
+
 def _compiled_energy_rows(kernel, trees, frames, q, *, chunk_size: int):
-    return jax.vmap(
-        lambda tree, frame, q_row: vmc_energy_custom_lap_compiled(
+    def energy_row(tree, frame, q_row):
+        n_sites = int(q_row.shape[-2])
+        frame = eqx.tree_at(
+            lambda value: value.w_levels,
+            frame,
+            tuple(build_W_levels(frame.custom_lap_J_eff, n_sites)),
+        )
+        return vmc_energy_custom_lap_compiled(
             kernel,
             tree,
             frame,
             q_row,
             chunk_size=chunk_size,
         )
-    )(trees, frames, q)
+
+    return jax.vmap(energy_row)(trees, frames, q)
+
+
+def _build_compiled_energy_rows(mesh: Mesh, chunk_size: int):
+    systems = P("systems")
+    system_batch = P("systems", None)
+
+    def local_energy(kernel, trees, frames, q):
+        outputs = _compiled_energy_rows(
+            kernel,
+            trees,
+            frames,
+            q,
+            chunk_size=int(chunk_size),
+        )
+        local_count = jnp.asarray(q.shape[0] * q.shape[1], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "systems")
+        guard = global_count.astype(jnp.float32) * jnp.asarray(0.0, jnp.float32)
+        return tuple(value + guard.astype(value.dtype) for value in outputs)
+
+    mapped = jax.shard_map(
+        local_energy,
+        mesh=mesh,
+        in_specs=(P(), systems, systems, system_batch),
+        out_specs=(system_batch,) * 4,
+        check_vma=False,
+    )
+    replicated_sharding = NamedSharding(mesh, P())
+    systems_sharding = NamedSharding(mesh, systems)
+    system_batch_sharding = NamedSharding(mesh, system_batch)
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            replicated_sharding,
+            systems_sharding,
+            systems_sharding,
+            system_batch_sharding,
+        ),
+        out_shardings=(system_batch_sharding,) * 4,
+    )
 
 
 _compile_single = jax.jit(compile_wavefunction)
@@ -335,6 +417,131 @@ def _single_state_sharding(mesh: Mesh, state):
     )
 
 
+def _single_state_specs(state):
+    walkers = P("batch")
+    replicated = P()
+    return type(state)(
+        q=walkers,
+        log_p=walkers,
+        grad_log_p=walkers,
+        beta=replicated,
+        sigma=replicated,
+        step=replicated,
+        key=walkers,
+        n_local_accept=walkers,
+        n_local=walkers,
+        n_swap_accept=walkers,
+        n_swap=walkers,
+        mask=replicated,
+        m=replicated,
+        n_haar_accept=walkers,
+        n_haar=walkers,
+    )
+
+
+def _row_state_specs(state):
+    systems = P("systems")
+    return jax.tree_util.tree_map(lambda _value: systems, state)
+
+
+def _row_model_specs(model):
+    return CompiledWaveFunctions(
+        kernel=jax.tree_util.tree_map(lambda _value: P(), model.kernel),
+        trees=jax.tree_util.tree_map(lambda _value: P("systems"), model.trees),
+    )
+
+
+def _row_model_sharding(mesh: Mesh, model):
+    return CompiledWaveFunctions(
+        kernel=_replicated_sharding(mesh, model.kernel),
+        trees=_system_sharding(mesh, model.trees),
+    )
+
+
+def _build_rows_mcmc_step(
+    mesh: Mesh,
+    state,
+    model,
+    context,
+    *,
+    replica_steps: int,
+    walker_chunk_size: int,
+):
+    state_specs = _row_state_specs(state)
+    model_specs = _row_model_specs(model)
+    context_specs = jax.tree_util.tree_map(lambda _value: P("systems"), context)
+    state_sharding = _system_sharding(mesh, state)
+    model_sharding = _row_model_sharding(mesh, model)
+    context_sharding = _system_sharding(mesh, context)
+
+    def local_step(state_local, model_local, context_local):
+        out = _step_compiled_rows(
+            state_local,
+            model_local,
+            context_local,
+            replica_steps=int(replica_steps),
+            walker_chunk_size=int(walker_chunk_size),
+        )
+        local_count = jnp.asarray(out.q.shape[0] * out.q.shape[1], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "systems")
+        guard = global_count.astype(out.q.dtype) * jnp.asarray(0.0, out.q.dtype)
+        return eqx.tree_at(lambda value: value.q, out, out.q + guard)
+
+    mapped = jax.shard_map(
+        local_step,
+        mesh=mesh,
+        in_specs=(state_specs, model_specs, context_specs),
+        out_specs=state_specs,
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(state_sharding, model_sharding, context_sharding),
+        out_shardings=state_sharding,
+        donate_argnums=(0,),
+    )
+
+
+def _build_singular_mcmc_step(
+    mesh: Mesh,
+    state,
+    state_sharding,
+    model_sharding,
+    context_sharding,
+    *,
+    replica_steps: int,
+    walker_chunk_size: int,
+):
+    state_specs = _single_state_specs(state)
+
+    def local_step(state_local, model_local, context_local):
+        out = _step_single(
+            state_local,
+            model_local,
+            context_local,
+            replica_steps=int(replica_steps),
+            walker_chunk_size=int(walker_chunk_size),
+        )
+        local_count = jnp.asarray(out.q.shape[0], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "batch")
+        guard = global_count.astype(out.q.dtype) * jnp.asarray(0.0, out.q.dtype)
+        return eqx.tree_at(lambda value: value.q, out, out.q + guard)
+
+    mapped = jax.shard_map(
+        local_step,
+        mesh=mesh,
+        in_specs=(state_specs, P(), P()),
+        out_specs=state_specs,
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(state_sharding, model_sharding, context_sharding),
+        out_shardings=state_sharding,
+        donate_argnums=(0,),
+    )
+
+
 def _system_sharding(mesh: Mesh, value):
     return jax.tree_util.tree_map(
         lambda array: NamedSharding(
@@ -440,6 +647,8 @@ class DefaultEvalBackend:
         self._energy_frames: dict[int, Any] = {}
         self._context_meshes: dict[int, Mesh] = {}
         self._contest_mesh: Mesh | None = None
+        self._contest_step_entries: dict[tuple[int, int], Any] = {}
+        self._contest_adapt_entry: Any | None = None
         self._singular_mesh: Mesh | None = None
         self._singular_state_sharding: Any | None = None
         self._singular_model_sharding: Any | None = None
@@ -447,6 +656,7 @@ class DefaultEvalBackend:
         self._singular_step_entries: dict[tuple[int, int], Any] = {}
         self._singular_adapt_entry: Any | None = None
         self._singular_energy_entries: dict[int, Any] = {}
+        self._contest_energy_entries: dict[int, Any] = {}
 
     def build_system(self, system, energy: EnergyConfig):
         context, energy_inputs = build_context_and_energy(
@@ -598,6 +808,9 @@ class DefaultEvalBackend:
         self._energy_frames[id(routed)] = frames
         self._context_meshes[id(routed)] = mesh
         self._contest_mesh = mesh
+        self._contest_step_entries.clear()
+        self._contest_adapt_entry = None
+        self._contest_energy_entries.clear()
         return routed
 
     def release_context(self, context) -> None:
@@ -605,6 +818,9 @@ class DefaultEvalBackend:
         self._energy_frames.pop(id(context), None)
         self._context_meshes.pop(id(context), None)
         self._contest_mesh = None
+        self._contest_step_entries.clear()
+        self._contest_adapt_entry = None
+        self._contest_energy_entries.clear()
 
     def beam_candidates(
         self,
@@ -850,31 +1066,35 @@ class DefaultEvalBackend:
             key = (int(replica_steps), int(walker_chunk_size))
             step = self._singular_step_entries.get(key)
             if step is None:
-                step = jax.jit(
-                    partial(
-                        _step_single,
-                        replica_steps=key[0],
-                        walker_chunk_size=key[1],
-                    ),
-                    in_shardings=(
-                        self._singular_state_sharding,
-                        self._singular_model_sharding,
-                        self._singular_context_sharding,
-                    ),
-                    out_shardings=self._singular_state_sharding,
-                    donate_argnums=(0,),
+                step = _build_singular_mcmc_step(
+                    self._singular_mesh,
+                    state,
+                    self._singular_state_sharding,
+                    self._singular_model_sharding,
+                    self._singular_context_sharding,
+                    replica_steps=key[0],
+                    walker_chunk_size=key[1],
                 )
                 self._singular_step_entries[key] = step
             return step(state, model, context)
         if not isinstance(model, CompiledWaveFunctions):
             raise TypeError("multirow eval MCMC requires compiled wavefunctions")
-        return _step_compiled_rows(
-            state,
-            model,
-            context,
-            replica_steps=int(replica_steps),
-            walker_chunk_size=int(walker_chunk_size),
-        )
+        mesh = self._context_meshes.get(id(context))
+        if mesh is None:
+            raise RuntimeError("contest MCMC mesh is unavailable")
+        key = (int(replica_steps), int(walker_chunk_size))
+        step = self._contest_step_entries.get(key)
+        if step is None:
+            step = _build_rows_mcmc_step(
+                mesh,
+                state,
+                model,
+                context,
+                replica_steps=key[0],
+                walker_chunk_size=key[1],
+            )
+            self._contest_step_entries[key] = step
+        return step(state, model, context)
 
     def adapt_mcmc(self, state, config: EvalMCMCConfig):
         arguments = (
@@ -885,7 +1105,23 @@ class DefaultEvalBackend:
             jnp.asarray(config.haar_target_acceptance, dtype=jnp.float32),
         )
         if state.q.ndim == 5:
-            return _adapt_rows(*arguments)
+            if self._contest_mesh is None:
+                raise RuntimeError("contest adaptation mesh is unavailable")
+            if self._contest_adapt_entry is None:
+                replicated = NamedSharding(self._contest_mesh, P())
+                state_sharding = _system_sharding(self._contest_mesh, state)
+                self._contest_adapt_entry = jax.jit(
+                    _adapt_rows,
+                    in_shardings=(
+                        state_sharding,
+                        replicated,
+                        replicated,
+                        replicated,
+                        replicated,
+                    ),
+                    out_shardings=state_sharding,
+                )
+            return self._contest_adapt_entry(*arguments)
         if self._singular_mesh is None or self._singular_state_sharding is None:
             raise RuntimeError("singular eval placement has not been prepared")
         if self._singular_adapt_entry is None:
@@ -969,12 +1205,18 @@ class DefaultEvalBackend:
         chunk_size = int(config.chunk_size)
         if isinstance(model, CompiledWaveFunctions):
             frames = self._frames(context)
-            return _compiled_energy_rows(
+            mesh = self._context_meshes.get(id(context))
+            if mesh is None:
+                raise RuntimeError("contest energy mesh is unavailable")
+            energy = self._contest_energy_entries.get(chunk_size)
+            if energy is None:
+                energy = _build_compiled_energy_rows(mesh, chunk_size)
+                self._contest_energy_entries[chunk_size] = energy
+            return energy(
                 model.kernel,
                 model.trees,
                 frames,
                 q,
-                chunk_size=chunk_size,
             )
         if not isinstance(model, CompiledWaveFunction):
             raise TypeError("singular eval energy requires a compiled wavefunction")
@@ -988,11 +1230,6 @@ class DefaultEvalBackend:
             self._singular_mesh,
             P("batch", None, None),
         )
-        energy_sharding = NamedSharding(
-            self._singular_mesh,
-            P(None, "batch"),
-        )
-        output_shardings = (energy_sharding,) * 4
         q = jax.device_put(q, q_sharding)
         frames = self._frames(context)
         frame_sharding = _replicated_sharding(
@@ -1002,26 +1239,18 @@ class DefaultEvalBackend:
         frames = jax.device_put(frames, frame_sharding)
         energy = self._singular_energy_entries.get(chunk_size)
         if energy is None:
-            energy = jax.jit(
-                partial(
-                    _compiled_energy_single,
-                    chunk_size=chunk_size,
-                ),
-                in_shardings=(
-                    self._singular_model_sharding.kernel,
-                    self._singular_model_sharding.tree,
-                    frame_sharding,
-                    q_sharding,
-                ),
-                out_shardings=output_shardings,
+            energy = _build_compiled_energy_single(
+                self._singular_mesh,
+                chunk_size,
             )
             self._singular_energy_entries[chunk_size] = energy
-        return energy(
+        outputs = energy(
             model.kernel,
             model.tree,
             frames,
             q,
         )
+        return tuple(jnp.expand_dims(value, axis=0) for value in outputs)
 
     def block_until_ready(self, value) -> None:
         jax.block_until_ready(value)

src/hamiltonzero/modes/finetune.py

@@ -80,21 +80,14 @@ def _burn_in(
     context,
     state: REState,
     config: FineTuneConfig,
+    step_mcmc,
+    adapt,
 ) -> REState:
-    replica_steps = config.mcmc.burn_in_replica_steps
-    step_mcmc = eqx.filter_jit(
-        functools.partial(
-            run_batched,
-            n_steps=replica_steps,
-            walker_chunk_size=config.mcmc.walker_chunk_size,
-        )
-    )
-    adapt = eqx.filter_jit(functools.partial(_adapt, config=config))
     for iteration in range(config.mcmc.burn_in):
         state = step_mcmc(
+            state,
             model,
             context,
-            state,
         )
         if iteration > 0 and iteration % config.mcmc.adapt_every == 0:
             state = adapt(state)
@@ -109,9 +102,35 @@ def _replicate(value, sharding: NamedSharding):
 
 
 def _place_state(state: REState, mesh: Mesh) -> REState:
+    return jax.device_put(state, _state_sharding(mesh))
+
+
+def _state_specs() -> REState:
+    walkers = P("batch")
+    replicated = P()
+    return REState(
+        q=walkers,
+        log_p=walkers,
+        grad_log_p=walkers,
+        beta=replicated,
+        sigma=replicated,
+        step=replicated,
+        key=walkers,
+        n_local_accept=walkers,
+        n_local=walkers,
+        n_swap_accept=walkers,
+        n_swap=walkers,
+        mask=replicated,
+        m=replicated,
+        n_haar_accept=walkers,
+        n_haar=walkers,
+    )
+
+
+def _state_sharding(mesh: Mesh) -> REState:
     replicated = NamedSharding(mesh, P())
     batched = NamedSharding(mesh, P("batch"))
-    shardings = REState(
+    return REState(
         q=batched,
         log_p=batched,
         grad_log_p=batched,
@@ -128,7 +147,83 @@ def _place_state(state: REState, mesh: Mesh) -> REState:
         n_haar_accept=batched,
         n_haar=batched,
     )
-    return jax.device_put(state, shardings)
+
+
+def _build_mcmc_entry(
+    mesh: Mesh,
+    state_sharding,
+    model_sharding,
+    context_sharding,
+    *,
+    replica_steps: int,
+    walker_chunk_size: int | None,
+):
+    specs = _state_specs()
+
+    def local_step(state, model, context):
+        out = run_batched(
+            model,
+            context,
+            state,
+            n_steps=int(replica_steps),
+            walker_chunk_size=walker_chunk_size,
+        )
+        local_count = jnp.asarray(out.q.shape[0], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "batch")
+        guard = global_count.astype(out.q.dtype) * jnp.asarray(0.0, out.q.dtype)
+        return eqx.tree_at(lambda value: value.q, out, out.q + guard)
+
+    mapped = jax.shard_map(
+        local_step,
+        mesh=mesh,
+        in_specs=(specs, P(), P()),
+        out_specs=specs,
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(state_sharding, model_sharding, context_sharding),
+        out_shardings=state_sharding,
+        donate_argnums=(0,),
+    )
+
+
+def _build_energy_entry(
+    mesh: Mesh,
+    model_sharding,
+    frame_sharding,
+    *,
+    chunk_size: int,
+):
+    q_spec = P("batch", None, None)
+    output_spec = P("batch")
+
+    def local_energy(model, frame, q):
+        outputs = vmc_energy_custom_lap_finetune(
+            model,
+            frame,
+            q,
+            chunk_size=int(chunk_size),
+        )
+        local_count = jnp.asarray(q.shape[0], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "batch")
+        guard = global_count.astype(jnp.float32) * jnp.asarray(0.0, jnp.float32)
+        return tuple(value + guard.astype(value.dtype) for value in outputs)
+
+    mapped = jax.shard_map(
+        local_energy,
+        mesh=mesh,
+        in_specs=(P(), P(), q_spec),
+        out_specs=(output_spec,) * 4,
+        check_vma=False,
+    )
+    q_sharding = NamedSharding(mesh, q_spec)
+    output_sharding = NamedSharding(mesh, output_spec)
+    return jax.jit(
+        mapped,
+        in_shardings=(model_sharding, frame_sharding, q_sharding),
+        out_shardings=(output_sharding,) * 4,
+    )
 
 
 def _metric(
@@ -175,10 +270,8 @@ def run_finetune(
         initial_m=config.mcmc.initial_haar_sites,
         initial_sigma=config.mcmc.initial_sigma,
     )
-    reused = False
     if config.mcmc.reuse_mcmc is not None:
         state = load_mcmc(config.mcmc.reuse_mcmc, state)
-        reused = True
     key, _route_key = jax.random.split(key)
     freeze_route = eqx.filter_jit(
         functools.partial(
@@ -222,6 +315,10 @@ def run_finetune(
         )
     mesh = Mesh(np.asarray(devices, dtype=object), ("batch",))
     replicated = NamedSharding(mesh, P())
+    state_sharding = _state_sharding(mesh)
+    model_sharding = jax.tree_util.tree_map(lambda _value: replicated, model)
+    context_sharding = jax.tree_util.tree_map(lambda _value: replicated, context)
+    frame_sharding = jax.tree_util.tree_map(lambda _value: replicated, energy_frame)
     model = _replicate(model, replicated)
     context = _replicate(context, replicated)
     energy_frame = _replicate(energy_frame, replicated)
@@ -234,6 +331,42 @@ def run_finetune(
         jnp.zeros((1, config.mcmc.batch_size), dtype=jnp.complex64),
         kfac_data,
     )
+    mcmc_entries = {}
+
+    def mcmc_entry(replica_steps: int):
+        entry_key = (int(replica_steps), config.mcmc.walker_chunk_size)
+        entry = mcmc_entries.get(entry_key)
+        if entry is None:
+            entry = _build_mcmc_entry(
+                mesh,
+                state_sharding,
+                model_sharding,
+                context_sharding,
+                replica_steps=entry_key[0],
+                walker_chunk_size=entry_key[1],
+            )
+            mcmc_entries[entry_key] = entry
+        return entry
+
+    adapt = jax.jit(
+        functools.partial(_adapt, config=config),
+        in_shardings=(state_sharding,),
+        out_shardings=state_sharding,
+    )
+    local_energy = _build_energy_entry(
+        mesh,
+        model_sharding,
+        frame_sharding,
+        chunk_size=config.energy.chunk_size,
+    )
+    target_entry = jax.jit(
+        functools.partial(
+            process_finetune_targets,
+            mad_width=config.kfac.mad_clip_width,
+        ),
+        in_shardings=(kfac_data, replicated),
+        out_shardings=kfac_data,
+    )
     kfac = init_finetune_kfac_state(
         config.kfac,
         model,
@@ -244,30 +377,23 @@ def run_finetune(
         key=jax.random.fold_in(key, 0xCAFE),
         multi_device=mesh.size > 1,
     )
-    if not reused:
-        state = _burn_in(model, context, state, config)
-    step_mcmc = eqx.filter_jit(
-        functools.partial(
-            run_batched,
-            n_steps=config.mcmc.steps,
-            walker_chunk_size=config.mcmc.walker_chunk_size,
-        )
-    )
-    adapt = eqx.filter_jit(functools.partial(_adapt, config=config))
-    local_energy = eqx.filter_jit(
-        functools.partial(
-            vmc_energy_custom_lap_finetune,
-            chunk_size=config.energy.chunk_size,
-        )
+    state = _burn_in(
+        model,
+        context,
+        state,
+        config,
+        mcmc_entry(config.mcmc.burn_in_replica_steps),
+        adapt,
     )
+    step_mcmc = mcmc_entry(config.mcmc.steps)
     run_started = time.perf_counter()
     last_metric = None
     for step in range(config.steps):
         step_started = time.perf_counter()
         state = step_mcmc(
+            state,
             model,
             context,
-            state,
         )
         if step > 0 and step % config.mcmc.adapt_every == 0:
             state = adapt(state)
@@ -277,10 +403,9 @@ def run_finetune(
             energy_frame,
             q_cold,
         )
-        target = process_finetune_targets(
+        target = target_entry(
             total[None],
             context_batch.s_norm,
-            mad_width=config.kfac.mad_clip_width,
         )
         key, key_kfac = jax.random.split(key)
         model, kfac = apply_finetune_kfac_step(

src/hamiltonzero/modes/train.py

@@ -13,15 +13,23 @@ import jax.numpy as jnp
 import numpy as np
 from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 
-from hamiltonzero.checkpoint import load_mcmc, save_model
-from hamiltonzero.compiled.tree import compile_physical_tree_reference
-from hamiltonzero.compiled.trunk import bind_shared_kernel, compile_shared_trunk
+from hamiltonzero.checkpoint import load_mcmc, load_model, save_model
+from hamiltonzero.compiled.tree import (
+    bind_physical_compiler_kernel,
+    compile_physical_tree_from_shared_trunk,
+)
+from hamiltonzero.compiled.trunk import (
+    bind_shared_kernel,
+    bind_trunk_compiler_kernel,
+    compile_shared_trunk_from_kernel,
+)
 from hamiltonzero.compiled.types import (
     CompiledWaveFunction,
 )
 from hamiltonzero.config import TrainConfig
 from hamiltonzero.data import build_context_and_energy, load_systems
 from hamiltonzero.energy import vmc_energy_custom_lap_compiled
+from hamiltonzero.energy.custom_lap import build_W_levels
 from hamiltonzero.energy.frame import compile_energy_frame
 from hamiltonzero.mcmc import (
     REState,
@@ -83,25 +91,22 @@ def _identity_perms(n_max: int):
     )
 
 
-def _route_sharding(mesh: Mesh, value):
-    return jax.tree_util.tree_map(
-        lambda x: NamedSharding(
-            mesh,
-            P("systems", *([None] * (x.ndim - 1)))
-            if x.ndim and x.shape[0] == ROUTE_SAMPLES
-            else P(),
-        ),
-        value,
-    )
+def _systems_sharding(mesh: Mesh, value):
+    sharding = NamedSharding(mesh, P("systems"))
+    return jax.tree_util.tree_map(lambda _value: sharding, value)
+
+
+def _replicated_sharding(mesh: Mesh, value):
+    sharding = NamedSharding(mesh, P())
+    return jax.tree_util.tree_map(lambda _value: sharding, value)
 
 
 def _replicate(mesh: Mesh, value):
-    replicated = NamedSharding(mesh, P())
-    return jax.device_put(value, jax.tree_util.tree_map(lambda _: replicated, value))
+    return jax.device_put(value, _replicated_sharding(mesh, value))
 
 
 def _place_routes(mesh: Mesh, value):
-    return jax.device_put(value, _route_sharding(mesh, value))
+    return jax.device_put(value, _systems_sharding(mesh, value))
 
 
 def _host_pool(value):
@@ -132,34 +137,195 @@ def _activate_system(mesh: Mesh, state: _SystemState) -> _SystemState:
     )
 
 
-def _compile_trees(model, trunk, perms):
-    return jax.vmap(lambda perm: compile_physical_tree_reference(model, trunk, perm))(
-        perms
+def _abstract(value):
+    return jax.tree_util.tree_map(
+        lambda leaf: jax.ShapeDtypeStruct(leaf.shape, leaf.dtype),
+        value,
     )
 
 
-def _compile_frames(inputs, mask, bmask, perms):
-    return jax.vmap(
-        lambda perm: compile_energy_frame(
-            inputs,
-            mask,
-            bmask,
-            perm,
+def _owner_reduce(value):
+    owner = jax.lax.axis_index("systems") == 0
+    if jnp.issubdtype(value.dtype, jnp.bool_):
+        return jax.lax.pmax(jnp.where(owner, value, jnp.zeros_like(value)), "systems")
+    return jax.lax.psum(jnp.where(owner, value, jnp.zeros_like(value)), "systems")
+
+
+def _build_owner_entry(mesh: Mesh, function, templates):
+    output_template = jax.eval_shape(function, *_abstract(templates))
+    input_specs = jax.tree_util.tree_map(lambda _value: P(), templates)
+    output_specs = jax.tree_util.tree_map(lambda _value: P(), output_template)
+
+    def local(*values):
+        owner = jax.lax.axis_index("systems") == 0
+        output = jax.lax.cond(
+            owner,
+            lambda args: function(*args),
+            lambda _args: jax.tree_util.tree_map(
+                lambda value: jnp.zeros(value.shape, value.dtype),
+                output_template,
+            ),
+            values,
         )
-    )(perms)
+        return jax.tree_util.tree_map(_owner_reduce, output)
+
+    mapped = jax.shard_map(
+        local,
+        mesh=mesh,
+        in_specs=input_specs,
+        out_specs=output_specs,
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=_replicated_sharding(mesh, templates),
+        out_shardings=_replicated_sharding(mesh, output_template),
+    )
 
 
-def _run_routes(kernel, trees, contexts, state, n_steps, chunk_size):
-    def one(tree, context, sampler):
-        return run_batched(
-            CompiledWaveFunction(kernel=kernel, tree=tree),
-            context,
-            sampler,
-            n_steps,
-            walker_chunk_size=chunk_size,
+def _compile_tree_local(physical_kernel, trunk, perms):
+    tree = compile_physical_tree_from_shared_trunk(
+        physical_kernel,
+        trunk,
+        perms[0],
+    )
+    return jax.tree_util.tree_map(lambda value: value[None], tree)
+
+
+def _build_compile_trees(mesh: Mesh, physical_kernel, trunk, perms):
+    local_perms = jax.ShapeDtypeStruct((1, perms.shape[1]), perms.dtype)
+    local_output = jax.eval_shape(
+        _compile_tree_local,
+        _abstract(physical_kernel),
+        _abstract(trunk),
+        local_perms,
+    )
+    mapped = jax.shard_map(
+        _compile_tree_local,
+        mesh=mesh,
+        in_specs=(P(), P(), P("systems")),
+        out_specs=jax.tree_util.tree_map(lambda _value: P("systems"), local_output),
+        check_vma=False,
+    )
+    output_template = jax.eval_shape(
+        mapped,
+        _abstract(physical_kernel),
+        _abstract(trunk),
+        _abstract(perms),
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            _replicated_sharding(mesh, physical_kernel),
+            _replicated_sharding(mesh, trunk),
+            NamedSharding(mesh, P("systems")),
+        ),
+        out_shardings=_systems_sharding(mesh, output_template),
+    )
+
+
+def _compile_frame_local(inputs, mask, bmask, perms):
+    frame = compile_energy_frame(inputs, mask, bmask, perms[0])
+    return jax.tree_util.tree_map(lambda value: value[None], frame)
+
+
+def _build_compile_frames(mesh: Mesh, inputs, mask, bmask, perms):
+    local_perms = jax.ShapeDtypeStruct((1, perms.shape[1]), perms.dtype)
+    local_output = jax.eval_shape(
+        _compile_frame_local,
+        _abstract(inputs),
+        _abstract(mask),
+        _abstract(bmask),
+        local_perms,
+    )
+    mapped = jax.shard_map(
+        _compile_frame_local,
+        mesh=mesh,
+        in_specs=(P(), P(), P(), P("systems")),
+        out_specs=jax.tree_util.tree_map(lambda _value: P("systems"), local_output),
+        check_vma=False,
+    )
+    output_template = jax.eval_shape(
+        mapped,
+        _abstract(inputs),
+        _abstract(mask),
+        _abstract(bmask),
+        _abstract(perms),
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            _replicated_sharding(mesh, inputs),
+            NamedSharding(mesh, P()),
+            NamedSharding(mesh, P()),
+            NamedSharding(mesh, P("systems")),
+        ),
+        out_shardings=_systems_sharding(mesh, output_template),
+    )
+
+
+def _run_routes_local(state, kernel, trees, n_steps, chunk_size):
+    sampler = jax.tree_util.tree_map(lambda value: value[0], state)
+    tree = jax.tree_util.tree_map(lambda value: value[0], trees)
+    sampler = run_batched(
+        CompiledWaveFunction(kernel=kernel, tree=tree),
+        None,
+        sampler,
+        int(n_steps),
+        walker_chunk_size=chunk_size,
+    )
+    return jax.tree_util.tree_map(lambda value: value[None], sampler)
+
+
+def _build_run_routes(
+    mesh: Mesh,
+    state,
+    kernel,
+    trees,
+    *,
+    n_steps: int,
+    chunk_size: int | None,
+):
+    state_specs = jax.tree_util.tree_map(lambda _value: P("systems"), state)
+    tree_specs = jax.tree_util.tree_map(lambda _value: P("systems"), trees)
+
+    def local(state_value, kernel_value, trees_value):
+        output = _run_routes_local(
+            state_value,
+            kernel_value,
+            trees_value,
+            int(n_steps),
+            chunk_size,
+        )
+        local_count = jnp.asarray(
+            output.q.shape[0] * output.q.shape[1],
+            dtype=jnp.int32,
+        )
+        global_count = jax.lax.psum(local_count, "systems")
+        guard = global_count.astype(output.q.dtype) * jnp.asarray(0.0, output.q.dtype)
+        return eqx.tree_at(
+            lambda value: value.q,
+            output,
+            output.q + guard,
         )
 
-    return jax.vmap(one)(trees, contexts, state)
+    mapped = jax.shard_map(
+        local,
+        mesh=mesh,
+        in_specs=(state_specs, P(), tree_specs),
+        out_specs=state_specs,
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            _systems_sharding(mesh, state),
+            _replicated_sharding(mesh, kernel),
+            _systems_sharding(mesh, trees),
+        ),
+        out_shardings=_systems_sharding(mesh, state),
+        donate_argnums=(0,),
+    )
 
 
 def _adapt_routes(state, config):
@@ -176,6 +342,12 @@ def _adapt_routes(state, config):
 
 def _sampled_energy(kernel, trees, frames, q, chunk_size):
     def one(tree, frame, q_row):
+        n_sites = int(q_row.shape[-2])
+        frame = eqx.tree_at(
+            lambda value: value.w_levels,
+            frame,
+            tuple(build_W_levels(frame.custom_lap_J_eff, n_sites)),
+        )
         return vmc_energy_custom_lap_compiled(
             kernel,
             tree,
@@ -187,6 +359,102 @@ def _sampled_energy(kernel, trees, frames, q, chunk_size):
     return jax.vmap(one)(trees, frames, q)
 
 
+def _build_sampled_energy(mesh: Mesh, chunk_size: int):
+    systems = P("systems")
+    system_batch = P("systems", None)
+
+    def local_energy(kernel, trees, frames, q):
+        outputs = _sampled_energy(
+            kernel,
+            trees,
+            frames,
+            q,
+            int(chunk_size),
+        )
+        local_count = jnp.asarray(q.shape[0] * q.shape[1], dtype=jnp.int32)
+        global_count = jax.lax.psum(local_count, "systems")
+        guard = global_count.astype(jnp.float32) * jnp.asarray(0.0, jnp.float32)
+        return tuple(value + guard.astype(value.dtype) for value in outputs)
+
+    mapped = jax.shard_map(
+        local_energy,
+        mesh=mesh,
+        in_specs=(P(), systems, systems, system_batch),
+        out_specs=(system_batch,) * 4,
+        check_vma=False,
+    )
+    replicated_sharding = NamedSharding(mesh, P())
+    systems_sharding = NamedSharding(mesh, systems)
+    system_batch_sharding = NamedSharding(mesh, system_batch)
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            replicated_sharding,
+            systems_sharding,
+            systems_sharding,
+            system_batch_sharding,
+        ),
+        out_shardings=(system_batch_sharding,) * 4,
+    )
+
+
+def _build_reframe(mesh: Mesh, state, context, perms):
+    state_specs = jax.tree_util.tree_map(lambda _value: P("systems"), state)
+    context_specs = jax.tree_util.tree_map(lambda _value: P("systems"), context)
+
+    def local(state_value, context_value, old_perms, new_perms):
+        return reframe_state_context(
+            state_value,
+            context_value,
+            old_perms,
+            new_perms,
+        )
+
+    mapped = jax.shard_map(
+        local,
+        mesh=mesh,
+        in_specs=(
+            state_specs,
+            context_specs,
+            P("systems"),
+            P("systems"),
+        ),
+        out_specs=(state_specs, context_specs),
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            _systems_sharding(mesh, state),
+            _systems_sharding(mesh, context),
+            NamedSharding(mesh, P("systems")),
+            NamedSharding(mesh, P("systems")),
+        ),
+        out_shardings=(
+            _systems_sharding(mesh, state),
+            _systems_sharding(mesh, context),
+        ),
+    )
+
+
+def _build_rebase(mesh: Mesh, q, perms):
+    mapped = jax.shard_map(
+        rebase_cold_samples,
+        mesh=mesh,
+        in_specs=(P("systems"), P("systems")),
+        out_specs=P("systems"),
+        check_vma=False,
+    )
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            NamedSharding(mesh, P("systems")),
+            NamedSharding(mesh, P("systems")),
+        ),
+        out_shardings=NamedSharding(mesh, P("systems")),
+    )
+
+
 def _mode_energy(kernel, tree, frame, q_canonical, mode_perm, chunk_size):
     q_mode = jnp.take(q_canonical, mode_perm, axis=-2)
 
@@ -207,6 +475,92 @@ def _mode_energy(kernel, tree, frame, q_canonical, mode_perm, chunk_size):
     return jax.vmap(one)(q_mode)
 
 
+def _build_mode_energy(mesh: Mesh, kernel, tree, frame, q, mode_perm, chunk_size):
+    system_batch = P("systems", None)
+
+    def local(kernel_value, tree_value, frame_value, q_value, perm_value):
+        return _mode_energy(
+            kernel_value,
+            tree_value,
+            frame_value,
+            q_value,
+            perm_value,
+            int(chunk_size),
+        )
+
+    mapped = jax.shard_map(
+        local,
+        mesh=mesh,
+        in_specs=(P(), P(), P(), system_batch, P()),
+        out_specs=(system_batch,) * 5,
+        check_vma=False,
+    )
+    output_sharding = NamedSharding(mesh, system_batch)
+    return jax.jit(
+        mapped,
+        in_shardings=(
+            _replicated_sharding(mesh, kernel),
+            _replicated_sharding(mesh, tree),
+            _replicated_sharding(mesh, frame),
+            NamedSharding(mesh, system_batch),
+            NamedSharding(mesh, P()),
+        ),
+        out_shardings=(output_sharding,) * 5,
+    )
+
+
+def _compile_mode(physical_kernel, trunk, inputs, mask, bmask, perm):
+    tree = compile_physical_tree_from_shared_trunk(physical_kernel, trunk, perm)
+    frame = compile_energy_frame(inputs, mask, bmask, perm)
+    return tree, frame
+
+
+def _build_exact_skip(mesh: Mesh):
+    def local(sampled, mode):
+        equal = jnp.all(sampled.astype(jnp.int32) == mode[None].astype(jnp.int32))
+        return jax.lax.pmin(equal.astype(jnp.int32), "systems").astype(jnp.bool_)
+
+    mapped = jax.shard_map(
+        local,
+        mesh=mesh,
+        in_specs=(P("systems"), P()),
+        out_specs=P(),
+        check_vma=False,
+    )
+    replicated = NamedSharding(mesh, P())
+    return jax.jit(
+        mapped,
+        in_shardings=(NamedSharding(mesh, P("systems")), replicated),
+        out_shardings=replicated,
+    )
+
+
+def _initial_model(config: TrainConfig, key):
+    model = build_model(config.model, key, n_max=config.n_max)
+    if config.checkpoint is not None:
+        model = load_model(config.checkpoint, model)
+    return model
+
+
+def _run_route_burn_in(
+    sampler,
+    kernel,
+    trees,
+    config,
+    run_routes,
+    adapt,
+):
+    for iteration in range(config.mcmc.burn_in):
+        sampler = run_routes(
+            sampler,
+            kernel,
+            trees,
+        )
+        if iteration and iteration % config.mcmc.adapt_every == 0:
+            sampler = adapt(sampler)
+    return sampler
+
+
 def _initial_system_state(
     model,
     context,
@@ -218,7 +572,7 @@ def _initial_system_state(
     mesh,
     compile_plan,
     compile_trees,
-    run_routes,
+    get_run_routes,
 ):
     walkers = config.mcmc.batch_size // ROUTE_SAMPLES
     cpu = jax.devices("cpu")[0]
@@ -253,21 +607,28 @@ def _initial_system_state(
     sampler = _place_routes(mesh, sampler)
     contexts = _place_routes(mesh, contexts)
     perms = _place_routes(mesh, perms)
-    trunk = compile_plan(model, context)[0]
-    kernel = bind_shared_kernel(model)
-    trees = compile_trees(model, trunk, perms)
-    if config.mcmc.reuse_mcmc is None:
-        for iteration in range(config.mcmc.burn_in):
-            sampler = run_routes(
-                kernel,
-                trees,
-                contexts,
-                sampler,
-                config.mcmc.burn_in_replica_steps,
-                config.mcmc.walker_chunk_size,
-            )
-            if iteration and iteration % config.mcmc.adapt_every == 0:
-                sampler = _adapt_routes(sampler, config)
+    trunk = compile_plan(bind_trunk_compiler_kernel(model), context)
+    physical_kernel = bind_physical_compiler_kernel(model)
+    trees = compile_trees(physical_kernel, trunk, perms)
+    run_routes = get_run_routes(
+        sampler,
+        bind_shared_kernel(model),
+        trees,
+        config.mcmc.burn_in_replica_steps,
+    )
+    adapt = jax.jit(
+        lambda state: _adapt_routes(state, config),
+        in_shardings=(_systems_sharding(mesh, sampler),),
+        out_shardings=_systems_sharding(mesh, sampler),
+    )
+    sampler = _run_route_burn_in(
+        sampler,
+        bind_shared_kernel(model),
+        trees,
+        config,
+        run_routes,
+        adapt,
+    )
     return _SystemState(sampler=sampler, context=contexts, perms=perms)
 
 
@@ -290,6 +651,8 @@ def run_train(
 ) -> TrainResult:
     if config.mcmc.batch_size % ROUTE_SAMPLES:
         raise ValueError("mcmc.batch_size must be divisible by K=8")
+    if config.mcmc.burn_in < 0:
+        raise ValueError("mcmc.burn_in must be non-negative")
     systems = load_systems(config.systems)
     if not systems:
         raise ValueError("training requires at least one system")
@@ -308,29 +671,53 @@ def run_train(
     energy_inputs = [energy for _context, energy in systems_data]
     key = jax.random.PRNGKey(config.seed)
     key_model, key_mcmc = jax.random.split(key)
-    model = build_model(config.model, key_model, n_max=config.n_max)
+    model = _initial_model(config, key_model)
     devices = tuple(jax.devices())
-    if len(devices) < ROUTE_SAMPLES:
-        raise ValueError(
-            "learned-router train requires eight devices for the K=8 systems mesh"
-        )
-    mesh = Mesh(np.asarray(devices[:ROUTE_SAMPLES], dtype=object), ("systems",))
+    if len(devices) != ROUTE_SAMPLES:
+        raise ValueError("learned-router train requires exactly eight visible devices")
+    mesh = Mesh(np.asarray(devices, dtype=object), ("systems",))
     model = _replicate(mesh, model)
-    compile_plan = jax.jit(
-        lambda model_value, context_value: (
-            compile_shared_trunk(model_value, context_value),
-        )
+    trunk_kernel = bind_trunk_compiler_kernel(model)
+    compile_plan = _build_owner_entry(
+        mesh,
+        compile_shared_trunk_from_kernel,
+        (trunk_kernel, contexts[0]),
     )
-    compile_trees = jax.jit(_compile_trees)
-    compile_frames = jax.jit(_compile_frames)
-    run_routes = jax.jit(
-        _run_routes,
-        static_argnums=(4, 5),
-        donate_argnums=(3,),
+    trunk_template = compile_plan(trunk_kernel, contexts[0])
+    physical_kernel = bind_physical_compiler_kernel(model)
+    perms_template = _place_routes(mesh, _identity_perms(config.n_max))
+    compile_trees = _build_compile_trees(
+        mesh,
+        physical_kernel,
+        trunk_template,
+        perms_template,
     )
-    sampled_energy = jax.jit(_sampled_energy, static_argnums=(4,))
-    mode_energy = jax.jit(_mode_energy, static_argnums=(5,))
-    reframe = jax.jit(reframe_state_context, donate_argnums=(0, 1))
+    compile_frames = _build_compile_frames(
+        mesh,
+        energy_inputs[0],
+        contexts[0].mask,
+        contexts[0].bmask,
+        perms_template,
+    )
+    sampled_energy = _build_sampled_energy(mesh, config.energy.chunk_size)
+    exact_skip = _build_exact_skip(mesh)
+    mcmc_entries = {}
+
+    def get_run_routes(state, kernel, trees, n_steps):
+        entry_key = (int(n_steps), config.mcmc.walker_chunk_size)
+        entry = mcmc_entries.get(entry_key)
+        if entry is None:
+            entry = _build_run_routes(
+                mesh,
+                state,
+                kernel,
+                trees,
+                n_steps=entry_key[0],
+                chunk_size=entry_key[1],
+            )
+            mcmc_entries[entry_key] = entry
+        return entry
+
     system_states: list[_SystemState | None] = [None] * len(systems)
 
     def get_system(index: int):
@@ -346,13 +733,19 @@ def run_train(
                 mesh=mesh,
                 compile_plan=compile_plan,
                 compile_trees=compile_trees,
-                run_routes=run_routes,
+                get_run_routes=get_run_routes,
             )
         return _activate_system(mesh, cached)
 
     first = get_system(0)
     q_seed = jax.vmap(cold_samples)(first.sampler)
-    energy_seed = _place_routes(mesh, np.zeros(q_seed.shape[:2], dtype=np.complex64))
+    system_batch_sharding = NamedSharding(mesh, P("systems", None))
+    systems_sharding = NamedSharding(mesh, P("systems"))
+    q_seed = jax.device_put(q_seed, system_batch_sharding)
+    energy_seed = jax.device_put(
+        np.zeros(q_seed.shape[:2], dtype=np.complex64),
+        system_batch_sharding,
+    )
     kfac = init_router_kfac_state(
         config.kfac,
         model,
@@ -365,6 +758,81 @@ def run_train(
         route_tau=config.router.temperature,
         route_loss_weight=config.router.loss_weight,
     )
+    reframe = _build_reframe(mesh, first.sampler, first.context, first.perms)
+    rebase = _build_rebase(mesh, q_seed, first.perms)
+    adapt_routes = jax.jit(
+        lambda state: _adapt_routes(state, config),
+        in_shardings=(_systems_sharding(mesh, first.sampler),),
+        out_shardings=_systems_sharding(mesh, first.sampler),
+    )
+    target_entry = jax.jit(
+        lambda sampled, baseline, sigma, weights: process_route_targets(
+            sampled,
+            baseline,
+            sigma,
+            weights,
+            mad_width=config.kfac.mad_clip_width,
+        ),
+        in_shardings=(
+            system_batch_sharding,
+            system_batch_sharding,
+            systems_sharding,
+            system_batch_sharding,
+        ),
+        out_shardings=(system_batch_sharding, systems_sharding),
+    )
+    snis_entry = jax.jit(
+        snis_mode_baseline,
+        in_shardings=(
+            system_batch_sharding,
+            system_batch_sharding,
+            system_batch_sharding,
+        ),
+        out_shardings=system_batch_sharding,
+    )
+    router_kernel_template = bind_router_kernel(model)
+    compile_router = _build_owner_entry(
+        mesh,
+        compile_router_static,
+        (
+            router_kernel_template,
+            trunk_template,
+            contexts[0].route_quotient_node_key,
+            contexts[0].route_quotient_edge_key,
+            contexts[0].needs_fwl2,
+        ),
+    )
+    router_static_template = compile_router(
+        router_kernel_template,
+        trunk_template,
+        contexts[0].route_quotient_node_key,
+        contexts[0].route_quotient_edge_key,
+        contexts[0].needs_fwl2,
+    )
+    route_sampler = build_route_sampler(
+        mesh,
+        router_kernel_template.decoder,
+        router_static_template,
+    )
+    mode_sampler = build_beam16(
+        mesh,
+        router_kernel_template.decoder,
+        router_static_template,
+    )
+    mode_perm_template = jnp.arange(config.n_max, dtype=jnp.int32)
+    compile_mode = _build_owner_entry(
+        mesh,
+        _compile_mode,
+        (
+            physical_kernel,
+            trunk_template,
+            energy_inputs[0],
+            contexts[0].mask,
+            contexts[0].bmask,
+            mode_perm_template,
+        ),
+    )
+    mode_energy = None
     system_states[0] = _host_system_state(first)
     del first, q_seed, energy_seed
     order_rng = np.random.default_rng(config.seed)
@@ -378,9 +846,12 @@ def run_train(
             order_rng.shuffle(order)
         system_index = int(order[step % len(order)])
         state = get_system(system_index)
-        trunk = compile_plan(model, contexts[system_index])[0]
+        trunk = compile_plan(
+            bind_trunk_compiler_kernel(model),
+            contexts[system_index],
+        )
         router_kernel = bind_router_kernel(model)
-        router_static = compile_router_static(
+        router_static = compile_router(
             router_kernel,
             trunk,
             contexts[system_index].route_quotient_node_key,
@@ -388,13 +859,11 @@ def run_train(
             contexts[system_index].needs_fwl2,
         )
         tau = jnp.asarray(config.router.temperature, dtype=jnp.float32)
-        router_kernel = _replicate(mesh, router_kernel)
-        router_static = _replicate(mesh, router_static)
         key, key_route = jax.random.split(key)
-        new_perms = build_route_sampler(mesh, router_kernel.decoder, router_static)(
+        new_perms = route_sampler(
             router_kernel.decoder, router_static, key_route, tau
         )
-        mode_perm = build_beam16(mesh, router_kernel.decoder, router_static)(
+        mode_perm = mode_sampler(
             router_kernel.decoder, router_static, tau
         )
         state.sampler, state.context = reframe(
@@ -402,75 +871,81 @@ def run_train(
         )
         state.perms = new_perms
         kernel = bind_shared_kernel(model)
-        trees = compile_trees(model, trunk, new_perms)
+        physical_kernel = bind_physical_compiler_kernel(model)
+        trees = compile_trees(
+            physical_kernel,
+            trunk,
+            new_perms,
+        )
         frames = compile_frames(
             energy_inputs[system_index],
             contexts[system_index].mask,
             contexts[system_index].bmask,
             new_perms,
         )
-        state.sampler = run_routes(
+        state.sampler = get_run_routes(
+            state.sampler,
             kernel,
             trees,
-            state.context,
-            state.sampler,
             config.mcmc.steps,
-            config.mcmc.walker_chunk_size,
-        )
+        )(state.sampler, kernel, trees)
         if step and step % config.mcmc.adapt_every == 0:
-            state.sampler = _adapt_routes(state.sampler, config)
+            state.sampler = adapt_routes(state.sampler)
         q_cold = jax.vmap(cold_samples)(state.sampler)
+        q_cold = jax.device_put(q_cold, system_batch_sharding)
         total, _exchange, _casimir, _field = sampled_energy(
             kernel,
             trees,
             frames,
             q_cold,
-            config.energy.chunk_size,
-        )
-        baseline_is_sampled = bool(
-            np.asarray(
-                jax.device_get(
-                    jnp.all(
-                        new_perms.astype(jnp.int32)
-                        == mode_perm[None, :].astype(jnp.int32)
-                    )
-                )
-            )
         )
+        baseline_is_sampled = bool(np.asarray(jax.device_get(exact_skip(new_perms, mode_perm))))
         if baseline_is_sampled:
             baseline_total = total
-            baseline_weights = jnp.full(
-                total.shape,
-                1.0 / total.shape[-1],
-                dtype=total.real.dtype,
-            )
+            baseline_weights = jnp.ones_like(total.real) / total.shape[-1]
         else:
-            mode_tree = compile_physical_tree_reference(model, trunk, mode_perm)
-            mode_frame = compile_energy_frame(
+            mode_tree, mode_frame = compile_mode(
+                physical_kernel,
+                trunk,
                 energy_inputs[system_index],
                 contexts[system_index].mask,
                 contexts[system_index].bmask,
                 mode_perm,
             )
-            q_canonical = rebase_cold_samples(q_cold, new_perms)
+            q_canonical = rebase(q_cold, new_perms)
+            if mode_energy is None:
+                mode_energy = _build_mode_energy(
+                    mesh,
+                    kernel,
+                    mode_tree,
+                    mode_frame,
+                    q_canonical,
+                    mode_perm,
+                    config.energy.chunk_size,
+                )
             baseline_total, _bx, _bc, _bf, candidate_log_p = mode_energy(
                 kernel,
                 mode_tree,
                 mode_frame,
                 q_canonical,
                 mode_perm,
-                config.energy.chunk_size,
             )
             sampled_log_p = state.sampler.log_p[..., -1]
-            baseline_weights = snis_mode_baseline(
+            sampled_log_p = jax.device_put(sampled_log_p, system_batch_sharding)
+            baseline_weights = snis_entry(
                 baseline_total, candidate_log_p, sampled_log_p
             )
-        target, advantage = process_route_targets(
+        target, advantage = target_entry(
             total,
             baseline_total,
             state.context.s_norm,
             baseline_weights,
-            mad_width=config.kfac.mad_clip_width,
+        )
+        target = jax.device_put(target, system_batch_sharding)
+        advantage = jax.device_put(advantage, systems_sharding)
+        state.context = jax.device_put(
+            state.context,
+            _systems_sharding(mesh, state.context),
         )
         key, key_kfac = jax.random.split(key)
         model, kfac = apply_router_kfac_step(

src/kfac_jax/_src/utils/staging.py

@@ -338,6 +338,7 @@ def staged(
             jax.tree_util.tree_map(_spec_for, a) for a in dynamic_args
         )
         cache_key = (
+            id(instance),
             instance.pmap_axis_name,
             tuple(mesh.axis_names),
             tuple(mesh.shape.items()),