Lightning-AI/lightning-thunder

Make autocast a Transform

Ouverte

#1 122 ouverte le 7 sept. 2024

 (2 commentaires) (0 réaction) (0 personne assignée)Python (114 forks)github user discovery
enhancementgood first issue

Métriques du dépôt

Stars
 (1 460 étoiles)
Métriques de merge PR
 (Aucune PR mergée en 30 j)

Description

Currently autocast (transforms/autocast.py) is carried over from the old-style transforms on callable and invoked like

jit(autocast(fn))

it would be good to make this a AutocastTransform to be invoked as

jit(fn, transforms=[Autocast()])

for better composability.

This should be vaguely doable with the transforms e.g. as explained in Thunder Sessions Episode 7 plus the bsym mapping provided by the existing mechanism, so it might be a good first or second issue.

Guide contributeur