diff --git a/.gitignore b/.gitignore index 3b475834..e244a554 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,4 @@ internal-tests/ internal-docs/ notes/ third_party/cutlass +/artifacts/omega_e0m3/*.pt diff --git a/docs/omega_pack_e0m3.md b/docs/omega_pack_e0m3.md new file mode 100644 index 00000000..41e3a546 --- /dev/null +++ b/docs/omega_pack_e0m3.md @@ -0,0 +1,160 @@ +# Omega-QVLA pack to FlashRT E0M3 format + +This document describes the **Milestone 1 offline toolchain** shipped in this +repository. It converts rank-0 Omega-QVLA `dit_svdquant_v1` records into the +packed E0M3/UE4M3 weight format consumed by FlashRT's existing SM110 GEMM. + +This change does not connect the artifact to a frontend or pipeline. It adds no +runtime route, server, CUDA graph, CMake source, binding, or public Python API. +Runtime-consumer results mentioned in development discussions were produced on +an external experimental branch and are not capabilities of this repository. + +## Included tools + +- `tools/convert_omega_pack_e0m3.py`: validates and converts a pack. +- `tools/check_omega_e0m3_layer.py`: CPU emulation diagnostics and a Thor + kernel round-trip that directly consumes converted `packed` and `sfb` data. +- `tools/gen_omega_pack_fixture.py`: creates a four-record synthetic fixture. + +## Input container + +The input is a plain `torch.save` mapping loadable with +`torch.load(..., weights_only=True)`. A full inspected pi0.5 pack contains 252 +records plus `__meta__`: 126 expert records and 126 PaliGemma records. + +Each selected record must have `format == "dit_svdquant_v1"` and satisfy: + +| Field | Required contract | +|---|---| +| `weight_res_q` | finite floating tensor `[out_features, in_features]` | +| `rank` | exactly `0` | +| `lowrank_A`, `lowrank_B` | `[out_features, 0]`, `[in_features, 0]` | +| `weight_bits`, `a_bits` | both `4` | +| `act_scale_table` | finite positive floating tensor `[steps, in_features]` | +| `duquant_rotation_blocks` | finite `[in_features / 64, 64, 64]` | +| `duquant_rotation_perm` | int64 permutation of `[0, in_features)` | +| `duquant_rotation_out_blocks` | finite `[out_features / 64, 64, 64]` | +| `in_features`, `out_features` | exact match for the weight tensor | + +Both feature dimensions must be positive multiples of 64. The converter +validates every selected record before importing the CUDA extension or writing +an artifact. Unsupported records are errors; they are never skipped. + +For an unfiltered non-fixture pack, the default coverage gate is 252 records. +`--layer-regex` and fixture packs expect every selected record. Use +`--expected-records` when intentionally converting a different complete pack. + +## Output container + +The output format is `omega_e0m3_v1`, schema version 1: + +```text +{ + "format": "omega_e0m3_v1", + "schema_version": 1, + "source_pack_meta": {...}, + "source_record_count": int, + "selected_record_count": int, + "selected_layers": [str, ...], + "fold": "none" | "mean" | "actnorm", + "weights": { + layer: { + "packed": uint8[N, K / 2], + "sfb": uint8[sfa_size_bytes(N, K, True)], + "N": int, + "K": int, + } + }, + "aux": {layer: {...}} +} +``` + +The converter writes through a sibling temporary file and atomically replaces +the destination only after all selected records convert successfully. The +`weights`, `aux`, and `selected_layers` sets must be identical. + +## Conversion math + +Omega stores `weight_res_q` as a dequantized fp16 tensor already in the +rotated and permuted domain. No GPTQ bitstream decoding is required. + +FlashRT converts each `[N, K]` weight with: + +```text +quantize_e0m3_dynamic_sfa_fp16(weight, packed, sfb, N, K, is_sfb=True) +``` + +This produces packed 4-bit E0M3 elements and tile-interleaved UE4M3 per-16 +scales. SFB storage is allocated with +`flash_rt_fp4.sfa_size_bytes(N, K, True)` and zero-initialized because the +tile-interleaved layout contains padding. + +The default `--fold none` strategy intentionally does not fold the +per-channel activation table into the weight. Thor experiments on four real +layers measured E0M3-vs-fp16 cosine around 0.9924-0.9932, comparable to the +source fake-quant path. `mean` and `actnorm` remain ablation modes; they are not +the recommended artifact format. + +## Reproduction + +### CPU fixture and emulation + +```bash +python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt + +python tools/convert_omega_pack_e0m3.py \ + --pack /tmp/fixture_pack.pt \ + --out /tmp/unused.pt \ + --validate-only + +python tools/check_omega_e0m3_layer.py \ + --pack /tmp/fixture_pack.pt \ + --mode emulate +``` + +The validation command performs no CUDA work and does not create its output. + +### Thor artifact round-trip + +```bash +python tools/convert_omega_pack_e0m3.py \ + --pack /tmp/fixture_pack.pt \ + --out /tmp/fixture_e0m3.pt \ + --fold none + +python tools/check_omega_e0m3_layer.py \ + --pack /tmp/fixture_pack.pt \ + --artifact /tmp/fixture_e0m3.pt \ + --mode kernel \ + --min-artifact-cos 0.98 +``` + +With `--artifact`, the checker quantizes only the activation. It loads the +weight's `packed` and `sfb` tensors directly from the converted artifact, +validates their shape and byte count, runs the GEMM, and fails if cosine is +below the requested threshold. This is the converter round-trip gate. + +### Full pack + +```bash +python tools/convert_omega_pack_e0m3.py \ + --pack /path/to/quantized.pt \ + --out /path/to/pi05_e0m3.pt \ + --fold none +``` + +Without a layer filter, this command fails unless all 252 records validate and +convert. A representative artifact layer should then be checked with the same +`--artifact --mode kernel` command above. Release evidence must include both +the direct packed/SFB GEMM result and `252/252` coverage. + +## Scope and roadmap + +Milestone 1 is limited to conversion, schema validation, a synthetic fixture, +CPU emulation, and direct artifact GEMM verification. + +A future runtime PR may load these artifacts in a pi0.5 Thor frontend, apply +the DuQuant input/output rotations, and establish end-to-end accuracy and +latency gates. Such a PR must independently add and test its frontend, +pipeline, graph-lifecycle, and serving contracts. None of that runtime surface +is provided here. diff --git a/tests/test_omega_e0m3_tools.py b/tests/test_omega_e0m3_tools.py new file mode 100644 index 00000000..87132fc8 --- /dev/null +++ b/tests/test_omega_e0m3_tools.py @@ -0,0 +1,203 @@ +import copy +import importlib.util +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + + +ROOT = Path(__file__).resolve().parents[1] + + +def _load_tool(name: str): + path = ROOT / "tools" / f"{name}.py" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +converter = _load_tool("convert_omega_pack_e0m3") +checker = _load_tool("check_omega_e0m3_layer") + + +def _record(n: int = 64, k: int = 64) -> dict: + return { + "format": "dit_svdquant_v1", + "weight_res_q": torch.randn(n, k, dtype=torch.float16), + "lowrank_A": torch.empty(n, 0, dtype=torch.float16), + "lowrank_B": torch.empty(k, 0, dtype=torch.float16), + "act_scale_table": torch.ones(2, k), + "duquant_rotation_blocks": torch.eye(64).half().view(1, 64, 64), + "duquant_rotation_perm": torch.arange(k, dtype=torch.int64), + "duquant_rotation_out_blocks": torch.eye(64).half().view(1, 64, 64), + "weight_bits": 4, + "a_bits": 4, + "rank": 0, + "in_features": k, + "out_features": n, + } + + +def test_converter_accepts_complete_rank_zero_record(): + assert converter.validate_record("layer", _record()) == (64, 64) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + (lambda r: r.update(format="other"), "unsupported format"), + (lambda r: r.update(rank=1), "only rank=0"), + (lambda r: r.update(in_features=128), "metadata shape mismatch"), + (lambda r: r.update(act_scale_table=torch.ones(2, 32)), + "act_scale_table"), + (lambda r: r.update(duquant_rotation_blocks=torch.eye(64).half()), + "input rotation shape"), + (lambda r: r.update(duquant_rotation_perm=torch.zeros(64, dtype=torch.int64)), + "not a permutation"), + ], +) +def test_converter_rejects_incomplete_or_inconsistent_records(mutation, message): + record = copy.deepcopy(_record()) + mutation(record) + with pytest.raises(ValueError, match=message): + converter.validate_record("layer", record) + + +def test_full_pack_count_defaults_to_252_but_fixture_and_subset_are_exact(): + assert converter.expected_record_count( + {}, layer_regex="", selected_count=251, explicit=None) == 252 + assert converter.expected_record_count( + {"recipe": "synthetic fixture (test)"}, layer_regex="", selected_count=4, + explicit=None) == 4 + assert converter.expected_record_count( + {}, layer_regex="layer.0", selected_count=1, explicit=None) == 1 + + +def test_validate_only_runs_without_cuda_and_writes_no_artifact(tmp_path): + pack_path = tmp_path / "fixture_pack.pt" + out_path = tmp_path / "must_not_exist.pt" + torch.save({ + "__meta__": {"recipe": "synthetic fixture (test)"}, + "layer": _record(), + }, pack_path) + + result = subprocess.run( + [ + sys.executable, + str(ROOT / "tools" / "convert_omega_pack_e0m3.py"), + "--pack", str(pack_path), + "--out", str(out_path), + "--validate-only", + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 0, result.stderr + assert "validated 1/1 records" in result.stdout + assert not out_path.exists() + + +def test_validate_only_fails_before_cuda_for_bad_records(tmp_path): + pack_path = tmp_path / "bad_pack.pt" + out_path = tmp_path / "must_not_exist.pt" + record = _record() + record["rank"] = 1 + torch.save({ + "__meta__": {"recipe": "synthetic fixture (test)"}, + "layer": record, + }, pack_path) + + result = subprocess.run( + [ + sys.executable, + str(ROOT / "tools" / "convert_omega_pack_e0m3.py"), + "--pack", str(pack_path), + "--out", str(out_path), + "--validate-only", + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 2 + assert "only rank=0 is supported" in result.stderr + assert not out_path.exists() + + +def test_converter_refuses_to_overwrite_source_pack(tmp_path): + pack_path = tmp_path / "fixture_pack.pt" + torch.save({ + "__meta__": {"recipe": "synthetic fixture (test)"}, + "layer": _record(), + }, pack_path) + result = subprocess.run( + [ + sys.executable, + str(ROOT / "tools" / "convert_omega_pack_e0m3.py"), + "--pack", str(pack_path), + "--out", str(pack_path), + ], + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 2 + assert "must not overwrite" in result.stderr + loaded = torch.load(pack_path, map_location="cpu", weights_only=True) + assert "layer" in loaded + + +def test_artifact_schema_requires_complete_direct_packed_sfb_coverage(): + class FakeFp4: + @staticmethod + def sfa_size_bytes(n, k, is_sfb): + assert is_sfb + return 32 + + artifact = { + "format": "omega_e0m3_v1", + "schema_version": 1, + "fold": "none", + "selected_record_count": 1, + "selected_layers": ["layer"], + "weights": { + "layer": { + "packed": torch.zeros(64, 32, dtype=torch.uint8), + "sfb": torch.zeros(32, dtype=torch.uint8), + "N": 64, + "K": 64, + } + }, + "aux": {"layer": {"fold": "none"}}, + } + entry, aux = checker.validate_artifact( + artifact, "layer", 64, 64, FakeFp4()) + assert entry["packed"].shape == (64, 32) + assert aux["fold"] == "none" + + broken = copy.deepcopy(artifact) + broken["aux"] = {} + with pytest.raises(ValueError, match="do not cover"): + checker.validate_artifact(broken, "layer", 64, 64, FakeFp4()) + + +def test_docs_and_gitignore_only_describe_milestone_one_surface(): + docs = (ROOT / "docs" / "omega_pack_e0m3.md").read_text() + gitignore = (ROOT / ".gitignore").read_text().splitlines() + + for nonexistent in ( + "omega_e0m3_linear.py", + "check_omega_e0m3_consumer.py", + "serve_omega_e0m3.py", + "omega_e0m3_graph.py", + "check_omega_e0m3_graph_smoke.py", + "start_e0m3_server.sh", + ): + assert nonexistent not in docs + assert "*.pt" not in gitignore + assert "/artifacts/omega_e0m3/*.pt" in gitignore diff --git a/tools/check_omega_e0m3_layer.py b/tools/check_omega_e0m3_layer.py new file mode 100644 index 00000000..1ab9e27c --- /dev/null +++ b/tools/check_omega_e0m3_layer.py @@ -0,0 +1,376 @@ +#!/usr/bin/env python3 +"""Single-layer cosine harness: Omega fake-quant vs. FlashRT E0M3 GEMM. + +Quantifies the fidelity cost of migrating one Omega-QVLA dit_svdquant_v1 +record (docs/omega_pack_e0m3.md) onto the FlashRT SM110 E0M3 path. + +References (torch, fp32 accumulate): + y_fp = x2 @ W^T — no activation quant (ceiling; + W is already fake-quant dequant) + y_omega = fakequant(x2, s_t) @ W^T — exact Omega consumer semantics + (per-channel table, int [-8,7]) + +Variants: + S0: A = e0m3(x2), B = e0m3(W) — drops the scale table + S1: A = e0m3(x2 / s_t), B = e0m3(W * diag(s_mean)) — mean-fold strategy + +x2 is the rotated activation (perm + 64x64 block rotation), computed in +fp16 exactly like the Omega runtime. The output rotation is an exact +orthonormal transform applied to BOTH compared vectors, so it is skipped +(cosines are invariant to it). + +Modes: + --mode emulate : pure-torch E0M3 emulation (per-16 amax/7, UE4M3-rounded + scales, int [-7,7]). Runs anywhere; approximates the + tcgen05 result up to accumulation order and UE4M3 + rounding corner cases. Use for local pre-checks. + --mode kernel : real kernels via flash_rt_fp4 (quantize + GEMM, + a_format=0). Requires CUDA + built extension (Thor). + +Usage: + python tools/check_omega_e0m3_layer.py \ + --pack packs_hf/pi05_long/quantized.pt --mode emulate + python tools/check_omega_e0m3_layer.py \ + --pack packs_hf/pi05_long/quantized.pt --mode kernel \ + --layer paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj + python tools/check_omega_e0m3_layer.py \ + --pack /tmp/fixture_pack.pt --artifact /tmp/fixture_e0m3.pt \ + --mode kernel --min-artifact-cos 0.98 +""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Mapping + +import torch + +DEFAULT_LAYER = ("paligemma_with_expert.gemma_expert.model.layers.0." + "self_attn.q_proj") + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--pack", required=True) + p.add_argument("--layer", default=DEFAULT_LAYER) + p.add_argument("--mode", choices=("emulate", "kernel"), default="emulate") + p.add_argument("--tokens", type=int, default=256, help="M (rows)") + p.add_argument("--step", type=int, default=0, + help="denoise step index into act_scale_table") + p.add_argument("--seed", type=int, default=0) + p.add_argument("--device", default="", + help="torch device for reference/emulation math " + "(default: cuda in kernel mode, cpu in emulate mode)") + p.add_argument("--artifact", + help="converted omega_e0m3_v1 artifact; in kernel mode, " + "use its packed/SFB weight directly instead of re-quantizing") + p.add_argument("--min-artifact-cos", type=float, default=0.98, + help="minimum artifact-vs-fp16 global cosine (default: 0.98)") + return p.parse_args() + + +# ──────────────────────────────────────────────────────────────────── +# Omega consumer semantics (mirror gr00t/quantization/gptq_layers.py) +# ──────────────────────────────────────────────────────────────────── +def apply_input_rotation(x: torch.Tensor, perm: torch.Tensor, + blocks: torch.Tensor) -> torch.Tensor: + """x[N, in] fp16 -> bmm(x[:, perm].view(N, nb, B), blocks). fp16 in/out.""" + nb, b, _ = blocks.shape + dev = x.device + x2 = x.index_select(dim=-1, index=perm.to(dev)) + x2 = x2.reshape(-1, nb, b) + x2 = torch.bmm(x2.transpose(0, 1).contiguous(), + blocks.to(device=dev, dtype=x.dtype)) + return x2.transpose(0, 1).contiguous().reshape(x.shape[0], nb * b) + + +def fake_quant_omega(x: torch.Tensor, scale: torch.Tensor, + bits: int = 4) -> torch.Tensor: + """clamp(round(x/s), -2^(b-1), 2^(b-1)-1) * s — Omega's asymmetric grid.""" + qmax = 2 ** (bits - 1) - 1 + return (torch.clamp(torch.round(x / scale), -qmax - 1, qmax) + * scale).float() + + +# ──────────────────────────────────────────────────────────────────── +# E0M3 emulation (approximates quantize_e0m3_dynamic_sfa_fp16) +# ──────────────────────────────────────────────────────────────────── +def ue4m3_round(x: torch.Tensor) -> torch.Tensor: + """Round positive tensor to the nearest UE4M3 value (E4M3 without sign, + exp bias 7, 3 mantissa bits, subnormals below 2^-6). Approximation for + emulation mode; the kernel's exact rounding may differ at bin edges.""" + x = x.clamp_min(1e-12) + e = torch.floor(torch.log2(x)) + # normals: value = 2^E * (1 + M/8), M in 0..7 + base = torch.pow(2.0, e) + m = torch.round(x / base - 1.0).clamp(0, 8) + overflow = m == 8 + e = e + overflow.float() + m = m * (~overflow).float() + normal = torch.pow(2.0, e) * (1.0 + m / 8.0) + # subnormals: step 2^-9 + sub = torch.round(x / 2 ** -9) * 2 ** -9 + return torch.where(x < 2 ** -6, sub, normal).clamp(max=480.0) + + +def e0m3_emulate(t: torch.Tensor) -> torch.Tensor: + """Per-16 dynamic E0M3 fake-quant along the last dim. Returns fp32 + dequantized tensor with the same shape.""" + *lead, d = t.shape + assert d % 16 == 0 + v = t.float().reshape(-1, d // 16, 16) + scale = ue4m3_round(v.abs().amax(dim=-1, keepdim=True) / 7.0) + # all-zero blocks round to scale 0; clamp to the smallest UE4M3 + # subnormal so emulation never divides by zero (kernel writes a real + # scale here, exact corner behavior is hardware-specific) + scale = scale.clamp_min(2 ** -9) + q = torch.clamp(torch.round(v / scale), -7, 7) + return (q * scale).reshape(*lead, d) + + +# ──────────────────────────────────────────────────────────────────── +# Kernel path (requires flash_rt_fp4, i.e. Thor) +# ──────────────────────────────────────────────────────────────────── +def e0m3_kernel_gemm(a_fp16: torch.Tensor, b_fp16: torch.Tensor, + fvk_fp4) -> torch.Tensor: + """Quantize A[M,K] and B[N,K] with the real E0M3 kernels and run the + tcgen05 block-scaled GEMM (a_format=0). Returns fp16 D[M, N].""" + a_fp16 = a_fp16.contiguous() + b_fp16 = b_fp16.contiguous() + m, k = a_fp16.shape + n, kb = b_fp16.shape + assert k == kb and k % 16 == 0 + + a_packed = torch.empty(m, k // 2, dtype=torch.uint8, device="cuda") + a_sfa = torch.zeros(fvk_fp4.sfa_size_bytes(m, k, False), + dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + a_fp16.data_ptr(), a_packed.data_ptr(), a_sfa.data_ptr(), + m, k, False, 0) + if rc != 0: + raise RuntimeError(f"A quantize failed rc={rc}") + + b_packed = torch.empty(n, k // 2, dtype=torch.uint8, device="cuda") + b_sfb = torch.zeros(fvk_fp4.sfa_size_bytes(n, k, True), + dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + b_fp16.data_ptr(), b_packed.data_ptr(), b_sfb.data_ptr(), + n, k, True, 0) + if rc != 0: + raise RuntimeError(f"B quantize failed rc={rc}") + + d = torch.empty(m, n, dtype=torch.float16, device="cuda") + rc = fvk_fp4.cutlass_fp4_gemm_e0m3w( + a_packed.data_ptr(), a_sfa.data_ptr(), + b_packed.data_ptr(), b_sfb.data_ptr(), d.data_ptr(), + m, n, k, 1.0, 0.0, 0, 0) + if rc != 0: + raise RuntimeError(f"cutlass_fp4_gemm_e0m3w failed rc={rc:#x}") + torch.cuda.synchronize() + return d + + +def validate_artifact(artifact: Mapping, layer: str, n: int, k: int, + fvk_fp4) -> tuple[Mapping, Mapping]: + if not isinstance(artifact, Mapping) or artifact.get("format") != "omega_e0m3_v1": + raise ValueError("artifact format must be 'omega_e0m3_v1'") + if artifact.get("schema_version") != 1: + raise ValueError( + f"artifact schema_version must be 1, got {artifact.get('schema_version')!r}") + weights = artifact.get("weights") + aux = artifact.get("aux") + if not isinstance(weights, Mapping) or not isinstance(aux, Mapping): + raise ValueError("artifact weights and aux must be mappings") + selected = artifact.get("selected_layers") + count = artifact.get("selected_record_count") + if not isinstance(selected, list) or count != len(selected): + raise ValueError("artifact selected layer metadata is inconsistent") + if set(selected) != set(weights) or set(selected) != set(aux): + raise ValueError("artifact weights/aux do not cover every selected layer") + if layer not in weights or layer not in aux: + raise ValueError(f"artifact does not contain layer {layer!r}") + entry = weights[layer] + aux_entry = aux[layer] + if not isinstance(entry, Mapping) or not isinstance(aux_entry, Mapping): + raise ValueError(f"artifact layer {layer!r} entries must be mappings") + if entry.get("N") != n or entry.get("K") != k: + raise ValueError( + f"artifact layer shape mismatch: expected N={n}, K={k}, got " + f"N={entry.get('N')!r}, K={entry.get('K')!r}") + packed = entry.get("packed") + sfb = entry.get("sfb") + if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8 \ + or tuple(packed.shape) != (n, k // 2): + raise ValueError( + f"artifact packed must be uint8 shape ({n},{k // 2})") + expected_sfb = fvk_fp4.sfa_size_bytes(n, k, True) + if not isinstance(sfb, torch.Tensor) or sfb.dtype != torch.uint8 \ + or sfb.numel() != expected_sfb: + raise ValueError( + f"artifact sfb must be uint8 with {expected_sfb} bytes") + fold = artifact.get("fold") + if fold not in {"none", "mean", "actnorm"} or aux_entry.get("fold") != fold: + raise ValueError("artifact fold metadata is missing or inconsistent") + return entry, aux_entry + + +def e0m3_artifact_gemm(a_fp16: torch.Tensor, entry: Mapping, + fvk_fp4, *, alpha: float = 1.0) -> torch.Tensor: + """Quantize A, then consume artifact packed/SFB without re-quantizing B.""" + a_fp16 = a_fp16.contiguous() + m, k = a_fp16.shape + n = int(entry["N"]) + a_packed = torch.empty(m, k // 2, dtype=torch.uint8, device="cuda") + a_sfa = torch.zeros(fvk_fp4.sfa_size_bytes(m, k, False), + dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + a_fp16.data_ptr(), a_packed.data_ptr(), a_sfa.data_ptr(), + m, k, False, 0) + if rc != 0: + raise RuntimeError(f"A quantize failed rc={rc}") + b_packed = entry["packed"].to(device="cuda", non_blocking=False).contiguous() + b_sfb = entry["sfb"].to(device="cuda", non_blocking=False).contiguous() + d = torch.empty(m, n, dtype=torch.float16, device="cuda") + rc = fvk_fp4.cutlass_fp4_gemm_e0m3w( + a_packed.data_ptr(), a_sfa.data_ptr(), + b_packed.data_ptr(), b_sfb.data_ptr(), d.data_ptr(), + m, n, k, alpha, 0.0, 0, 0) + if rc != 0: + raise RuntimeError(f"artifact cutlass_fp4_gemm_e0m3w failed rc={rc:#x}") + torch.cuda.synchronize() + return d + + +# ──────────────────────────────────────────────────────────────────── +def cosine_stats(a: torch.Tensor, b: torch.Tensor) -> tuple: + """(global cos, per-row cos mean, per-row cos min), fp32 inputs.""" + a = a.float() + b = b.float() + glob = (torch.dot(a.flatten(), b.flatten()) / ( + a.norm() * b.norm())).item() + per = torch.nn.functional.cosine_similarity(a, b, dim=-1) + return glob, per.mean().item(), per.min().item() + + +def report(tag: str, a: torch.Tensor, b: torch.Tensor) -> None: + g, mean, mn = cosine_stats(a, b) + print(f" {tag:<28} global {g:.6f} per-token mean {mean:.6f} " + f"min {mn:.6f}") + + +def main() -> int: + args = parse_args() + if args.artifact and args.mode != "kernel": + print("error: --artifact requires --mode kernel", file=sys.stderr) + return 2 + pack = torch.load(args.pack, map_location="cpu", weights_only=True) + if args.layer not in pack: + print(f"error: layer '{args.layer}' not in pack", file=sys.stderr) + return 2 + rec = pack[args.layer] + table = rec["act_scale_table"].float() + if not 0 <= args.step < table.shape[0]: + print(f"error: --step {args.step} out of range " + f"(table has {table.shape[0]} steps)", file=sys.stderr) + return 2 + s_t = table[args.step] + s_mean = table.mean(dim=0) + + w = rec["weight_res_q"].float() # (out, in), rotated+permuted domain + out_f, in_f = w.shape + print(f"layer: {args.layer} N(out)={out_f} K(in)={in_f} " + f"table=({table.shape[0]},{table.shape[1]}) step={args.step}") + + # Synthetic activations with realistic per-channel heterogeneity: + # lognormal gains plus a few strong outlier channels (DuQuant's target). + g = torch.Generator().manual_seed(args.seed) + gains = torch.exp(torch.randn(in_f, generator=g)) + outlier_idx = torch.randperm(in_f, generator=g)[: in_f // 128 + 1] + gains[outlier_idx] *= 10.0 + x = (torch.randn(args.tokens, in_f, generator=g) * gains).half() + + dev = args.device or ("cuda" if args.mode == "kernel" else "cpu") + x = x.to(dev) + x2 = apply_input_rotation( + x, rec["duquant_rotation_perm"], + rec["duquant_rotation_blocks"].to(dev)) + w = w.to(dev) + s_t = s_t.to(dev) + s_mean = s_mean.to(dev) + + # Calibrate synthetic activations to the pack's scale table: the table + # is q99.9(|x2|)/7 on real (rotated) activations, so rescale synthetic + # x2 per channel to match. Without this the fake-quant clips almost + # everything and the comparison measures clipping, not format migration. + q999 = torch.quantile(x2.abs().float(), 0.999, dim=0).clamp_min(1e-8) + x2 = (x2.float() * (7.0 * s_t / q999)).half() + + # References (fp32 accumulate). + y_fp = x2.float() @ w.t() + y_omega = fake_quant_omega(x2.float(), s_t) @ w.t() + print("\nreferences:") + report("omega vs fp (own quant cost)", y_omega, y_fp) + + if args.mode == "emulate": + y_s0 = e0m3_emulate(x2) @ e0m3_emulate(w).t() + y_s1 = (e0m3_emulate((x2.float() / s_t).half().float()) + @ e0m3_emulate(w * s_mean).t()) + else: + if not torch.cuda.is_available(): + print("error: --mode kernel requires CUDA", file=sys.stderr) + return 2 + try: + import flash_rt.flash_rt_fp4 as fvk_fp4 + except ImportError: + print("error: flash_rt_fp4 not importable — run on Thor", + file=sys.stderr) + return 2 + if args.artifact: + artifact = torch.load( + args.artifact, map_location="cpu", weights_only=True) + try: + entry, aux_entry = validate_artifact( + artifact, args.layer, out_f, in_f, fvk_fp4) + except ValueError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + fold = artifact["fold"] + if fold == "none": + artifact_input, alpha = x2, 1.0 + elif fold == "mean": + artifact_input, alpha = (x2.float() / s_t).half(), 1.0 + else: + static = aux_entry.get("act_scale_static") + alpha = aux_entry.get("act_out_scale") + if not isinstance(static, torch.Tensor) \ + or tuple(static.shape) != (in_f,) or not isinstance(alpha, float): + print("error: invalid actnorm metadata", file=sys.stderr) + return 2 + artifact_input = (x2.float() / static.to(x2.device)).half() + y_artifact = e0m3_artifact_gemm( + artifact_input, entry, fvk_fp4, alpha=alpha).float() + artifact_cos = cosine_stats(y_artifact, y_fp)[0] + print("\nconverted artifact round-trip:") + report(f"artifact ({fold}) vs omega", y_artifact, y_omega) + report(f"artifact ({fold}) vs fp", y_artifact, y_fp) + if artifact_cos < args.min_artifact_cos: + print( + f"error: artifact cosine {artifact_cos:.6f} is below " + f"{args.min_artifact_cos:.6f}", file=sys.stderr) + return 1 + return 0 + y_s0 = e0m3_kernel_gemm(x2, w.half(), fvk_fp4).float() + y_s1 = e0m3_kernel_gemm((x2.float() / s_t).half(), + (w * s_mean).half(), fvk_fp4).float() + + print(f"\nvariants vs references (mode={args.mode}):") + report("S0 (drop table) vs omega", y_s0, y_omega) + report("S1 (mean fold) vs omega", y_s1, y_omega) + report("S0 vs fp", y_s0, y_fp) + report("S1 vs fp", y_s1, y_fp) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/convert_omega_pack_e0m3.py b/tools/convert_omega_pack_e0m3.py new file mode 100644 index 00000000..d5a6f720 --- /dev/null +++ b/tools/convert_omega_pack_e0m3.py @@ -0,0 +1,315 @@ +#!/usr/bin/env python3 +r"""Offline converter: Omega-QVLA dit_svdquant_v1 pack -> FlashRT E0M3 weights. + +Reads an Omega-QVLA quantized pack (see docs/omega_pack_e0m3.md for the +record schema) and re-quantizes every `weight_res_q` tensor into the +FlashRT SM110 E0M3 operand format: packed 4-bit elements [N, K/2] plus +tile-interleaved UE4M3 SFB scales (per-16, amax/7), via the +`quantize_e0m3_dynamic_sfa_fp16` kernel. No GPTQ bitstream decoding is +needed — the pack stores weights as dequantized fp16. + +Scale-fold strategies (--fold): + none : B = e0m3(W). The activation-side per-channel calibration table is + not represented anywhere (strategy S0 in the doc). + mean : B = e0m3(W * diag(s_mean)), s_mean = act_scale_table.mean(dim=0). + The runtime must then divide activations by s_t per step before + quantization (strategy S1). Exact for the mean step; residual is + the table's step-to-step spread. BROKEN on hardware: raw s_mean + (~1e-2) shrinks weight columns, pressing per-16 block scales + below the UE4M3 subnormal floor (2^-9). Kept for reference. + actnorm : floor-safe S1. Decompose s_mean = c * r with c = geomean + (per-layer scalar) and r = s_mean / c (geomean 1, O(1) entries): + weights fold r (magnitudes preserved, no floor issue), + activations are divided by s_mean at runtime (static — no + per-step dispatch), and c is absorbed into the GEMM alpha. + Identity: (x/s_mean) @ (W*r)^T * c == x @ W^T. + +Auxiliary tensors needed by a runtime consumer (input/output rotations, +permutation, scale tables) are copied through unchanged into the output. + +Requires: CUDA + the compiled flash_rt_fp4 extension (i.e. run on Thor; +the quantize kernels are plain CUDA but the GEMM they feed is SM110). + +Usage: + python tools/convert_omega_pack_e0m3.py \ + --pack /path/to/Omega-QVLA/packs_hf/pi05_long/quantized.pt \ + --out pi05_long_e0m3.pt --fold none + # subset for bring-up: + python tools/convert_omega_pack_e0m3.py --pack ... --out /tmp/one.pt \ + --fold none --layer-regex 'layers\.0\.self_attn\.q_proj' +""" + +from __future__ import annotations + +import argparse +import os +import re +import sys +import tempfile +import time +from collections.abc import Mapping +from pathlib import Path + +import torch + +OUTPUT_FORMAT = "omega_e0m3_v1" + +# Record fields copied verbatim into the output's per-layer aux entry. +AUX_TENSORS = ( + "duquant_rotation_blocks", + "duquant_rotation_perm", + "duquant_rotation_out_blocks", + "act_scale_table", +) +AUX_SCALARS = ("weight_bits", "a_bits", "in_features", "out_features", "rank") +EXPECTED_FULL_PACK_RECORDS = 252 + + +def _tensor(rec: Mapping, name: str, field: str) -> torch.Tensor: + value = rec.get(field) + if not isinstance(value, torch.Tensor): + raise ValueError(f"{name}: {field} must be a tensor") + return value + + +def validate_record(name: str, rec: Mapping) -> tuple[int, int]: + """Validate the complete rank-0 ``dit_svdquant_v1`` record contract.""" + if not isinstance(rec, Mapping): + raise ValueError(f"{name}: record must be a mapping") + if rec.get("format") != "dit_svdquant_v1": + raise ValueError( + f"{name}: unsupported format {rec.get('format')!r}; " + "expected 'dit_svdquant_v1'") + if rec.get("rank") != 0: + raise ValueError(f"{name}: only rank=0 is supported, got {rec.get('rank')!r}") + if rec.get("weight_bits") != 4 or rec.get("a_bits") != 4: + raise ValueError( + f"{name}: expected weight_bits=a_bits=4, got " + f"{rec.get('weight_bits')!r}/{rec.get('a_bits')!r}") + + weight = _tensor(rec, name, "weight_res_q") + if weight.ndim != 2 or not weight.is_floating_point(): + raise ValueError(f"{name}: weight_res_q must be a floating [N,K] tensor") + if not bool(torch.isfinite(weight).all()): + raise ValueError(f"{name}: weight_res_q contains NaN or Inf") + n, k = weight.shape + if n <= 0 or k <= 0 or n % 64 or k % 64: + raise ValueError( + f"{name}: N and K must be positive multiples of 64, got N={n}, K={k}") + if rec.get("out_features") != n or rec.get("in_features") != k: + raise ValueError( + f"{name}: metadata shape mismatch: tensor=({n},{k}), " + f"out_features/in_features={rec.get('out_features')!r}/" + f"{rec.get('in_features')!r}") + + lowrank_a = _tensor(rec, name, "lowrank_A") + lowrank_b = _tensor(rec, name, "lowrank_B") + if tuple(lowrank_a.shape) != (n, 0) or tuple(lowrank_b.shape) != (k, 0): + raise ValueError( + f"{name}: rank-0 lowrank shapes must be ({n},0)/({k},0), got " + f"{tuple(lowrank_a.shape)}/{tuple(lowrank_b.shape)}") + + table = _tensor(rec, name, "act_scale_table") + if table.ndim != 2 or table.shape[0] < 1 or table.shape[1] != k: + raise ValueError( + f"{name}: act_scale_table must have shape [steps,{k}], got " + f"{tuple(table.shape)}") + if not table.is_floating_point() or not bool(torch.isfinite(table).all()) \ + or not bool((table > 0).all()): + raise ValueError(f"{name}: act_scale_table must contain finite positive values") + + rotation_in = _tensor(rec, name, "duquant_rotation_blocks") + rotation_out = _tensor(rec, name, "duquant_rotation_out_blocks") + if tuple(rotation_in.shape) != (k // 64, 64, 64): + raise ValueError( + f"{name}: input rotation shape must be ({k // 64},64,64), got " + f"{tuple(rotation_in.shape)}") + if tuple(rotation_out.shape) != (n // 64, 64, 64): + raise ValueError( + f"{name}: output rotation shape must be ({n // 64},64,64), got " + f"{tuple(rotation_out.shape)}") + if not bool(torch.isfinite(rotation_in).all()) \ + or not bool(torch.isfinite(rotation_out).all()): + raise ValueError(f"{name}: rotation tensors contain NaN or Inf") + + perm = _tensor(rec, name, "duquant_rotation_perm") + if perm.dtype != torch.int64 or tuple(perm.shape) != (k,): + raise ValueError( + f"{name}: permutation must be int64 shape ({k},), got " + f"{perm.dtype} {tuple(perm.shape)}") + if not torch.equal(torch.sort(perm).values, torch.arange(k, dtype=torch.int64)): + raise ValueError(f"{name}: duquant_rotation_perm is not a permutation of [0,{k})") + return n, k + + +def expected_record_count(meta: Mapping, *, layer_regex: str, + selected_count: int, + explicit: int | None) -> int: + if explicit is not None: + if explicit < 1: + raise ValueError("--expected-records must be positive") + return explicit + recipe = meta.get("recipe", "") + if layer_regex or (isinstance(recipe, str) + and recipe.startswith("synthetic fixture")): + return selected_count + return EXPECTED_FULL_PACK_RECORDS + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + p.add_argument("--pack", required=True, help="input Omega quantized.pt") + p.add_argument("--out", required=True, help="output .pt path") + p.add_argument("--fold", choices=("none", "mean", "actnorm"), + default="none", + help="scale-table fold strategy (default: none = S0; " + "mean/actnorm are ablation-only, see docstring)") + p.add_argument("--layer-regex", default="", + help="only convert layers matching this regex") + p.add_argument("--keep-fp16", action="store_true", + help="also store the (possibly folded) fp16 weight, " + "for offline reference checks") + p.add_argument("--expected-records", type=int, + help="required converted record count; defaults to 252 for " + "a full pack and to the selected count for fixtures/subsets") + p.add_argument("--validate-only", action="store_true", + help="validate schema and record coverage without CUDA") + return p.parse_args() + + +def main() -> int: + args = parse_args() + if not args.validate_only \ + and Path(args.pack).resolve() == Path(args.out).resolve(): + print("error: --out must not overwrite the source pack", file=sys.stderr) + return 2 + + pack = torch.load(args.pack, map_location="cpu", weights_only=True) + if not isinstance(pack, Mapping): + print("error: pack must be a mapping", file=sys.stderr) + return 2 + meta = pack.get("__meta__", {}) + if not isinstance(meta, Mapping): + print("error: __meta__ must be a mapping", file=sys.stderr) + return 2 + names = sorted(k for k in pack if k != "__meta__") + if args.layer_regex: + rx = re.compile(args.layer_regex) + names = [n for n in names if rx.search(n)] + if not names: + print("error: no layers matched", file=sys.stderr) + return 2 + + try: + expected = expected_record_count( + meta, layer_regex=args.layer_regex, selected_count=len(names), + explicit=args.expected_records) + if len(names) != expected: + raise ValueError( + f"record coverage mismatch: selected {len(names)}, expected {expected}") + shapes = {name: validate_record(name, pack[name]) for name in names} + except (KeyError, TypeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(f"validated {len(names)}/{expected} records") + if args.validate_only: + return 0 + + if not torch.cuda.is_available(): + print("error: CUDA is required (quantize kernels run on GPU)", + file=sys.stderr) + return 2 + try: + import flash_rt.flash_rt_fp4 as fvk_fp4 + except ImportError: + print("error: flash_rt_fp4 extension not importable — run this on a " + "machine with FlashRT built (Thor)", file=sys.stderr) + return 2 + + device = torch.device("cuda") + weights: dict = {} + aux: dict = {} + t0 = time.time() + for i, name in enumerate(names): + rec = pack[name] + w = rec["weight_res_q"].to(device=device, dtype=torch.float16, + non_blocking=False).contiguous() + table = rec["act_scale_table"].float() + act_out_scale = None + if args.fold == "mean": + s_mean = table.mean(dim=0) # (in,) + w = (w * s_mean.to(device=device, dtype=torch.float16) + .unsqueeze(0)).contiguous() + elif args.fold == "actnorm": + s_mean = table.mean(dim=0).clamp_min(1e-12) # (in,) + c = float(torch.exp(torch.log(s_mean).mean())) + r = s_mean / c # geomean 1, O(1) entries + w = (w * r.to(device=device, dtype=torch.float16) + .unsqueeze(0)).contiguous() + act_out_scale = c + n, k = shapes[name] + + packed = torch.empty(n, k // 2, dtype=torch.uint8, device=device) + # Zero-init: tile-interleaved SFB pads K to 64-element atoms and the + # kernel never writes padding entries (see fp4_utils.py). + sfb = torch.zeros(fvk_fp4.sfa_size_bytes(n, k, True), + dtype=torch.uint8, device=device) + rc = fvk_fp4.quantize_e0m3_dynamic_sfa_fp16( + w.data_ptr(), packed.data_ptr(), sfb.data_ptr(), n, k, True, 0) + if rc != 0: + raise RuntimeError(f"quantize_e0m3_dynamic_sfa_fp16 failed on " + f"{name}: rc={rc}") + + entry = {"packed": packed.cpu(), "sfb": sfb.cpu(), "N": n, "K": k} + if args.keep_fp16: + entry["weight_fp16_folded"] = w.cpu() + weights[name] = entry + + aux_entry = {f: rec[f].clone() for f in AUX_TENSORS if f in rec} + aux_entry.update({f: rec[f] for f in AUX_SCALARS if f in rec}) + aux_entry["fold"] = args.fold + if args.fold == "actnorm": + # Consumer contract: divide activations by act_scale_static + # (post-rotation, pre-quantize) and pass act_out_scale as the + # GEMM alpha. See the --fold actnorm note in the docstring. + aux_entry["act_scale_static"] = s_mean.clone() + aux_entry["act_out_scale"] = act_out_scale + aux[name] = aux_entry + + if (i + 1) % 21 == 0 or i + 1 == len(names): + print(f"[{i + 1}/{len(names)}] {name} N={n} K={k} " + f"({time.time() - t0:.1f}s)") + + torch.cuda.synchronize() + if set(weights) != set(names) or set(aux) != set(names): + raise RuntimeError( + "internal coverage error: not all validated records were converted") + out = { + "format": OUTPUT_FORMAT, + "schema_version": 1, + "source_pack_meta": meta, + "source_record_count": len(pack) - int("__meta__" in pack), + "selected_record_count": len(names), + "selected_layers": names, + "fold": args.fold, + "weights": weights, + "aux": aux, + } + out_path = Path(args.out) + out_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + prefix=out_path.name + ".", suffix=".tmp", + dir=out_path.parent, delete=False) as tmp: + tmp_path = Path(tmp.name) + try: + torch.save(out, tmp_path) + os.replace(tmp_path, out_path) + finally: + tmp_path.unlink(missing_ok=True) + print(f"wrote {args.out}: {len(weights)} layers, fold={args.fold}, " + f"{time.time() - t0:.1f}s total") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/gen_omega_pack_fixture.py b/tools/gen_omega_pack_fixture.py new file mode 100644 index 00000000..0deeefc4 --- /dev/null +++ b/tools/gen_omega_pack_fixture.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +"""Generate a synthetic miniature dit_svdquant_v1 pack (fixture). + +Lets anyone exercise the converter + consumer round-trip WITHOUT the +real 4.8GB Omega-QVLA pack: + + python tools/gen_omega_pack_fixture.py --out /tmp/fixture_pack.pt + python tools/convert_omega_pack_e0m3.py --pack /tmp/fixture_pack.pt \ + --out /tmp/fixture_e0m3.pt --fold none --keep-fp16 # Thor + python tools/check_omega_e0m3_layer.py --pack /tmp/fixture_pack.pt \ + --artifact /tmp/fixture_e0m3.pt --mode kernel # Thor + +Design notes (what makes the fixture a real test and not a toy): + +- Schema is field-for-field identical to the real pack (checked against + packs_hf/pi05_long/quantized.pt): format tag, fp16 `weight_res_q`, + empty rank-0 lowrank tensors, fp16 64x64 rotation blocks, int64 perm, + float32 act_scale_table (expert style: T=10 rows; paligemma style: + T=1), plus the calibration metadata scalars. +- Rotations are RANDOM ORTHOGONAL (QR of gaussian), not identity — + otherwise the consumer's bmm rotate path would pass trivially. +- Weights carry per-input-channel lognormal scale outliers, the regime + that stresses E0M3 per-16 quantization and the UE4M3 subnormal floor + (the S1 failure mode in docs/omega_pack_e0m3.md). +- Shapes mix one full-size realistic layer (16384x2048) with small ones + (256x1024) to cover kernel edge sizes; all N,K are multiples of 64 as + the 64x64 rotation blocks require. + +Pure CPU, deterministic under --seed. Output ~85MB with defaults. +""" + +from __future__ import annotations + +import argparse + +import torch + +# (name, N=out, K=in, table_rows) — names/shapes mirror the real pack. +LAYERS = [ + ("paligemma_with_expert.gemma_expert.model.layers.0.self_attn.q_proj", + 2048, 1024, 10), + ("paligemma_with_expert.gemma_expert.model.layers.0.self_attn.v_proj", + 256, 1024, 10), + ("paligemma_with_expert.gemma_expert.model.layers.0.mlp.down_proj", + 1024, 4096, 10), + ("paligemma_with_expert.paligemma.model.language_model.layers.0." + "mlp.gate_proj", 16384, 2048, 1), +] + + +def _orthogonal_blocks(n_blocks: int, g: torch.Generator) -> torch.Tensor: + """(nb, 64, 64) random orthogonal blocks, fp16 (as the real pack).""" + a = torch.randn(n_blocks, 64, 64, generator=g) + q, r = torch.linalg.qr(a) + # QR sign convention: make diag(R) positive so Q is uniform-ish. + sign = torch.sign(torch.diagonal(r, dim1=1, dim2=2)) + q = q * sign[:, None, :] + return q.to(torch.float16) + + +def make_record(name: str, n: int, k: int, t_rows: int, + g: torch.Generator) -> dict: + assert k % 64 == 0 and n % 64 == 0, "64x64 rotation blocks need it" + # Weight: gaussian body x per-input-channel lognormal outliers. + chan_scale = torch.exp(torch.randn(k, generator=g) * 0.8) + w = (torch.randn(n, k, generator=g) + * chan_scale[None, :] * 0.02).to(torch.float16) + + # Act scale table: positive, ~1.0, mild per-step + per-channel jitter. + table = torch.exp(torch.randn(t_rows, k, generator=g) * 0.10 + + torch.randn(k, generator=g)[None, :] * 0.05) + + return { + "format": "dit_svdquant_v1", + "weight_res_q": w, + "lowrank_A": torch.empty(n, 0, dtype=torch.float16), + "lowrank_B": torch.empty(k, 0, dtype=torch.float16), + "act_scale_table": table.float(), + "duquant_rotation_blocks": _orthogonal_blocks(k // 64, g), + "duquant_rotation_perm": torch.randperm(k, generator=g), + "duquant_rotation_out_blocks": _orthogonal_blocks(n // 64, g), + "weight_bits": 4, + "a_bits": 4, + "rank": 0, + "in_features": k, + "out_features": n, + "n_calib_total": 100 * t_rows, + "n_calib_per_step": [100] * t_rows, + "act_percentile": 99.9, + "gptq_damp_percent": 0.05, + } + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--out", required=True, help="output .pt path") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + g = torch.Generator().manual_seed(args.seed) + pack = {"__meta__": { + "recipe": "synthetic fixture (tools/gen_omega_pack_fixture.py)", + "suite": "fixture", + "fresh": True, + "seed": args.seed, + }} + for name, n, k, t in LAYERS: + pack[name] = make_record(name, n, k, t, g) + print(f" {name} N={n} K={k} table=({t},{k})") + + torch.save(pack, args.out) + print(f"wrote {args.out}: {len(LAYERS)} layers, seed={args.seed}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())