[Afloat16]

What does this PR do?

Fixes evaluation and prediction when preprocess_logits_for_metrics returns FP8 tensors to reduce the size of cached logits. The hook is documented as running before prediction caching, but nested_numpify currently only promotes BF16. NumPy conversion of the cached FP8 tensors raises TypeError: Got unsupported ScalarType Float8_e4m3fn.

Extend that existing promotion to single-byte floating-point tensors at the NumPy boundary. FP8 tensors can remain in their compact dtype during caching; the resulting NumPy arrays use FP32. Ordinary floating-point and integer conversion paths remain unchanged.

Reproducer

This fails on main 41f9d3992023970fb5a463e70d56232506c80fbc and succeeds with this change, using the default GPT-2 configuration and the real Trainer on CPU:

import tempfile
import torch
from transformers import GPT2Config, GPT2LMHeadModel, Trainer, TrainingArguments

model = GPT2LMHeadModel(GPT2Config())
rows = [{"input_ids": torch.tensor([1, 7, 3, 9, 4, 8, 6, 2]) + i} for i in range(3)]
for row in rows:
    row["labels"] = row["input_ids"].clone()

with tempfile.TemporaryDirectory() as output_dir:
    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir=output_dir,
            use_cpu=True,
            per_device_eval_batch_size=2,
            eval_accumulation_steps=1,
            report_to=[],
        ),
        preprocess_logits_for_metrics=lambda logits, labels: logits.to(torch.float8_e4m3fn),
    )
    predictions = trainer.predict(rows).predictions
    assert predictions.shape == (3, 8, 50257)

Regression test and validation

The single added regression in the existing Trainer evaluation test file reuses get_regression_trainer. Its 14 subcases cover four FP8 formats plus BF16/FP32/FP64 controls, with and without per-step CPU accumulation. It checks both nested prediction outputs, labels, metrics and evaluation/prediction loss parity with a partial final batch.

• Original source: eight FP8 subcases fail; six dtype controls pass.
• Patched source: all 14 subcases pass.
• Four existing evaluation/prediction tests pass.
• The separate default-config GPT-2 run verifies every returned quantized logit against native model output and checks the metric and labels.
• Repository make style passes.

Native validation tests exact commit bf9a3cbe64372541b9be46b691f033301e6fc1b8 with PyTorch 2.14 CPU. This validation does not claim FP8 distributed-collective support or a change to model quantization accuracy.

[github-actions]

Thank you for your contribution 🤗!

CI Security Gate — waiting for maintainer approval

CI is approved automatically only for authors who already have a contribution merged in this repository. A maintainer will review this PR and approve CI to run on it.