[BUG] LayerNorm without affine params causes errors

[ishan-modi]

What does this PR do?

Error occurs when we check m.weights.is_meta here for module of type nn.LayerNorm with affine params as False, because it initializes weights as None see here

This change checks for weights being None and bypasses the module if it holds true

[meenchen]

Hi @ishan-modi, could you provide a code snippet to produce the error? I want to add some tests to cover the use case.

[ishan-modi]

Thanks for response @meenchen, please find the code snippet

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 = "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 = SanaTransformer2DModel.from_pretrained(checkpoint)

diffusers == 0.33.1
modelopt == 0.31.0

[ishan-modi]

@meenchen, any updates on this ?

[ishan-modi]

@meenchen @kevalmorabia97, this is little critical for enabling modelopt in huggingface diffusers. Can we release a patch for this because this bug might also be occurring in generative models using layernorm without affine params ?

[kevalmorabia97]

Hi @ishan-modi sorry we couldnt get this in the previous release. We'll include this fix in the upcoming release 0.35.0 planned in about 2 weeks. Would that timeline be fine?

[ishan-modi]

Thanks @kevalmorabia97! Patch release would be great for faster integration into diffusers, though next release is fine too.

[kevalmorabia97]

Fix is merged in main branch but not in the official release. Temporary workaround would be to install modelopt from main branch instead of wheels pip install git+https://github.com/NVIDIA/TensorRT-Model-Optimizer.git

[kevalmorabia97]

@ishan-modi 0.33.1 release is now live which includes this fix

[ishan-modi]

Thanks @kevalmorabia97