Prototyping partial graph capture

[lucaslie]

Work in Progress for partial graph capture as a prerequisite to do cudagraph capture for VLMs

Summary

There is two major changes required:

1. Nested graph inside larger nn.Module --> make all transforms aware and differentiate between transforms on the full model and on the individual subgraphs
2. Switching infra to kwargs-only from positional args for inputs. Using positional args is not a scalable way forward as the subgraphs are most certainly going to be called with positional arguments.

Other thoughts:

1. Thinking about auto-capturing subgraph based on inputs to the sub-forward function --> hard-coded right now
2. Better args/kwargs handling and adding extra flexibility there...

config.yaml to play around with

# model: meta-llama/Meta-Llama-3.1-8B-Instruct
# model: mistralai/Magistral-Small-2507
# model: Qwen/Qwen2.5-VL-7B-Instruct
# model: meta-llama/Llama-4-Scout-17B-16E-Instruct
model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
args:
  # mode: graph
  world_size: 0
  runtime: demollm
  compile_backend: torch-opt
  attn_page_size: 64
  max_input_len: 4096
  max_seq_len: 8192
  attn_backend: flashinfer
  # model_factory: AutoModelForImageTextToText
  # model_factory: AutoModelForCausalLM
  model_factory: Mistral3VLM
  skip_loading_weights: true
  model_kwargs:
    text_config:
        num_hidden_layers: 2
    # tp_plan: auto
    # tp_plan: null
    # device_map: cuda
    # num_hidden_layers: 3
    # _attn_implementation: eager
benchmark:
  enabled: false
dry_run: false
# prompt:
#   batch_size: 4
#   queries:
#     - "How big is the universe? "
#     - {"prompt": "In simple words and a single sentence, explain the concept of gravity: "}
#     # see for chat template format: https://huggingface.co/docs/transformers/en/chat_templating_multimodal
#     - - role: user
#         content:
#           - type: text
#             text: How to fix slicing in golf?
#     - - role: user
#         content:
#           - type: text
#             text: Please describe the natural scenery you see in the following images
#           - type: image
#             url: https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png
#           - type: image
#             url: https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png

[lucaslie]

see NVIDIA#8203