latent_moe_surgery

Weight-only surgery for latent-projected mixture-of-experts layers.

The layer descriptor supplies tensor names; this module owns the mathematical contract. A full-rank transform is function preserving and orders the latent coordinates so a smaller child can be realized by taking a prefix.

Classes

`LatentMoETensorLayout`	Logical latent-MoE tensors in a checkpoint.
`LatentMoETransform`	Full-rank input basis and output factorization for a latent MoE.

Functions

`apply_latent_moe_transform`	Apply a precomputed full-rank transform to tensors resident in one shard.
`apply_latent_moe_sort`	Apply a function-preserving full-rank latent transform in place.
`compute_latent_moe_transform`	Compute the proper Deci sensitivity-weighted full-rank transform.
`reverse_latent_moe_transform`	Reverse both ordered latent bases without changing the represented function.

class LatentMoETensorLayout

Bases: object

Logical latent-MoE tensors in a checkpoint.

Split HF checkpoints use one key per expert. A native grouped checkpoint adapter can expose views using the same logical sequence, keeping the decomposition independent of the storage format.

__init__(fc1_key, fc2_key, expert_up_keys, expert_down_keys)

Parameters:

fc1_key (str)
fc2_key (str)
expert_up_keys (tuple[str, ...])
expert_down_keys (tuple[str, ...])

Return type:

None

expert_down_keys: tuple[str, ...]

expert_up_keys: tuple[str, ...]

fc1_key: str

fc2_key: str

validate(state_dict)

Parameters:: state_dict (Mapping[str, Tensor])
Return type:: int

class LatentMoETransform

Bases: object

Full-rank input basis and output factorization for a latent MoE.

__init__(input_basis, output_basis, output_compressor, transformed_fc2)

Parameters:

input_basis (Tensor)
output_basis (Tensor)
output_compressor (Tensor)
transformed_fc2 (Tensor)

Return type:

None

input_basis: Tensor

output_basis: Tensor

output_compressor: Tensor

transformed_fc2: Tensor

apply_latent_moe_sort(state_dict, layout, *, latent_cov_in, expert_weights, latent_cov_out, reverse=False, compute_device='cpu', tensor_loader=None)

Apply a function-preserving full-rank latent transform in place.

state_dict may be one checkpoint shard. tensor_loader supplies the original logical tensors needed to compute a single global transform; only tensors resident in state_dict are rewritten.

Parameters:

state_dict (MutableMapping[str, Tensor])
layout (LatentMoETensorLayout)
latent_cov_in (Tensor)
expert_weights (Tensor)
latent_cov_out (Tensor)
reverse (bool)
compute_device (device | str)
tensor_loader (Callable[[str], Tensor | None] | None)

Return type:

LatentMoETransform

apply_latent_moe_transform(state_dict, layout, transform)

Apply a precomputed full-rank transform to tensors resident in one shard.

Parameters:

state_dict (MutableMapping[str, Tensor])
layout (LatentMoETensorLayout)
transform (LatentMoETransform)

Return type:

None

compute_latent_moe_transform(fc1, fc2, expert_ups, expert_weights, latent_cov_in, latent_cov_out, *, compute_device='cpu')

Compute the proper Deci sensitivity-weighted full-rank transform.

Parameters:

fc1 (Tensor)
fc2 (Tensor)
expert_ups (Sequence[Tensor])
expert_weights (Tensor)
latent_cov_in (Tensor)
latent_cov_out (Tensor)
compute_device (device | str)

Return type:

LatentMoETransform

reverse_latent_moe_transform(transform)

Reverse both ordered latent bases without changing the represented function.

Parameters:: transform (LatentMoETransform)
Return type:: LatentMoETransform