Open kshitij12345 opened 3 weeks ago
that's right, in autocast we don't consider device
https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/transforms.py#L3788
does this have practical impacts on target models?
AFAIK, NeMo does use autocast. With our current implementation, we may silently add conversions when user asked to apply autocast only on a certain device and if there are operations happening on both CPU and GPU in that context. Honestly, I don't think it happens in practice.
@tfogal do you know if NeMo does both CPU and GPU operations (which are affected by autocast ctx manager) within a single autocast context?
@tfogal do you know if NeMo does both CPU and GPU operations (which are affected by autocast ctx manager) within a single autocast context?
I don't know, sorry :-( @athitten might.
But I agree with you that it is unlikely---we could just not support it for now. But I would ask that we 'loudly' not support mixed-device autocast: can we check for this case and error out when it happens?
I’m 100% for failing loudly if it’s not a beaten path (and this one looks like it’s not)
triage review:
From the example below, the autocast is applied only for device cuda, however thunder.jit still applies it to CPU inputs.
Output
cc @crcrpar