Function build_cross_layer_pairs

scripts/train_predictor.py:285–311

Build Fate-style cross-layer training pairs: (h[l], experts[l+1]). Records are token-major: layers 0..num_layers-1 per token in sequence. For each token t and each layer l in 0..num_layers-2, we pair: input: hiddens[t, l] (gate input at layer l) target: experts[t, l+1] (

(layers, hiddens, experts, num_layers)

def build_cross_layer_pairs(layers, hiddens, experts, num_layers):
    """Build Fate-style cross-layer training pairs: (h[l], experts[l+1]).

    Records are token-major: layers 0..num_layers-1 per token in sequence.
    For each token t and each layer l in 0..num_layers-2, we pair:
      input:  hiddens[t, l]      (gate input at layer l)
      target: experts[t, l+1]   (expert activations at layer l+1)
      label:  l+1                (target layer index for layer embedding)

    Returns (X, L, E) arrays with shape (N_tokens*(num_layers-1), ...).
    """
    n_total = len(layers)
    assert n_total % num_layers == 0, \
        f"Expected multiple of {num_layers} records, got {n_total}"
    n_tokens = n_total // num_layers

    X_list, L_list, E_list = [], [], []
    for t in range(n_tokens):
        base = t * num_layers
        for l in range(num_layers - 1):
            X_list.append(hiddens[base + l])       # h[l]
            L_list.append(layers[base + l + 1])    # target layer l+1
            E_list.append(experts[base + l + 1])   # experts at l+1

    return (np.stack(X_list),
            np.array(L_list, dtype=np.int32),
            np.stack(E_list))


def analyze_cosine_similarity(layers, hiddens, num_layers):