Lightning-AI/lightning-thunder

Make autocast a Transform

Offen

#1.122 geöffnet am 07.09.2024

 (2 Kommentare) (0 Reaktionen) (0 zugewiesene Personen)Python (114 Forks)github user discovery
enhancementgood first issue

Repository-Metriken

Stars
 (1.460 Sterne)
PR-Merge-Metriken
 (Keine gemergten PRs in 30 T)

Beschreibung

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.

Contributor Guide