Support converting SymTypes Node to input proxy

[kiya00]

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

What does this PR do?

While exploring the split reasons and fallback operators in issue #2169, it was found that the current implementation of get_proxy_inputs_from_node does not support SymInt types. It causes graph splits when running models like Undi95/dbrx-base and llama-moe/LLaMA-MoE-v1-3_5B-4_16.

SplitReason(reason_type=<SplitReasonType.EXCEPTION_PROXY_THUNDER_OP: 3>, info='Failed while creating proxy for node with name: expert_w1 and target: <built-in method select of type object at 0x7f10f0184760>, see exception field', exception="`make_input_proxy` received example_value which wasn't Tensor or Tuple")

from transformers import AutoConfig
import torch
from torch.testing import make_tensor
from functools import partial
import time
import argparse


model_id = "Undi95/dbrx-base"

config = AutoConfig.from_pretrained(
        model_id,
        # Scaled down for testing
        vocab_size=16,
        pad_token_id=15,
        max_position_embeddings=32,
        trust_remote_code=True,
        num_hidden_layers=1,)
batch_size = 1
seq_length = 4096
shape = (batch_size, seq_length)

from transformers import AutoModelForCausalLM

with torch.device("cuda"):
    model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16,trust_remote_code=True)
    model.requires_grad_(True)
    make = partial(
        make_tensor,
        low=0,
        high=config.vocab_size,
        device="cuda",
        dtype=torch.int64,
        requires_grad=False,
    )
    a = make(shape)
    inputs = {"input_ids": a, "labels": a}

    from thunder.dynamo import thunderfx
    cmodel = thunderfx(model)

cmodel(**inputs)
infos = cmodel._backend.subgraph_infos
print(f"{len(infos)} segmentations")
for info in infos:
    rs = info.split_reasons
    print(f"{len(rs)} split reasons")
    for i,r in enumerate(rs):
        print(f"    reason {i}: {r}")

[mruberry]

@IvanYashchuk would you like to review this?

[kiya00]

Hi @kshitij12345 @IvanYashchuk ,could you help to take a look

[kshitij12345]

Thanks @kiya00 , I just have one question.

Also, we should add a test, maybe something like the following. I think we should also verify that thunder caching works correctly and calling cfn with different idx value also returns the correct result.

import torch

from thunder.dynamo import thunderfx

def fn(x, idx):
    return torch.select(x, 0, idx)

x = torch.randn(10, 10)
idx = 0
cfn = thunderfx(fn, dynamic=True)
cfn(x, idx)

assert cfn._backend.subgraph_infos[0].split_reasons == []

[kshitij12345]

LGTM, thanks @kiya00

[kiya00]

Hi @IvanYashchuk @t-vi , could you help take a look, I think it's ready to merge

[IvanYashchuk]

Looks good!

[kiya00]

Hi @t-vi @Borda , it's ready to merge

[kiya00]

Hi @t-vi @mruberry , it's ready to merge

[mruberry]