Compilation

Aperçu

Pytorch 2.0 a introduit torch.compile, une fonctionnalité puissante qui accélère l’exécution du code PyTorch en le compilant à la volée (JIT) en noyaux optimisés. Les principales fonctionnalités de torch.compile incluent :

• Amélioration des performances : Accélère considérablement l'exécution du modèle en optimisant le graphe de calcul.
• Facilité d'utilisation : Nécessite des modifications de code minimales pour être mis en œuvre, ce qui le rend très accessible.
• Compatibilité : Fonctionne de manière transparente avec le code et les modèles PyTorch existants.

Lorsqu’il est utilisé avec Accelerate, torch.compile s’intègre sans accroc aux flux de travail d’entraînement distribué, vous permettant de tirer parti à la fois de l’exécution distribuée et des optimisations de compilation.

La première exécution du code compilé prend généralement plus de temps car elle inclut le temps de compilation, mais les exécutions suivantes sont nettement plus rapides. Pour des performances optimales dans différents scénarios, torch.compile propose divers modes tels que "default", "reduce-overhead" (qui utilise les graphes CUDA pour réduire davantage les surcoûts) et "max-autotune" (qui effectue un autotuning approfondi pour trouver les meilleurs noyaux pour votre modèle).

Utiliser torch.compile avec Accelerate

Accelerate fournit TorchDynamoPlugin pour une intégration simple et transparente de torch.compile dans vos scripts d’entraînement.

from accelerate import Accelerator
from accelerate.utils import TorchDynamoPlugin

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    backend="inductor",  # Options: "inductor", "aot_eager", "aot_nvfuser", etc.
    mode="default",      # Options: "default", "reduce-overhead", "max-autotune"
    fullgraph=True,
    dynamic=False
)

# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply torch.compile to your model
model = accelerator.prepare(model)

Il est compatible avec toutes les autres fonctionnalités et plugins d'Accelerate, y compris la précision mixte, l'entraînement distribué (DDP, FSDP, Deepspeed), etc.

Compilation régionale

Plutôt que de tenter de compiler le modèle entier, ce qui présente généralement un espace de problèmes d’optimisation important, la compilation régionale cible les blocs répétés de la même classe et les compile séquentiellement afin de tirer parti du cache du compilateur. Par exemple, dans GPT2LMHeadModel, le bloc/classe répété est GPT2Block, et il est accessible via model.transformer.h[0]. Le reste du modèle (par exemple model.lm_head) est compilé séparément.

Cela nous permet d’accélérer le surcoût de compilation / le démarrage à froid de modèles tels que les LLM et les Transformers en général. Voir https://pytorch.org/tutorials/recipes/regional_compilation.html pour plus de détails.

Comment utiliser la compilation régionale

Il peut être activé en définissant use_regional_compilation=True dans la configuration de TorchDynamoPlugin :

# Configure the compilation backend
dynamo_plugin = TorchDynamoPlugin(
    use_regional_compilation=True,
    ... # other parameters
)
# Initialize accelerator with the plugin
accelerator = Accelerator(dynamo_plugin=dynamo_plugin)
# This will apply compile_regions to your model
model = accelerator.prepare(model)

Vous pouvez également utiliser directement l'utilitaire accelerate.utils.compile_regions de la même manière que vous utiliseriez torch.compile.

Avantages de la compilation régionale

Nous avons réalisé des benchmarks approfondis comparant la compilation complète et la compilation régionale à l'aide de la fonctionnalité torch.compile de PyTorch. Les résultats complets sont disponibles dans le dépôt accelerate. Les principaux constats issus de nos benchmarks sont les suivants :

1. Performances comparables : la compilation régionale offre des gains de vitesse similaires à ceux de la compilation complète, en particulier pour les modèles de grande taille.
2. Compilation plus rapide : la compilation régionale réduit considérablement le temps de compilation des modèles, ce qui en fait un choix plus efficace pour le déploiement.
3. Impact de la taille du lot : L'écart de performance entre les stratégies de compilation diminue lorsque la taille du lot augmente, ce qui indique que le surcoût lié à la compilation a un impact moindre dans ces scénarios.
4. Considérations relatives à la taille du modèle : Les avantages de la compilation régionale sont plus marqués pour les modèles de grande taille, où les gains de temps de compilation peuvent être considérables.
5. Application pratique : Pour les applications réelles, la compilation régionale est un choix pratique pour optimiser les temps de démarrage à froid de l’entraînement, en particulier lors du travail avec de grands modèles.

Conclusion

La compilation complète et régionale peut accélérer considérablement vos modèles. La compilation régionale offre un équilibre pratique entre le temps de compilation et les performances d'exécution, en particulier pour l'entraînement de grands modèles avec des tailles de lot importantes.