Function evaluate_ppl

TransMLA_NeurIPS_2025/transmla/utils.py:206–257

Evaluate the model's perplexity on the test set using batch processing. It is expected that model is already on the correct device.

(
    model: torch.nn.Module, 
    pad_token_id: int | None, 
    testloader: DataLoader[dict[str, torch.Tensor]], 
    message: str = "Evaluating perplexity"
)

    
@torch.no_grad()
def evaluate_ppl(
    model: torch.nn.Module, 
    pad_token_id: int | None, 
    testloader: DataLoader[dict[str, torch.Tensor]], 
    message: str = "Evaluating perplexity"
) -> float:
    """
    Evaluate the model's perplexity on the test set using batch processing.
    It is expected that model is already on the correct device.
    """
    sync_gpus()

    start_time = time.time()

    model.eval()

    if pad_token_id:
        loss_fn = torch.nn.CrossEntropyLoss(reduction="none", ignore_index=pad_token_id)
    else:
        loss_fn = torch.nn.CrossEntropyLoss(reduction="none")

    nlls = []

    logging.info(message)
    for batch in tqdm(testloader, desc=message):
        logging.debug(f"Evaluating batch {len(nlls)}")
        batch = map_tensors(batch, model.model.embed_tokens.weight.device)
        logits = model(**batch, use_cache=False).logits

        # shift outputs and labels autoregressively.
        logits = logits[:, :-1, :]
        shift_labels = batch["input_ids"][:, 1:]

        # CrossEntropyLoss demands data dimension is dimension 1.
        nll = loss_fn(logits.permute(0, 2, 1), shift_labels).float()

        mask = shift_labels != loss_fn.ignore_index
        nll_means = (nll * mask).sum(dim=1) / mask.sum(dim=1)
        nlls.append(nll_means)

    nlls_tensor = torch.cat(nlls)
    ppl = torch.exp(nlls_tensor.mean())

    sync_gpus()

    elapsed = time.time() - start_time
    logging.info(
        "Time spent on evaluation: %s",
        time.strftime("%H:%M:%S.{}".format(str(elapsed % 1)[2:])[:13], time.gmtime(elapsed)),
    )

    return ppl.item()

def insert_qkv_hooks(model):
    query_hooks = []