What is model surgery in a machine-learning checkpoint?

It is editing the parameters of a saved model — the pytree of weight arrays — so that weights stored under one structure land under the structure a different model expects. No training, no forward pass: just a rewrite of a dict of arrays. Orbax’s experimental model_surgery module gives you six composable transforms for this — rename keys by regex, re-nest a flat dict into a pytree, fuse several params into one, repeat a param along a dimension, and stack an indexed family of params onto a new axis. Real reasons you reach for it: importing a HuggingFace checkpoint into a JAX model, fusing gate_proj and up_proj into one gate_up_proj for a faster kernel, expanding grouped-query KV heads to full multi-head, collapsing 64 separately-saved experts into one stacked MoE tensor, or fixing up quantization key names.

How is a model-surgery transform designed in Orbax?

Every operation is a factory that returns a closure. types.py defines one Protocol — Transformation.__call__(*source_array_trees) -> PyTreeOf[jax.Array] — and functions like stack(pattern, expected_count=64) do not do the work; they bake the configuration in and hand back a transform(params) you apply later. That currying is the whole design: configs are fixed at construction time and the returned transforms compose into a pipeline. Two conventions run through all six — they operate on flat dotted-key dicts (“model.layers.0.gate_proj.weight” is a literal string they do surgery on), and each guards against more than one input tree and copies rather than mutates unless inplace=True is set to bound peak memory.

How does Orbax fuse gate_proj and up_proj into one tensor?

fuse_by_pattern takes a regex that matches the candidate keys, an ordered list of the unique parts to combine, and the replacement part for the fused key. For gate/up fusion the pattern is r"^(.*)\.(gate_proj|up_proj)\.weight$", unique_parts is ["gate_proj", "up_proj"], and fused_unique_part is "gate_up_proj". It groups keys by their would-be fused name, and only fuses a group when ALL of the unique parts are present — otherwise it logs a warning and leaves them alone. The actual combine is jnp.concatenate along the given axis; the source keys are deleted and the concatenated value is stored under the fused key.

How does the stack transform collapse MoE experts into one tensor?

stack(pattern, ...) takes a regex with exactly one capture group holding the integer index — e.g. r"mlp\.experts\.(\d+\.)". For each key it reads the index from the capture group, strips that group out of the key to form a base_key, and groups values by base_key. Then per base key it stacks the values in index order onto a new axis, so experts.0.weight … experts.63.weight become one experts.weight with a leading axis of 64. expected_count defaults to max index + 1; sort_by_size=True stacks the largest base keys first to manage peak headroom.

What happens if some expert indices are missing when stacking?

It depends on whether you gave it a filler. If padding is disabled (no filler_mapping and no default_filler) and the indices found are not exactly 0..expected_count-1, stack raises a ValueError naming the base key and the indices it found — it refuses to guess. If a filler value is available, it allocates a full array pre-filled with that value and writes each present index into its slot, logging a warning that it padded the missing indices. So a partial expert set is either an explicit, loud error or an explicit, logged pad — never a silently wrong tensor.

Why is my layer-stacked checkpoint slower to load?

Stacking weights into one tensor per layer is what jax.lax.scan wants, but it works against the loader. A public Orbax issue reports a layer-stacked checkpoint loading roughly four times slower on eight GPUs than the same weights saved per layer, because a single fat tensor cannot be fetched in parallel the way many smaller per-layer tensors can. The stacking that helps the compiler hurts the read. The direction out is to treat stack and unstack as operations of the load itself, so each device shard reads exactly the byte ranges it needs and assembles its piece directly, the same per-shard idea stacking.py already uses in its _streaming_stack path, generalized so on-disk layout no longer dictates load speed.

Model Surgery: Rewriting a Checkpoint’s Parameters, Line by Line

A model is saved under one structure and needed under another — HuggingFace names, a split gate/up projection, 64 separately-saved experts. We walk Orbax’s real model-surgery module: six transforms that rename, fuse, repeat, and stack a pytree of weights, and the host-memory care that lets them run on tens of gigabytes.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Model Surgery: Rewriting a Checkpoint’s Parameters, Line by Line — built to be understood, not skimmed.

Questions

What is model surgery in a machine-learning checkpoint?: It is editing the parameters of a saved model — the pytree of weight arrays — so that weights stored under one structure land under the structure a different model expects. No training, no forward pass: just a rewrite of a dict of arrays. Orbax’s experimental model_surgery module gives you six composable transforms for this — rename keys by regex, re-nest a flat dict into a pytree, fuse several params into one, repeat a param along a dimension, and stack an indexed family of params onto a new axis. Real reasons you reach for it: importing a HuggingFace checkpoint into a JAX model, fusing gate_proj and up_proj into one gate_up_proj for a faster kernel, expanding grouped-query KV heads to full multi-head, collapsing 64 separately-saved experts into one stacked MoE tensor, or fixing up quantization key names.
How is a model-surgery transform designed in Orbax?: Every operation is a factory that returns a closure. types.py defines one Protocol — Transformation.__call__(*source_array_trees) -> PyTreeOf[jax.Array] — and functions like stack(pattern, expected_count=64) do not do the work; they bake the configuration in and hand back a transform(params) you apply later. That currying is the whole design: configs are fixed at construction time and the returned transforms compose into a pipeline. Two conventions run through all six — they operate on flat dotted-key dicts (“model.layers.0.gate_proj.weight” is a literal string they do surgery on), and each guards against more than one input tree and copies rather than mutates unless inplace=True is set to bound peak memory.
How does Orbax fuse gate_proj and up_proj into one tensor?: fuse_by_pattern takes a regex that matches the candidate keys, an ordered list of the unique parts to combine, and the replacement part for the fused key. For gate/up fusion the pattern is r"^(.*)\.(gate_proj|up_proj)\.weight
quot;, unique_parts is ["gate_proj", "up_proj"], and fused_unique_part is "gate_up_proj". It groups keys by their would-be fused name, and only fuses a group when ALL of the unique parts are present — otherwise it logs a warning and leaves them alone. The actual combine is jnp.concatenate along the given axis; the source keys are deleted and the concatenated value is stored under the fused key.
How does the stack transform collapse MoE experts into one tensor?: stack(pattern, ...) takes a regex with exactly one capture group holding the integer index — e.g. r"mlp\.experts\.(\d+\.)". For each key it reads the index from the capture group, strips that group out of the key to form a base_key, and groups values by base_key. Then per base key it stacks the values in index order onto a new axis, so experts.0.weight … experts.63.weight become one experts.weight with a leading axis of 64. expected_count defaults to max index + 1; sort_by_size=True stacks the largest base keys first to manage peak headroom.
What happens if some expert indices are missing when stacking?: It depends on whether you gave it a filler. If padding is disabled (no filler_mapping and no default_filler) and the indices found are not exactly 0..expected_count-1, stack raises a ValueError naming the base key and the indices it found — it refuses to guess. If a filler value is available, it allocates a full array pre-filled with that value and writes each present index into its slot, logging a warning that it padded the missing indices. So a partial expert set is either an explicit, loud error or an explicit, logged pad — never a silently wrong tensor.
Why is my layer-stacked checkpoint slower to load?: Stacking weights into one tensor per layer is what jax.lax.scan wants, but it works against the loader. A public Orbax issue reports a layer-stacked checkpoint loading roughly four times slower on eight GPUs than the same weights saved per layer, because a single fat tensor cannot be fetched in parallel the way many smaller per-layer tensors can. The stacking that helps the compiler hurts the read. The direction out is to treat stack and unstack as operations of the load itself, so each device shard reads exactly the byte ranges it needs and assembles its piece directly, the same per-shard idea stacking.py already uses in its _streaming_stack path, generalized so on-disk layout no longer dictates load speed.

Model Surgery: Rewriting a Checkpoint’s Parameters, Line by Line

Questions

Related explainers