vocab: int,

layers: int,

d: int,

latent_mode: str,

latent_dim: int,

n_groups: int,

codebook_size: int,

) -> dict[str, int]:

components: dict[str, int] = {}

if latent_mode == "discrete":

code_bits = int(math.log2(codebook_size))

group_dim = latent_dim // n_groups

components["token_code_bytes"] = packed_code_bytes(vocab, n_groups, code_bits)

components["codebook_bytes"] = n_groups * codebook_size * group_dim * 2

elif latent_mode == "continuous":

components["latent_table_bytes"] = vocab * latent_dim * 2

else:

raise ValueError(f"unknown latent mode: {latent_mode}")

components["decoder_weight_bytes"] = layers * d * latent_dim * 2

components["decoder_bias_bytes"] = layers * d * 2

components["storage_bytes"] = sum(components.values())

components["raw_bytes"] = vocab * layers * d * 2

codes = np.ascontiguousarray(codes, dtype=np.int64)

if codes.min() < 0 or codes.max() >= (1 << code_bits):

raise SystemExit("codes out of range for the declared code bit width")

shifts = np.arange(code_bits - 1, -1, -1, dtype=np.int64)

bits = ((codes.reshape(-1, 1) >> shifts) & 1).astype(np.uint8).reshape(-1)

packed = np.packbits(bits)

expected = packed_code_bytes(int(codes.shape[0]), int(codes.shape[1]), code_bits)

if int(packed.nbytes) != expected:

raise SystemExit("packed code byte count does not match closed-form accounting")

total_bits = vocab * n_groups * code_bits

bits = np.unpackbits(packed, count=total_bits).reshape(vocab, n_groups, code_bits)

weights = (1 << np.arange(code_bits - 1, -1, -1, dtype=np.int64))

"""G2[l] = mean over pool positions of g_l g_l^T (fp32, [layers, d, d])."""

device = gate_by_layer_device[0].device

second_moment = torch.zeros((layers, d, d), dtype=torch.float32, device=device)

n_positions = int(pool_sample_ids.numel())

for layer in range(layers):

for start in range(0, n_positions, chunk):

stop = min(n_positions, start + chunk)

gate = gate_by_layer_device[layer].index_select(0, pool_sample_ids[start:stop]).to(torch.float32)

second_moment[layer] += gate.T @ gate

del gate

second_moment /= float(n_positions)

"""G[l] = (W_d_l^T W_d_l) * E[g g^T] elementwise; exact position-averaged pullback."""

metric = torch.empty_like(gate_second_moment)

for layer in range(layers):

proj_gram = projection_weights[layer].T @ projection_weights[layer]

metric[layer] = proj_gram * gate_second_moment[layer]

metric[layer] = 0.5 * (metric[layer] + metric[layer].T)

layers = int(metric.shape[0])

sqrt_metric = torch.empty_like(metric)

invsqrt_metric = torch.empty_like(metric)

for layer in range(layers):

evals, evecs = torch.linalg.eigh(metric[layer])

floor = float(evals.max().clamp_min(1e-30)) * eigenvalue_floor

evals = evals.clamp_min(floor)

sqrt_metric[layer] = (evecs * evals.sqrt()) @ evecs.T

invsqrt_metric[layer] = (evecs * evals.rsqrt()) @ evecs.T

vocab = int(packed_device.shape[0])

token_vector_dim = int(packed_device.shape[1])

device = packed_device.device

mean_accum = torch.zeros(token_vector_dim, dtype=torch.float64, device=device)

gram = torch.zeros((token_vector_dim, token_vector_dim), dtype=torch.float32, device=device)

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

x = packed_device[start:stop].to(torch.float32)

mean_accum += x.sum(dim=0).double()

gram += x.T @ x

del x

mean = (mean_accum / vocab).float()

covariance = gram / vocab - torch.outer(mean, mean)

covariance = 0.5 * (covariance + covariance.T)

del gram

"""Top-latent_dim eigenpairs of blockdiag(G^1/2) C blockdiag(G^1/2)."""

whitened = torch.empty_like(covariance)

for layer in range(layers):

rows = slice(layer * d, (layer + 1) * d)

whitened[rows, :] = sqrt_metric[layer] @ covariance[rows, :]

for layer in range(layers):

cols = slice(layer * d, (layer + 1) * d)

whitened[:, cols] = whitened[:, cols] @ sqrt_metric[layer]

whitened = 0.5 * (whitened + whitened.T)

evals, evecs = torch.linalg.eigh(whitened)

del whitened

evals = evals.flip(0).clamp_min(0.0)

evecs = evecs.flip(1)

vocab = int(packed_device.shape[0])

latent_dim = int(basis.shape[1])

device = packed_device.device

latents = torch.empty((vocab, latent_dim), dtype=torch.float32, device=device)

mean_rows = mean.view(layers, d)

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

x = packed_device[start:stop].to(torch.float32).view(stop - start, layers, d)

whitened = torch.einsum("bld,lde->ble", x - mean_rows, sqrt_metric)

latents[start:stop] = whitened.reshape(stop - start, layers * d) @ basis

del x, whitened

"""OPQ-style allocation: balance per-group sums of log-eigenvalue.

group_log_sums = [0.0] * n_groups

group_slots: list[list[int]] = [[] for _ in range(n_groups)]

for dim in range(len(eigenvalues)):

open_groups = [g for g in range(n_groups) if len(group_slots[g]) < group_dim]

target = min(open_groups, key=lambda g: (group_log_sums[g], g))

group_slots[target].append(dim)

group_log_sums[target] += math.log(max(float(eigenvalues[dim]), 1e-12))

vocab = int(latents.shape[0])

device = latents.device

codebooks = torch.empty((n_groups, codebook_size, group_dim), dtype=torch.float32, device=device)

codes = torch.empty((vocab, n_groups), dtype=torch.int64, device=device)

for group in range(n_groups):

data = latents[:, group * group_dim : (group + 1) * group_dim].contiguous()

init_rows = torch.randperm(vocab, generator=generator)[:codebook_size].to(device)

centroids = data.index_select(0, init_rows).clone()

assign = torch.empty(vocab, dtype=torch.int64, device=device)

for _ in range(iterations):

min_dist = torch.empty(vocab, dtype=torch.float32, device=device)

cent_sq = centroids.square().sum(dim=1)

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

dist = cent_sq[None, :] - 2.0 * (data[start:stop] @ centroids.T)

best = dist.argmin(dim=1)

assign[start:stop] = best

min_dist[start:stop] = dist.gather(1, best[:, None]).squeeze(1) + data[start:stop].square().sum(dim=1)

del dist, best

counts = torch.bincount(assign, minlength=codebook_size).to(torch.float32)

sums = torch.zeros_like(centroids)

sums.index_add_(0, assign, data)

nonempty = counts > 0

centroids[nonempty] = sums[nonempty] / counts[nonempty, None]

empty = (~nonempty).nonzero(as_tuple=False).squeeze(1)

if int(empty.numel()) > 0:

far_rows = torch.topk(min_dist, int(empty.numel())).indices

centroids[empty] = data.index_select(0, far_rows)

del min_dist, counts, sums

cent_sq = centroids.square().sum(dim=1)

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

dist = cent_sq[None, :] - 2.0 * (data[start:stop] @ centroids.T)

assign[start:stop] = dist.argmin(dim=1)

del dist

codebooks[group] = centroids

codes[:, group] = assign

del data, centroids, assign

torch: Any,

packed_device,

codes,

codebooks,

decoder_weight,

decoder_bias,

output_metric,

layers: int,

d: int,

n_groups: int,

group_dim: int,

icd_sweeps: int,

ridge: float,

chunk: int,

):

"""Re-assign every token's codes to minimize (z - z*)^T A_z (z - z*).

vocab = int(packed_device.shape[0])

latent_dim = n_groups * group_dim

device = packed_device.device

with torch.no_grad():

weight = decoder_weight.detach()

bias_rows = decoder_bias.detach()

codebooks_det = codebooks.detach()

latent_metric = torch.zeros((latent_dim, latent_dim), dtype=torch.float32, device=device)

metric_weight = torch.empty((layers, d, latent_dim), dtype=torch.float32, device=device)

for layer in range(layers):

metric_weight[layer] = output_metric[layer] @ weight[layer]

latent_metric += weight[layer].T @ metric_weight[layer]

latent_metric = 0.5 * (latent_metric + latent_metric.T)

ridge_value = ridge * float(latent_metric.diagonal().mean().clamp_min(1e-30))

latent_metric += ridge_value * torch.eye(latent_dim, device=device)

group_arange = torch.arange(n_groups, device=device)

z_star = torch.empty((vocab, latent_dim), dtype=torch.float32, device=device)

quantization_error = torch.empty(vocab, dtype=torch.float32, device=device)

codes_changed = 0

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

x = packed_device[start:stop].to(torch.float32).view(stop - start, layers, d)

linear = torch.einsum("bld,lde->be", x - bias_rows, metric_weight)

z_star[start:stop] = torch.linalg.solve(latent_metric, linear.T).T

codes_chunk = codes[start:stop]

z_cur = codebooks_det[group_arange[None, :], codes_chunk].reshape(stop - start, latent_dim)

h = z_cur @ latent_metric - linear

for _ in range(icd_sweeps):

for group in range(n_groups):

cols = slice(group * group_dim, (group + 1) * group_dim)

block = latent_metric[cols, cols]

entries = codebooks_det[group]

quad = ((entries @ block) * entries).sum(dim=1)

lin = h[:, cols] - z_cur[:, cols] @ block

cost = quad[None, :] + 2.0 * (lin @ entries.T)

new_codes = cost.argmin(dim=1)

delta = entries.index_select(0, new_codes) - z_cur[:, cols]

h += delta @ latent_metric[cols.start : cols.stop, :]

z_cur[:, cols] = entries.index_select(0, new_codes)

codes_changed += int((new_codes != codes_chunk[:, group]).sum().item())

codes_chunk[:, group] = new_codes

del quad, lin, cost, new_codes, delta

diff = z_cur - z_star[start:stop]

quantization_error[start:stop] = ((diff @ latent_metric) * diff).sum(dim=1)

del x, linear, codes_chunk, z_cur, h, diff

stats = {

"codes_changed": codes_changed,

"codes_changed_fraction": codes_changed / float(vocab * n_groups * icd_sweeps),

"mean_output_metric_quantization_error": float(quantization_error.mean().item()),

"max_output_metric_quantization_error": float(quantization_error.max().item()),

}

torch: Any,

optimizer,

codes,

codebooks,

z_star,

quantization_error,

n_groups: int,

group_dim: int,

codebook_size: int,

min_usage: int,

):

"""Move unused codebook entries onto the output-metric-optimal latents of the

total_reinitialized = 0

with torch.no_grad():

error_order = torch.argsort(quantization_error, descending=True)

state = optimizer.state.get(codebooks, {})

cursor = 0

for group in range(n_groups):

usage = torch.bincount(codes[:, group], minlength=codebook_size)

dead = (usage < min_usage).nonzero(as_tuple=False).squeeze(1)

n_dead = int(dead.numel())

if n_dead == 0:

continue

donor_tokens = error_order[cursor : cursor + n_dead]

cursor += n_dead

cols = slice(group * group_dim, (group + 1) * group_dim)

codebooks.data[group, dead] = z_star[donor_tokens][:, cols]

codes[donor_tokens, group] = dead

for key in ("exp_avg", "exp_avg_sq"):

if key in state:

state[key][group, dead] = 0.0

total_reinitialized += n_dead

if "metadata" not in npz.files:

raise SystemExit("compressed PLE representation file missing scalar JSON metadata")

raw = npz["metadata"]

if raw.shape != ():

raise SystemExit("compressed PLE representation metadata must be a scalar JSON string")

metadata = json.loads(str(raw.item()))

if not isinstance(metadata, dict):

raise SystemExit("compressed PLE representation metadata JSON must be an object")

if name not in npz.files:

raise SystemExit(f"compressed PLE representation file missing required array {name!r}")

value = npz[name]

if value.dtype != np.dtype(dtype):

raise SystemExit(f"resident array {name!r} dtype {value.dtype} != {np.dtype(dtype)}")

if tuple(value.shape) != shape:

raise SystemExit(f"resident array {name!r} shape {tuple(value.shape)} != {shape}")

np: Any,

compressed_path: Path,

vocab: int,

layers: int,

d: int,

latent_mode: str,

latent_dim: int,

n_groups: int,

codebook_size: int,

expected_storage_bytes: int,

) -> dict[str, Any]:

expected_files = {"metadata", "decoder_weight_fp16", "decoder_bias_fp16"}

if latent_mode == "discrete":

expected_files.update({"token_codes_packed_uint8", "codebooks_fp16"})

else:

expected_files.add("latent_table_fp16")

with np.load(str(compressed_path), allow_pickle=False) as npz:

if set(npz.files) != expected_files:

raise SystemExit(f"resident file entries {sorted(npz.files)} != {sorted(expected_files)}")

metadata = load_npz_metadata_object(np, npz)

if metadata.get("compressed_ple_format_id") != FORMAT_ID:

raise SystemExit("resident metadata compressed_ple_format_id does not match this builder")

if metadata.get("method") != METHOD_ID or metadata.get("kind") != METHOD_ID:

raise SystemExit("resident metadata method/kind does not match this builder")

if metadata.get("ple_shape") != [vocab, layers * d]:

raise SystemExit("resident metadata ple_shape does not match expected full PLE shape")

if metadata.get("latent_mode") != latent_mode or int(metadata.get("latent_dim", -1)) != latent_dim:

raise SystemExit("resident metadata latent mode/dim does not match launch arguments")

payload: dict[str, Any] = {"metadata": metadata}

payload["decoder_weight_fp16"] = require_npz_array(

np, npz, "decoder_weight_fp16", np.float16, (layers, d, latent_dim)

)

payload["decoder_bias_fp16"] = require_npz_array(np, npz, "decoder_bias_fp16", np.float16, (layers, d))

if latent_mode == "discrete":

code_bits = int(math.log2(codebook_size))

group_dim = latent_dim // n_groups

payload["token_codes_packed_uint8"] = require_npz_array(

np, npz, "token_codes_packed_uint8", np.uint8, (packed_code_bytes(vocab, n_groups, code_bits),)

)

payload["codebooks_fp16"] = require_npz_array(

np, npz, "codebooks_fp16", np.float16, (n_groups, codebook_size, group_dim)

)

else:

payload["latent_table_fp16"] = require_npz_array(

np, npz, "latent_table_fp16", np.float16, (vocab, latent_dim)

)

resident_bytes = sum(int(value.nbytes) for key, value in payload.items() if key != "metadata")

if resident_bytes != expected_storage_bytes:

raise SystemExit(

f"resident tensor bytes {resident_bytes} != declared storage bytes {expected_storage_bytes}"

)

if latent_mode == "discrete":

code_bits = int(math.log2(codebook_size))

codes_np = unpack_codes(np, payload["token_codes_packed_uint8"], vocab, n_groups, code_bits)

codes = torch.from_numpy(codes_np).to(device)

codebooks = torch.from_numpy(payload["codebooks_fp16"]).to(device=device, dtype=torch.float32)

group_arange = torch.arange(n_groups, device=device)

return lambda start, stop: codebooks[group_arange[None, :], codes[start:stop]].reshape(

stop - start, latent_dim

)

latent_table = torch.from_numpy(payload["latent_table_fp16"]).to(device=device, dtype=torch.float32)

torch: Any,

np: Any,

payload: dict[str, Any],

recon_path: Path,

vocab: int,

layers: int,

d: int,

latent_mode: str,

latent_dim: int,

n_groups: int,

codebook_size: int,

device,

chunk: int,

) -> None:

from safetensors.torch import save_file

latent_chunk_fn = latents_from_payload(

torch, np, payload, vocab, latent_mode, latent_dim, n_groups, codebook_size, device

)

decoder_weight = torch.from_numpy(payload["decoder_weight_fp16"]).to(device=device, dtype=torch.float32)

decoder_bias = torch.from_numpy(payload["decoder_bias_fp16"]).to(device=device, dtype=torch.float32)

recon = torch.empty((vocab, layers * d), dtype=torch.bfloat16)

for layer in range(layers):

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

decoded = decode_payload_chunk(torch, latent_chunk_fn, decoder_weight, decoder_bias, layer, start, stop)

recon[start:stop, layer * d : (layer + 1) * d] = decoded.to(torch.bfloat16).cpu()

del decoded

recon_path.parent.mkdir(parents=True, exist_ok=True)

torch: Any,

np: Any,

safe_open: Any,

compressed_path: Path,

recon_path: Path,

vocab: int,

layers: int,

d: int,

latent_mode: str,

latent_dim: int,

n_groups: int,

codebook_size: int,

expected_storage_bytes: int,

device,

chunk: int,

atol: float,

) -> float:

payload = load_resident_payload(

np, compressed_path, vocab, layers, d, latent_mode, latent_dim, n_groups, codebook_size, expected_storage_bytes

)

latent_chunk_fn = latents_from_payload(

torch, np, payload, vocab, latent_mode, latent_dim, n_groups, codebook_size, device

)

decoder_weight = torch.from_numpy(payload["decoder_weight_fp16"]).to(device=device, dtype=torch.float32)

decoder_bias = torch.from_numpy(payload["decoder_bias_fp16"]).to(device=device, dtype=torch.float32)

max_abs_diff = 0.0

with safe_open(str(recon_path), framework="pt", device="cpu") as recon_file:

if RECON_KEY not in set(recon_file.keys()):

raise SystemExit(f"missing dense reconstruction tensor key in {recon_path}: {RECON_KEY}")

recon_slice = recon_file.get_slice(RECON_KEY)

if list(recon_slice.get_shape()) != [vocab, layers * d]:

raise SystemExit("resident roundtrip recon shape does not match expected full PLE shape")

for layer in range(layers):

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

decoded = decode_payload_chunk(

torch, latent_chunk_fn, decoder_weight, decoder_bias, layer, start, stop

).to(torch.bfloat16).to(torch.float32)

saved = recon_slice[start:stop, layer * d : (layer + 1) * d].to(device).to(torch.float32)

if not bool(torch.isfinite(decoded).all().item()):

raise SystemExit("resident file reload/decode produced non-finite values")

if not bool(torch.isfinite(saved).all().item()):

raise SystemExit("saved BF16 dense PLE reconstruction contains non-finite values")

diff = float((decoded - saved).abs().max().item())

max_abs_diff = max(max_abs_diff, diff)

if max_abs_diff > atol:

raise SystemExit(

f"resident file reload/decode max_abs_diff {max_abs_diff} exceeds tolerance {atol}"

)

del decoded, saved

torch: Any,

safe_open: Any,

recon_path: Path,

packed_original_device,

layers: int,

d: int,

chunk: int,

) -> float:

"""Row-space NMSE of the saved recon file against the device-resident

device = packed_original_device.device

vocab = int(packed_original_device.shape[0])

err = torch.zeros((), dtype=torch.float64, device=device)

norm = torch.zeros((), dtype=torch.float64, device=device)

with safe_open(str(recon_path), framework="pt", device="cpu") as recon_file:

recon_slice = recon_file.get_slice(RECON_KEY)

for layer in range(layers):

for start in range(0, vocab, chunk):

stop = min(vocab, start + chunk)

recon = recon_slice[start:stop, layer * d : (layer + 1) * d].to(device).to(torch.float32)

original = packed_original_device[start:stop, layer * d : (layer + 1) * d].to(torch.float32)

err += (recon - original).double().square().sum()

norm += original.double().square().sum()

del recon, original

torch: Any,

safe_open: Any,

recon_path: Path,

packed_original_device,

token_ids_device,

gate_by_layer_device,

projection_weights,

runtime_token_scale: float,

layers: int,

d: int,

position_chunk: int,

) -> float:

"""NMSE of the exact forward path over every locked-cache position and layer."""

device = packed_original_device.device

n_positions = int(token_ids_device.numel())

scale = float(runtime_token_scale)

err = 0.0

sig = 0.0

with safe_open(str(recon_path), framework="pt", device="cpu") as recon_file:

recon_slice = recon_file.get_slice(RECON_KEY)

for layer in range(layers):

recon_layer = recon_slice[:, layer * d : (layer + 1) * d].to(device).to(torch.float32)

original_layer = packed_original_device[:, layer * d : (layer + 1) * d]

for start in range(0, n_positions, position_chunk):

stop = min(n_positions, start + position_chunk)

tokens = token_ids_device[start:stop]

gate = gate_by_layer_device[layer][start:stop].to(torch.float32)

original_rows = original_layer.index_select(0, tokens).to(torch.float32)

diff_rows = recon_layer.index_select(0, tokens) - original_rows

err += float(

(((gate * diff_rows) * scale) @ projection_weights[layer].T).square().sum().item()

)

sig += float(

(((gate * original_rows) * scale) @ projection_weights[layer].T).square().sum().item()

)

del tokens, gate, original_rows, diff_rows

del recon_layer

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument("--it-weights", required=True)

parser.add_argument("--num-layers", type=int, required=True)

parser.add_argument("--d-ple", type=int, default=256)

parser.add_argument("--latent-mode", choices=LATENT_MODES, default="discrete")

parser.add_argument("--latent-dim", type=int, required=True)

parser.add_argument("--n-groups", type=int, default=0)

parser.add_argument("--codebook-size", type=int, default=0)

parser.add_argument("--group-dim-allocation", choices=GROUP_DIM_ALLOCATION_MODES, default="eigenvalue_balanced")

parser.add_argument("--alternations", type=int, required=True)

parser.add_argument("--steps-per-alternation", type=int, required=True)

parser.add_argument("--assignment-icd-sweeps", type=int, default=2)

parser.add_argument("--assignment-ridge", type=float, default=1e-8)

parser.add_argument("--assignment-chunk-size", type=int, default=65536)

parser.add_argument("--dead-code-min-usage", type=int, default=1)

parser.add_argument("--kmeans-iterations", type=int, default=25)

parser.add_argument("--ste-warmup-steps", type=int, default=0)

parser.add_argument("--ste-commitment-beta", type=float, default=0.25)

parser.add_argument("--ste-ema-gamma", type=float, default=0.99)

parser.add_argument("--ste-ema-eps", type=float, default=1e-5)

parser.add_argument("--metric-eigenvalue-floor", type=float, default=1e-6)

parser.add_argument("--lr", type=float, default=3e-4)

parser.add_argument("--optimizer", choices=OPTIMIZER_MODES, default="adamw")

parser.add_argument("--adamw-beta1", type=float, default=0.9)

parser.add_argument("--adamw-beta2", type=float, default=0.999)

parser.add_argument("--adamw-eps", type=float, default=1e-8)

parser.add_argument(

"--weight-decay",

type=float,

default=0.0,

help="AdamW weight decay applied to the decoder weight/bias only; codebook entries and the continuous latent table always train without weight decay",

)

parser.add_argument("--grad-clip-norm", type=float, default=1.0)

parser.add_argument("--lr-schedule", choices=LR_SCHEDULE_MODES, default="constant")

parser.add_argument("--warmup-steps", type=int, default=0)

parser.add_argument("--batch-size", type=int, default=8192)

parser.add_argument("--training-gate-cache", required=True)

parser.add_argument("--training-gate-cache-config", required=True)

parser.add_argument("--heldout-eval-size", type=int, default=0)

parser.add_argument("--heldout-eval-seed", type=int, default=1)

parser.add_argument("--heldout-eval-interval-steps", type=int, default=0)

parser.add_argument("--eval-chunk-size", type=int, default=8192)

parser.add_argument("--init-covariance-chunk-size", type=int, default=4096)

parser.add_argument("--gate-moment-chunk-size", type=int, default=16384)

parser.add_argument("--forward-path-eval-position-chunk-size", type=int, default=16384)

parser.add_argument("--resident-roundtrip-atol", type=float, default=0.0)

parser.add_argument("--device", required=True)

parser.add_argument("--seed", type=int, default=0)

parser.add_argument("--save-recon", required=True)

parser.add_argument("--save-compressed-ple-npz", required=True)

parser.add_argument("--output", required=True)

if args.num_layers <= 0 or args.d_ple != 256:

raise SystemExit("this builder supports positive Gemma PLE layer counts with d=256")

if args.latent_dim <= 0:

raise SystemExit("--latent-dim must be positive")

# Non-positive chunk sizes would leave the materialize and verify loops

# empty and save an uninitialized reconstruction.

for flag_name in (

"assignment_chunk_size",

"eval_chunk_size",

"init_covariance_chunk_size",

"gate_moment_chunk_size",

"forward_path_eval_position_chunk_size",

):

if getattr(args, flag_name) <= 0:

raise SystemExit(f"--{flag_name.replace('_', '-')} must be positive")

if not math.isfinite(args.resident_roundtrip_atol) or args.resident_roundtrip_atol < 0.0:

raise SystemExit("--resident-roundtrip-atol must be finite and nonnegative")

if args.latent_mode == "discrete":

if args.n_groups <= 0:

raise SystemExit("discrete latent mode requires positive --n-groups")

if args.latent_dim % args.n_groups != 0:

raise SystemExit("--latent-dim must be divisible by --n-groups")

if (

args.codebook_size < 2

or args.codebook_size > 65536

or (args.codebook_size & (args.codebook_size - 1)) != 0

):

raise SystemExit("--codebook-size must be a power of two in [2, 65536]")

if args.kmeans_iterations <= 0:

raise SystemExit("--kmeans-iterations must be positive")

if args.assignment_icd_sweeps <= 0:

raise SystemExit("--assignment-icd-sweeps must be positive")

if args.dead_code_min_usage < 0:

raise SystemExit("--dead-code-min-usage must be nonnegative")

else:

if args.n_groups != 0 or args.codebook_size != 0:

raise SystemExit("continuous latent mode must not receive --n-groups or --codebook-size")

if args.ste_warmup_steps != 0:

raise SystemExit("continuous latent mode has no vector-quantization warm start")

if args.alternations <= 0 or args.steps_per_alternation <= 0 or args.batch_size <= 0:

raise SystemExit("--alternations, --steps-per-alternation, and --batch-size must be positive")

if args.ste_warmup_steps < 0:

raise SystemExit("--ste-warmup-steps must be nonnegative")

if not (0.0 < args.ste_ema_gamma < 1.0):

raise SystemExit("--ste-ema-gamma must be in (0, 1)")

if args.ste_commitment_beta < 0 or not math.isfinite(args.ste_commitment_beta):

raise SystemExit("--ste-commitment-beta must be finite and nonnegative")

if args.lr <= 0 or not math.isfinite(args.lr):

raise SystemExit("--lr must be positive and finite")

if not (0.0 <= args.adamw_beta1 < 1.0) or not (0.0 <= args.adamw_beta2 < 1.0):

raise SystemExit("--adamw-beta1 and --adamw-beta2 must be in [0, 1)")

if args.adamw_eps <= 0 or not math.isfinite(args.adamw_eps):

raise SystemExit("--adamw-eps must be positive and finite")

if args.weight_decay < 0 or not math.isfinite(args.weight_decay):

raise SystemExit("--weight-decay must be finite and nonnegative")

if args.grad_clip_norm <= 0 or not math.isfinite(args.grad_clip_norm):

raise SystemExit("--grad-clip-norm must be positive and finite")

total_steps = args.ste_warmup_steps + args.alternations * args.steps_per_alternation

if args.warmup_steps < 0 or args.warmup_steps > total_steps:

raise SystemExit("--warmup-steps must be between 0 and the total Adam step count")

if args.heldout_eval_size < 0 or args.heldout_eval_interval_steps < 0:

raise SystemExit("--heldout-eval-size and --heldout-eval-interval-steps must be nonnegative")

if args.metric_eigenvalue_floor <= 0 or args.metric_eigenvalue_floor >= 1:

raise SystemExit("--metric-eigenvalue-floor must be in (0, 1)")

if args.assignment_ridge < 0 or not math.isfinite(args.assignment_ridge):

args = parse_args()

validate_args(args)

import numpy as np

import torch

from safetensors import safe_open

t0 = time.time()

torch.manual_seed(args.seed)

np.random.seed(args.seed)

if torch.cuda.is_available():

torch.cuda.manual_seed_all(args.seed)

device = torch.device(args.device)

layers = args.num_layers

d = args.d_ple

latent_dim = args.latent_dim

n_groups = args.n_groups

codebook_size = args.codebook_size

group_dim = latent_dim // n_groups if args.latent_mode == "discrete" else 0

code_bits = int(math.log2(codebook_size)) if args.latent_mode == "discrete" else 0

total_adam_steps = args.ste_warmup_steps + args.alternations * args.steps_per_alternation

components = storage_components(

262144, layers, d, args.latent_mode, latent_dim, n_groups, codebook_size

)

# ----- load original PLE and the locked training inputs ----- #

original_packed_cpu, ple_key, ple_shape = load_ple_tensor(

safe_open, args.it_weights, layers, d, allow_recon_key=False

)

vocab = int(original_packed_cpu.shape[0])

packed = original_packed_cpu.to(device=device, dtype=torch.bfloat16)

del original_packed_cpu

(

token_ids_cpu,

gate_by_layer_cpu,

runtime_token_scale,

training_loss_metadata,

) = build_exact_gate_cache_forward_path_training_inputs(

torch, safe_open, args.training_gate_cache, args.training_gate_cache_config, vocab, layers, d

)

training_loss_metadata = {

**training_loss_metadata,

"training_objective_mode": TRAINING_OBJECTIVE_MODE,

"training_forward_path_formula": (

"mean || W_d_l((g_l * (decoded_row - original_row)) * runtime_token_scale) ||^2"

),

"training_gate_cache_materialization": (

"device-resident per-layer float16 gate tensors loaded from the fixed D_out gate-cache file"

),

}

projection_weights, projection_metadata = load_projection_weights(

torch, safe_open, args.it_weights, layers, d, device

)

token_ids_device = token_ids_cpu.to(device)

gate_by_layer = [gate.to(device=device) for gate in gate_by_layer_cpu]

del gate_by_layer_cpu

n_positions = int(token_ids_device.numel())

scale = float(runtime_token_scale)

# ----- heldout split over gate-cache sample positions ----- #

heldout_eval_enabled = args.heldout_eval_size > 0

heldout_sample_ids_cpu = None

heldout_layer_ids_cpu = None

train_sample_pool_cpu = torch.arange(n_positions, dtype=torch.int64)

heldout_split_scope = "disabled"