Lightning-AI/lightning-thunder

Make autocast a Transform

開放

#1,122 建立於 2024年9月7日

 (2 則留言) (0 個反應) (0 位負責人)Python (114 個分叉)github user discovery
enhancementgood first issue

倉庫指標

星標
 (1,460 顆星)
PR 合併指標
 (30 天內沒有已合併 PR)

描述

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.

貢獻者指南