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.

贡献者指南