grad transform introduces reduction operation with NumberProxy in axes about lightning-thunder HOT 4 OPEN

triage review —

• we think we can change the backward trace to use constant numbers instead of number proxies for now
• but we do want number proxies in the backward trace soon, so we should carefully consider their integration with nvfuser and grad, and imposing restraints like "this must be a constant value"
• maybe nvFuser's cache could detect the value of the NumberProxy?

from lightning-thunder.

so we should carefully consider their integration with nvfuser and grad, and imposing restraints like "this must be a constant value"

Imposing those constraints universally in thunder is what I'm planning to do.
When you mentioned explicitly about integration, I'm under the impression that you are trying to allow backend to propose constraints. I'm open to that discussion and have an issue opened already. That might be a bit trickier to plumb through back to prologue trace.

Meanwhile, I think an easy way to unblock us for now is to just let integration reject things when seeing NumberProxy on those that must be a constant value. i.e. reduction axes for example.

from lightning-thunder.

NumberProxy shouldn't show up in trace.

Do you think NumberProxies shouldn't be used in traces, in general, today and in the future, or only today because of how it's not dynamic and only a wrapper around a constant today?

It's possible to modify the trace to unwrap all NumberProxies into actual values by modifying backward_trace after this line https://github.com/Lightning-AI/lightning-thunder/blob/4663b87e67955edec371c6f37a763a7ec358d835/thunder/core/transforms.py#L3961C5-L3961C49.

from lightning-thunder.

Do you think NumberProxies shouldn't be used in traces, in general, today and in the future, or only today because of how it's not dynamic and only a wrapper around a constant today?

I should be more careful with my wording 😆
Yes NumberProxies should be used in the trace in general. There're some rough corners (i.e. nvfuser's handling of baking in static number from proxies), but we should just patch those.

But I'm surprised to see NumberProxy show up with thunder.jit today. IIUC, that's what TestExecutor.make_callable uses for compilation. Current default caching logic is CONSTANT_VALUES, I thought that means everything gets baked in as static and I'm not sure where the NumberProxy is coming from.

Totally not pointing finger at grad transform. But I think in general NumberProxy should be treated differently with a plain Number. i.e. A NumberProxy means dynamic value and that should be respected by transforms.

from lightning-thunder.