torch.compile pour l'entraînement

torch.compile compile le code PyTorch en noyaux fusionnés pour accélérer son exécution. Pour l’entraînement, il trace à la fois le passage avant et le passage arrière, puis les compile en noyaux optimisés, réduisant ainsi le surcoût lié au lancement individuel des opérateurs et fusionnant les opérations afin de diminuer l’utilisation de la bande passante mémoire.

Définissez torch_compile=True dans TrainingArguments pour l’activer. L’entraînement compile à la fois le passage avant et le passage arrière, contrairement à l’inférence qui ne compile que le passage avant. La compilation a lieu à la première étape d’entraînement, il faut donc s’attendre à ce qu’elle soit nettement plus lente que les étapes suivantes.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    torch_compile=True,
    torch_compile_backend="inductor",
    torch_compile_mode="reduce-overhead",
)

Backend

Lorsqu’aucun backend n’est spécifié, TrainingArguments en sélectionne un en fonction de votre matériel. Sur la plupart des CPU et GPU, la valeur par défaut est inductor, qui compile vers des noyaux Triton via AOTAutograd et convient à la plupart des charges de travail d’entraînement. Sur Intel Gaudi (HPU), la valeur par défaut est hpu_backend. Sur AWS Trainium et Inferentia (Neuron), la valeur par défaut est neuron.

Utilisez cudagraphs pour des entrées de forme fixe.

Mode de compilation

Utilisez le tableau ci-dessous pour vous aider à choisir un mode torch.compile.

Étapes suivantes

• Consultez le guide torch.compile pour l’inférence pour obtenir des détails sur la compilation fullgraph et les benchmarks d’inférence.