Add generic TorchAOTensor extra_repr for nn.Modules

[jcaip]

Summary:

This fixes extra_repr for generic modules by introducing a new helper function, _module_extra_repr, which will add TorchAOTensor info to all the parameters in a module that are an instance of TorchAOBaseTensor to the original extra_repr of the module.

The configs supporting parameter quantization have been update to use _module_extra_repr.

Also renamed new_weight -> quantized_tensor to be more consistent.

For example, we will see the following output

    (language_model): Qwen3VLMoeTextModel(
      (embed_tokens): Embedding(151936, 2048)
      (layers): ModuleList(
        (0-47): 48 x Qwen3VLMoeTextDecoderLayer(
          (self_attn): Qwen3VLMoeTextAttention(
            (q_proj): Linear(in_features=2048, out_features=4096, bias=False, weight=Float8Tensor(self.act_quant_kwargs=QuantizeTensorToFloat8Kwargs(float8_dtype=torch.float8_e4m3fn, granularity=PerRow(dim=-1), mm_config=None, hp_value_lb=None, hp_value_ub=None, kernel_preference=<KernelPreference.AUTO: 'auto'>), self.block_size=[1, 2048], self.mm_config=Float8MMConfig(emulate=False, use_fast_accum=True, pad_inner_dim=False), self.scale.shape=torch.Size([4096, 1]), self.kernel_preference=<KernelPreference.AUTO: 'auto'>))
            (k_proj): Linear(in_features=2048, out_features=512, bias=False, weight=Float8Tensor(self.act_quant_kwargs=QuantizeTensorToFloat8Kwargs(float8_dtype=torch.float8_e4m3fn, granularity=PerRow(dim=-1), mm_config=None, hp_value_lb=None, hp_value_ub=None, kernel_preference=<KernelPreference.AUTO: 'auto'>), self.block_size=[1, 2048], self.mm_config=Float8MMConfig(emulate=False, use_fast_accum=True, pad_inner_dim=False), self.scale.shape=torch.Size([512, 1]), self.kernel_preference=<KernelPreference.AUTO: 'auto'>))
            (v_proj): Linear(in_features=2048, out_features=512, bias=False, weight=Float8Tensor(self.act_quant_kwargs=QuantizeTensorToFloat8Kwargs(float8_dtype=torch.float8_e4m3fn, granularity=PerRow(dim=-1), mm_config=None, hp_value_lb=None, hp_value_ub=None, kernel_preference=<KernelPreference.AUTO: 'auto'>), self.block_size=[1, 2048], self.mm_config=Float8MMConfig(emulate=False, use_fast_accum=True, pad_inner_dim=False), self.scale.shape=torch.Size([512, 1]), self.kernel_preference=<KernelPreference.AUTO: 'auto'>))
            (o_proj): Linear(in_features=4096, out_features=2048, bias=False)
            (q_norm): Qwen3VLMoeTextRMSNorm((128,), eps=1e-06)
            (k_norm): Qwen3VLMoeTextRMSNorm((128,), eps=1e-06)
          )
          (mlp): Qwen3VLMoeTextSparseMoeBlock(
            (gate): Linear(in_features=2048, out_features=128, bias=False)
            (experts): Qwen3VLMoeTextExperts(
              gate_up_proj=Float8Tensor(self.act_quant_kwargs=QuantizeTensorToFloat8Kwargs(float8_dtype=torch.float8_e4m3fn, granularity=PerRow(dim=-1), mm_config=None, hp_value_lb=None, hp_value_ub=None, kernel_preference=<KernelPreference.AUTO: 'auto'>), self.block_size=[1, 2048, 1], self.mm_config=Float8MMConfig(emulate=False, use_fast_accum=True, pad_inner_dim=False), self.scale.shape=torch.Size([128, 1, 1536]), self.kernel_preference=<KernelPreference.AUTO: 'auto'>), down_proj=Float8Tensor(self.act_quant_kwargs=QuantizeTensorToFloat8Kwargs(float8_dtype=torch.float8_e4m3fn, granularity=PerRow(dim=-1), mm_config=None, hp_value_lb=None, hp_value_ub=None, kernel_preference=<KernelPreference.AUTO: 'auto'>), self.block_size=[1, 768, 1], self.mm_config=Float8MMConfig(emulate=False, use_fast_accum=True, pad_inner_dim=False), self.scale.shape=torch.Size([128, 1, 2048]), self.kernel_preference=<KernelPreference.AUTO: 'auto'>)
              (act_fn): SiLUActivation()
            )
          )
          (input_layernorm): Qwen3VLMoeTextRMSNorm((2048,), eps=1e-06)
          (post_attention_layernorm): Qwen3VLMoeTextRMSNorm((2048,), eps=1e-06)
        )
      )
      (norm): Qwen3VLMoeTextRMSNorm((2048,), eps=1e-06)
      (rotary_emb): Qwen3VLMoeTextRotaryEmbedding()
    )
  )
  (lm_head): Linear(in_features=2048, out_features=151936, bias=False)

when we run the following code:

import logging

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig, Qwen3VLMoeForConditionalGeneration, AutoProcessor

from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    PerRow,
    FqnToConfig,
    quantize_
)

# Configure logging to see warnings and debug information
logging.basicConfig(
    level=logging.INFO, format="%(name)s - %(levelname)s - %(message)s"
)

# Enable specific loggers that might contain the serialization warnings
logging.getLogger("transformers").setLevel(logging.INFO)
logging.getLogger("torchao").setLevel(logging.INFO)
logging.getLogger("safetensors").setLevel(logging.INFO)
logging.getLogger("huggingface_hub").setLevel(logging.INFO)

config = Float8DynamicActivationFloat8WeightConfig(
    granularity=PerRow(),
)
expert_config = Float8DynamicActivationFloat8WeightConfig(
    granularity=[PerRow(), PerRow(1)],
)


# only quantize language model
quant_config = FqnToConfig({
    r"re:model.language_model.*.gate_up_proj": expert_config,
    r"re:model.language_model.*.down_proj": expert_config,
    r"re:model.language_model.*.q_proj": config,
    r"re:model.language_model.*.k_proj": config,
    r"re:model.language_model.*.v_proj": config,
})
quantization_config = TorchAoConfig(quant_type=quant_config)
model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
    "Qwen/Qwen3-VL-30B-A3B-Instruct",
    dtype="auto",
    device_map="auto",
    quantization_config=quantization_config
)

print(model)

Test Plan:

pytest test/quantization/test_quant_api.py -k test_fqn_to_config_repr

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3328

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 6 Pending

As of commit e1c23ba with merge base b4ec4cb ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

is there a test we can add print(module) to to prevent regressions?

[jerryzh168]

might be a bit fragile? since small changes will break it, maybe use FileCheck()?

[jerryzh168]

same here

[jerryzh168]

looks good, had some comments inline