Your current environment

N/A

Model Input Dumps

No response

🐛 Describe the bug

When we use cpu offloading together with torch.compile, it will error:

torch._dynamo.exc.Unsupported: builtin: setattr [<class 'torch._dynamo.variables.dicts.ConstDictVariable'>, <class 'torch._dynamo.variables.constant.ConstantVariable'>, <class 'torch._dynamo.variables.dicts.ConstDictVariable'>] False

The error is caused by this line:

Creating a state dict during forward will error.

I tried another approach of using tensor subclasses in #10609 . It works well for unquantized models, but does not work for quantized models.

The problem with quantized models, is that we have some classes inherits torch.nn.Parameter, e.g.

Using both tensor subclasses and parameter subclasses is a known problem in pytorch. See https://github.com/albanD/subclass_zoo/blob/main/custom_parameter.py for example.

To make torch.compile compatible with cpu offloading and quantization, we need to refactor the weight loading logic and how we create/store weights.

Take the GPTQ linear layer for example:

        qweight = PackedvLLMParameter(
            data=torch.empty(
                input_size_per_partition // self.quant_config.pack_factor,
                output_size_per_partition,
                dtype=torch.int32,
            ),
            input_dim=0,
            output_dim=1,
            packed_dim=0,
            packed_factor=self.quant_config.pack_factor,
            weight_loader=weight_loader)

We should avoid using nn.Parameter, and directly register the tensor as buffer:

self.qweight = torch.empty(
                input_size_per_partition // self.quant_config.pack_factor,
                output_size_per_partition,
                dtype=torch.int32,
            )
self.qweight.weight_loader = weight_loader
self.qweight_packed_factor = self.quant_config.pack_factor

The key ideas are:

1. No nn.Parameter, no class inheritance. We directly assign a tensor to the module. The tensor should be registered as a buffer.
2. Every tensor can have a weight_loader attribute. We can bind arguments needed for weight loading, e.g. self.qweight.weight_loader = partial(generic_weight_loader, args)
3. For the information we need during runtime (i.e. computing the forward pass), we store them in the module object, e.g. self.qweight_packed_factor = self.quant_config.pack_factor.

With all these changes, we should be able to use cpu offloading with quantization and torch.compile .

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.