[BUG] modelopt restore using 'from_pretrained' doesn't work in diffusers

[ishan-modi]

Describe the bug

Restore functionality doesnt work for from_pretrained function of huggingface/diffusers

Steps/Code to reproduce bug

from diffusers import SanaTransformer2DModel
import modelopt.torch.quantization as mtq

checkpoint = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
model = SanaTransformer2DModel.from_pretrained(checkpoint, subfolder="transformer")

config = mtq.FP8_DEFAULT_CFG
config['quant_cfg']['*weight_quantizer'].update({'fake_quant': False})
model = mtq.quantize(model, config)
mtq.compress(model)
model.save_pretrained("test_modelopt_quant")

model = SanaTransformer2DModel.from_pretrained("test_modelopt_quant")

Error Trace

    155 if force_quantize:
    156     m.weight_quantizer._dequantize = False
--> 157 assert not m.weight.is_meta, (
    158     "Real quantization does not support tensors on meta device."
    159 )
    160 real_quant_tensor = m.weight_quantizer(m.weight)
    161 m.weight = QTensorWrapper(real_quant_tensor)

AssertionError: Real quantization does not support tensors on meta device.

It seems like weights are on meta device for diffusers unlike transformers where they are on cpu when we restore modelopt_state

[kevalmorabia97]

Is this transformers>4.51? I think with latest transformers they do model restore on meta device. We have a fix for it coming up in the 0.29.0 release (early May). Meanwhile if possible can you try with transformers<=4.51?

[ishan-modi]

Got it, actually I am using latest version of diffusers, but this makes sense hope these changes will fix the diffuser as well

[kevalmorabia97]

Can you please retry with nvidia-modelopt==0.29.0 to see if this works now?

[ishan-modi]

Works now, Thanks !!

[ishan-modi]

@kevalmorabia97, on a closer inspection I found that error is still there.

from transformers import AutoModelForCausalLM
# from diffusers import SanaTransformer2DModel
import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.config import FP8_DEFAULT_CFG
from modelopt.torch.opt import enable_huggingface_checkpointing

enable_huggingface_checkpointing()

checkpoint = "HuggingFaceTB/SmolLM2-135M"
model = AutoModelForCausalLM.from_pretrained(checkpoint)

# checkpoint = "Efficient-Large-Model/Sana_600M_1024px_diffusers"
# model = SanaTransformer2DModel.from_pretrained(checkpoint, subfolder="transformer")

FP8_DEFAULT_CFG['quant_cfg']['*weight_quantizer'].update({'fake_quant': False})
FP8_DEFAULT_CFG['quant_cfg']['*input_quantizer'].update({'enable': False})
mtq.quantize(model, FP8_DEFAULT_CFG)
mtq.compress(model)
model.save_pretrained('test_modelopt_quant')


checkpoint = "./test_modelopt_quant"
model = AutoModelForCausalLM.from_pretrained(checkpoint)
# model = SanaTransformer2DModel.from_pretrained(checkpoint)

Library Versions

modelopt == 0.29.0
transformers == 4.51.1
diffusers == 0.33.1