Function load_routing_data

scripts/train_predictor.py:86–132

Load a single binary routing file into numpy arrays.

(path, hidden_dim)

def load_routing_data(path, hidden_dim):
    """Load a single binary routing file into numpy arrays."""
    data = Path(path).read_bytes()
    offset = 0
    layers = []
    hiddens = []
    experts = []

    while offset < len(data) - 8:  # at least header
        layer_idx = struct.unpack_from('<i', data, offset)[0]
        offset += 4
        K = struct.unpack_from('<i', data, offset)[0]
        offset += 4

        # Check for truncated record (e.g. server killed mid-write)
        needed = hidden_dim * 4 + K * 4
        if offset + needed > len(data):
            break

        h = np.frombuffer(data, dtype=np.float32, count=hidden_dim, offset=offset).copy()
        offset += hidden_dim * 4

        ei = np.frombuffer(data, dtype=np.int32, count=K, offset=offset).copy()
        offset += K * 4

        layers.append(layer_idx)
        hiddens.append(h)
        experts.append(ei)

    # Empty or fully-truncated file: return well-formed empty arrays rather
    # than crashing in np.stack([]) / max([]).
    if not experts:
        return (np.array([], dtype=np.int32),
                np.zeros((0, hidden_dim), dtype=np.float32),
                np.zeros((0, 0), dtype=np.int32), 0)

    layers = np.array(layers, dtype=np.int32)
    hiddens = np.stack(hiddens)
    max_K = max(len(e) for e in experts)
    # Pad with -1 (not 0): 0 is a valid expert index, so a 0 sentinel would be
    # indistinguishable from a real top-0 routing and corrupt labels/eval when
    # K varies across records. Consumers filter out negative indices.
    experts_padded = np.full((len(experts), max_K), -1, dtype=np.int32)
    for i, e in enumerate(experts):
        experts_padded[i, :len(e)] = e

    return layers, hiddens, experts_padded, max_K


def load_multiple_routing_files(paths, hidden_dim, num_layers, seed=42):