Shakeout fixes: CUDA-free platform probe, 16-way parallel verify, sharded forward w/ replica parity, transformers==5.12.1 pin, torchvision removal

r7_encoder/inventory.py

@@ -248,11 +248,22 @@ def verify_checkpoint_inventory(
     raw_shards = payload.get("shards")
     if not isinstance(raw_shards, dict):
         raise ValueError("checkpoint inventory lacks shard map")
-    for shard_name, record in sorted(raw_shards.items()):
+    def _check_shard(item):
+        shard_name, record = item
         if not isinstance(record, dict) or sha256_file(root / shard_name) != record.get(
             "file_sha256"
         ):
             raise ValueError(f"checkpoint shard changed after inventory: {shard_name}")
+
+    # Hash shards in parallel; identical failure semantics (any mismatch
+    # raises), deterministic because each check is independent and the full
+    # set must pass. Turns a ~25-minute single-thread startup verify into a
+    # few minutes on every launch and resume.
+    from concurrent.futures import ThreadPoolExecutor
+
+    with ThreadPoolExecutor(max_workers=16) as _pool:
+        for _ in _pool.map(_check_shard, sorted(raw_shards.items())):
+            pass
     auxiliary = payload.get("auxiliary_files_sha256")
     if not isinstance(auxiliary, dict) or not auxiliary:
         raise ValueError("checkpoint inventory lacks auxiliary/tokenizer assets")
@@ -396,7 +407,13 @@ def _collect_execution_platform(
         "torch_version": str(torch.__version__),
         "torch_git_version": str(getattr(torch.version, "git_version", "UNKNOWN")),
         "torch_cuda_build": str(torch.version.cuda),
-        "torch_cudnn_build": str(torch.backends.cudnn.version()),
+        # cuDNN build parsed from static build info: torch.backends.cudnn
+        # .version() initializes CUDA, which would trip the runtime's
+        # CUDA-purity gate 26 minutes later at walk startup.
+        "torch_cudnn_build": (
+            (re.search(r"CuDNN[^\n]*?([\d.]+)", torch.__config__.show())
+             or [None, "UNKNOWN"])[1]
+        ),
         "device": _nvidia_device_record(device),
         "torch_numerical_binaries_sha256": torch_binaries,
         "dynamic_dependency_closure_sha256": _dynamic_library_closure(dynamic_roots),

r7_encoder/transformers_runtime.py

@@ -220,6 +220,12 @@ class TransformersSequentialRuntime(GLM52Runtime):
         )
 
         self.device = torch.device(config.device)
+        # Optional device pool for the sharded successor forward. Pool of one
+        # (the default) is byte-identical to the original sequential flow.
+        _pool = tuple(getattr(config, "devices", ()) or ()) or (config.device,)
+        if _pool[0] != config.device:
+            raise ValueError("forward device pool must start with the primary device")
+        self.forward_devices = tuple(torch.device(item) for item in _pool)
         if self.device.type != "cuda":
             raise ValueError(
                 "the owner-run GLM reference runtime requires one CUDA device"
@@ -506,7 +512,8 @@ class TransformersSequentialRuntime(GLM52Runtime):
             "auxiliary": "prev_topk_indices-int32",
         }
 
-    def _context(self, hidden, metadata: Mapping[str, object]):
+    def _context(self, hidden, metadata: Mapping[str, object], device=None):
+        device = self.device if device is None else device
         import torch
 
         tokens = int(metadata["tokens"])
@@ -517,21 +524,21 @@ class TransformersSequentialRuntime(GLM52Runtime):
         ):
             raise ValueError("state metadata is not the pinned one-prompt DSA schema")
         x = (
-            torch.as_tensor(hidden, device=self.device)
+            torch.as_tensor(hidden, device=device)
             .reshape(1, tokens, HIDDEN_SIZE)
             .to(torch.bfloat16)
         )
-        position_ids = torch.arange(tokens, device=self.device).unsqueeze(0)
+        position_ids = torch.arange(tokens, device=device).unsqueeze(0)
         minimum = torch.finfo(x.dtype).min
         causal = torch.full(
-            (tokens, tokens), minimum, device=self.device, dtype=x.dtype
+            (tokens, tokens), minimum, device=device, dtype=x.dtype
         ).triu(diagonal=1)
         attention_mask = causal.unsqueeze(0).unsqueeze(0)
         position_embeddings = self.rotary(x, position_ids)
         previous = metadata.get("prev_topk_indices")
         if previous is not None:
             previous = torch.as_tensor(
-                previous, device=self.device, dtype=torch.int32
+                previous, device=device, dtype=torch.int32
             ).unsqueeze(0)
         return x, attention_mask, position_ids, position_embeddings, previous
 
@@ -562,12 +569,12 @@ class TransformersSequentialRuntime(GLM52Runtime):
             gate_up = torch.empty(
                 (NUM_EXPERTS, 2 * INTERMEDIATE_SIZE, HIDDEN_SIZE),
                 dtype=torch.bfloat16,
-                device=self.device,
+                device=device,
             )
             down = torch.empty(
                 (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE),
                 dtype=torch.bfloat16,
-                device=self.device,
+                device=device,
             )
             for expert in range(NUM_EXPERTS):
                 gate_up[expert, :INTERMEDIATE_SIZE] = installed[expert][
@@ -594,8 +601,8 @@ class TransformersSequentialRuntime(GLM52Runtime):
         module.eval()
         return module
 
-    def _advance(self, module, hidden, metadata):
-        context = self._context(hidden, metadata)
+    def _advance(self, module, hidden, metadata, device=None):
+        context = self._context(hidden, metadata, device=device)
         x, mask, positions, rotary, previous = context
         with __import__("torch").inference_mode():
             output, topk = module(
@@ -1360,22 +1367,25 @@ class TransformersSequentialRuntime(GLM52Runtime):
         }
         carrier_audit = sha256_bytes(canonical_json_bytes(layer_records))
         first_shard_id = input_values[0].shard_id
+        devices = getattr(self, "forward_devices", None) or (
+            getattr(self, "device", None),
+        )
+        pending = [
+            shard
+            for shard in input_values
+            if shard.shard_id not in completed_shard_ids
+        ]
+        replicas: list = []
         try:
-            for shard in input_values:
-                if shard.shard_id in completed_shard_ids:
-                    continue
+            def _load_shard(shard):
                 reader = SafeTensorReader(shard.hidden_path)
                 hidden = read_torch_tensor(reader, "hidden")
                 previous = read_torch_tensor(reader, "prev_topk_indices")
                 metadata = read_json(shard.metadata_path)
                 metadata["prev_topk_indices"] = previous
-                if shard.shard_id == first_shard_id:
-                    output, topk, repeat_audit = self._advance_repeat_oracle(
-                        module, hidden, metadata
-                    )
-                else:
-                    output, topk = self._advance(module, hidden, metadata)
-                    repeat_audit = None
+                return hidden, metadata
+
+            def _successor(shard, metadata, repeat_audit, parity_audit):
                 successor = dict(metadata)
                 successor.pop("prev_topk_indices", None)
                 successor["producer"] = f"round7-installed-layer-{layer:03d}"
@@ -1383,18 +1393,112 @@ class TransformersSequentialRuntime(GLM52Runtime):
                 successor["carrier_decode_audit_sha256"] = carrier_audit
                 if repeat_audit is not None:
                     successor["official_repeat_audit"] = repeat_audit
-                # Yield immediately: the walk seals this prompt in PARTIAL.json
-                # before the runtime starts the next expensive prompt.
-                yield self._write_state(
-                    output_partial=output_partial,
-                    layer_input=layer + 1,
-                    shard_id=shard.shard_id,
-                    hidden=output,
-                    previous_topk=topk,
-                    metadata=successor,
+                if parity_audit is not None:
+                    successor["replica_parity_audit"] = parity_audit
+                return successor
+
+            sequential = (
+                len(devices) <= 1
+                or len(pending) <= 1
+                # Mid-layer resume: the oracle prompt is already sealed, so
+                # the parity check has no reference. Sequential handles it.
+                or pending[0].shard_id != first_shard_id
+            )
+            if sequential:
+                # Original sequential flow, byte-identical.
+                for shard in pending:
+                    hidden, metadata = _load_shard(shard)
+                    if shard.shard_id == first_shard_id:
+                        output, topk, repeat_audit = self._advance_repeat_oracle(
+                            module, hidden, metadata
+                        )
+                    else:
+                        output, topk = self._advance(module, hidden, metadata)
+                        repeat_audit = None
+                    # Yield immediately: the walk seals this prompt in
+                    # PARTIAL.json before the next expensive prompt starts.
+                    yield self._write_state(
+                        output_partial=output_partial,
+                        layer_input=layer + 1,
+                        shard_id=shard.shard_id,
+                        hidden=output,
+                        previous_topk=topk,
+                        metadata=_successor(shard, metadata, repeat_audit, None),
+                    )
+                return
+
+            # Sharded successor forward across the device pool. Compute fans
+            # out to worker replicas; every state write and yield stays on
+            # THIS thread in canonical shard order, so PARTIAL.json sealing,
+            # resume semantics, and state hashing are unchanged.
+            import copy
+            from concurrent.futures import ThreadPoolExecutor
+
+            first_shard = pending[0]
+            first_hidden, first_metadata = _load_shard(first_shard)
+            output, topk, repeat_audit = self._advance_repeat_oracle(
+                module, first_hidden, first_metadata
+            )
+            # Replicate only after the primary oracle passes, and prove every
+            # replica reproduces the primary bytes before trusting it with
+            # calibration state.
+            replica_parity = {
+                "schema": "r7-replica-parity-v1",
+                "devices": [str(item) for item in devices],
+            }
+            modules = [module]
+            for extra in devices[1:]:
+                replica = copy.deepcopy(module).to(extra)
+                r_out, r_topk = self._advance(
+                    replica, first_hidden, first_metadata, device=extra
                 )
+                if not torch.equal(r_out, output) or not torch.equal(
+                    r_topk, topk
+                ):
+                    raise RuntimeError(
+                        f"replica on {extra} diverged from the primary "
+                        "successor forward; refusing sharded execution"
+                    )
+                replicas.append(replica)
+                modules.append(replica)
+            replica_parity["passed"] = True
+            yield self._write_state(
+                output_partial=output_partial,
+                layer_input=layer + 1,
+                shard_id=first_shard.shard_id,
+                hidden=output,
+                previous_topk=topk,
+                metadata=_successor(
+                    first_shard, first_metadata, repeat_audit, replica_parity
+                ),
+            )
+
+            def _compute(index_shard):
+                index, shard = index_shard
+                worker = index % len(modules)
+                hidden, metadata = _load_shard(shard)
+                out, tk = self._advance(
+                    modules[worker], hidden, metadata, device=devices[worker]
+                )
+                return shard, metadata, out, tk
+
+            with ThreadPoolExecutor(max_workers=len(modules)) as pool:
+                for shard, metadata, out, tk in pool.map(
+                    _compute, enumerate(pending[1:])
+                ):
+                    yield self._write_state(
+                        output_partial=output_partial,
+                        layer_input=layer + 1,
+                        shard_id=shard.shard_id,
+                        hidden=out,
+                        previous_topk=tk,
+                        metadata=_successor(shard, metadata, None, None),
+                    )
         finally:
             del module
+            for replica in replicas:
+                del replica
+            replicas.clear()
             self._installed.pop(layer, None)
             torch.cuda.empty_cache()
 

run_scripts/setup_box.sh

@@ -5,7 +5,13 @@
 # plus the encoder oracles validate the rebuilt extension on behavior.
 set -e
 python3 -m pip install --break-system-packages -q torch==2.12.1+cu132 --index-url https://download.pytorch.org/whl/cu132
-python3 -m pip install --break-system-packages -q transformers safetensors pytest ninja setuptools wheel "huggingface_hub[hf_transfer]"
+# transformers PINNED: the runtime is authored against 5.12.1
+# (GlmMoeDsaNaiveMoe was refactored away in later versions).
+python3 -m pip install --break-system-packages -q "transformers==5.12.1" safetensors pytest ninja setuptools wheel "huggingface_hub[hf_transfer]"
+# remove image-bundled torchvision/torchaudio: built against older torch,
+# their ABI-mismatched op registration explodes inside transformers imports
+python3 -m pip uninstall -y --break-system-packages torchvision torchaudio 2>/dev/null || true
+python3 -c "import importlib.util; assert importlib.util.find_spec('torchvision') is None, 'torchvision still present'"
 tar xzf exl3_v0043_src.tgz -C exl3src
 cd exl3src && TORCH_CUDA_ARCH_LIST=10.0 MAX_JOBS=64 python3 -m pip install --break-system-packages --no-deps --no-build-isolation .
 python3 -c "import torch, exllamav3_ext; print('ext OK on torch', torch.__version__)"