thunderfx : detecting parameters and buffers on thunderfx path

The FXGraph provided by Dynamo takes in Parameters and Buffers as arguments, however thunder.jit currently only determines a TensorProxy to be a parameter if it is unpacked from a Module. So, on thunderfx path, we don't tag these parameters with STATIC_MEMORY_LOCATION, leading to problem with CUDAGraphTransform and ExtraionOnlyPrologueTransform which depend on these tags.

https://github.com/Lightning-AI/lightning-thunder/blob/35ca2e9c20fdc2a2a58b7c3106ac8c617710b5a5/thunder/core/jit_ext.py#L1493-L1500

Potential Solution for Parameters : For parameters, maybe thunder.jit tag Proxies based on isinstance(obj, nn.Parameter).

Sample:

import torch

def backend(gm, sample_arg):
    gm.print_readable()
    print(sample_arg)
    return gm

model = torch.nn.Linear(2, 2)

cmodel = torch.compile(model, backend=backend)
cmodel(torch.randn(1, 2))

Output

class GraphModule(torch.nn.Module):
    def forward(self, L_fn_parameters_weight_: "f32[2, 2]", L_fn_parameters_bias_: "f32[2]", L_args_0_: "f32[1, 2]"):
        l_fn_parameters_weight_ = L_fn_parameters_weight_
        l_fn_parameters_bias_ = L_fn_parameters_bias_
        l_args_0_ = L_args_0_
        
         # File: /home/kkalambarkar/git/pytorch/torch/_dynamo/external_utils.py:31 in inner, code: return fn(*args, **kwargs)
        linear: "f32[1, 2]" = torch._C._nn.linear(l_args_0_, l_fn_parameters_weight_, l_fn_parameters_bias_);  l_args_0_ = l_fn_parameters_weight_ = l_fn_parameters_bias_ = None
        return (linear,)
        
[Parameter containing:
tensor([[ 0.6478,  0.6590],
        [-0.5319, -0.3303]], requires_grad=True), Parameter containing:
tensor([0.3209, 0.6565], requires_grad=True), tensor([[ 0.7282, -0.3549]])]
{}

Start at thunder/core/jit_ext.py lines 1493-1500 and inspect how thunder.jit identifies parameters on the thunderfx path. Reproduce the provided torch.compile example, then trace how CUDAGraphTransform and ExtraionOnlyPrologueTransform consume STATIC_MEMORY_LOCATION. Done means FXGraph parameters and buffers receive the required tags and the affected transforms no longer miss them.