model.from_pretrained( . . . ) not loading needed weights/parameters
Author: lorsonblair
Pull request: no

I am performing quantization of a PatchTSTForPrediction model and attempting to load a saved quantized model for testing. Model is saved using `model.save_pretrained( . . . )`. Testing proceeds perfectly once performed immediately after QAT (Hugging face trainer's handles loading at the end of training); however, when attempting to load a saved quantized (trained) model, the error below occurs. I perform all the pre-quantization preparation so that the model contains all the necessary parameters (untrained) and then try to load the saved checkpoint. How can I force `from_pretrained( . . . )` to load ALL required weights?  

`Some weights of the model checkpoint at ./checkpoints/ . . . were not used when initializing PatchTSTForPrediction: ['head.projection.calib_counter', 'head.projection.num_module_called', 'head.projection.obsrv_clipval', 'head.projection.obsrv_clipvaln', 'head.projection.obsrv_w_clipval', 'head.projection.quantize_feature.clip_val', 'head.projection.quantize_feature.clip_valn', 'head.projection.quantize_weight.clip_val', 'model.encoder.layers.0.ff.0.calib_counter', 'model.encoder.layers.0.ff.0.num_module_called', 'model.encoder.layers.0.ff.0.obsrv_clipval', 'model.encoder.layers.0.ff.0.obsrv_clipvaln', 'model.encoder.layers.0.ff.0.obsrv_w_clipval', 'model.encoder.layers.0.ff.0.quantize_feature.clip_val', 'model.encoder.layers.0.ff.0.quantize_feature.clip_valn', 'model.encoder.layers.0.ff.0.quantize_weight.clip_val', 'model.encoder.layers.0.ff.3.calib_counter', 'model.encoder.layers.0.ff.3.num_module_called', 'model.encoder.layers.0.ff.3.obsrv_clipval', 'model.encoder.layers.0.ff.3.obsrv_clipvaln', 'model.encoder.layers.0.ff.3.obsrv_w_clipval', 'model.encoder.layers.0.ff.3.quantize_feature.clip_val', 'model.encoder.layers.0.ff.3.quantize_feature.clip_valn', 'model.encoder.layers.0.ff.3.quantize_weight.clip_val', 'model.encoder.layers.0.self_attn.QBmm52.num_module_called', 'model.encoder.layers.0.self_attn.QBmm52.quantize_m1.clip_val', 'model.encoder.layers.0.self_attn.QBmm52.quantize_m1.clip_valn', 'model.encoder.layers.0.self_attn.QBmm52.quantize_m2.clip_val', 'model.encoder.layers.0.self_attn.QBmm52.quantize_m2.clip_valn', 'model.encoder.layers.0.self_attn.QBmm62.num_module_called', 'model.encoder.layers.0.self_attn.QBmm62.quantize_m1.clip_val', 'model.encoder.layers.0.self_attn.QBmm62.quantize_m1.clip_valn', 'model.encoder.layers.0.self_attn.QBmm62.quantize_m2.clip_val', 'model.encoder.layers.0.self_attn.QBmm62.quantize_m2.clip_valn', 'model.encoder.layers.0.self_attn.k_proj.calib_counter', 'model.encoder.layers.0.self_attn.k_proj.num_module_called', 'model.encoder.layers.0.self_attn.k_proj.obsrv_clipval', 'model.encoder.layers.0.self_attn.k_proj.obsrv_clipvaln', 'model.encoder.layers.0.self_attn.k_proj.obsrv_w_clipval', 'model.encoder.layers.0.self_attn.k_proj.quantize_feature.clip_val', 'model.encoder.layers.0.self_attn.k_proj.quantize_feature.clip_valn', 'model.encoder.layers.0.self_attn.k_proj.quantize_weight.clip_val', 'model.encoder.layers.0.self_attn.out_proj.calib_counter', . . .]

This IS expected if you are initializing PatchTSTForPrediction from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
This IS NOT expected if you are initializing PatchTSTForPrediction from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).` 

NB: QAT is simulated. Additional parameters are added to the model after qmodel_prep is called and QAT proceeds as normal. I am using IBM's fms-model-optimizer.  

Comments:
cc @MekkCyber for quantization but I think we'll need complete code that reproduces the error!

Thank you for your response. 

Below is the code to reproduce the warning. You'll need to install IBM's fms-model-optimizer [https://github.com/foundation-model-stack/fms-model-optimizer] using the following: 

`git clone https://github.com/foundation-model-stack/fms-model-optimizer
cd fms-model-optimizer
pip install -e .`

This installs transformers 4.55.4. I'm currently using Pytorch 2.6 with cuda 12.4. Any version of Pytorch >= 2.3 with cuda >= 12.1 works. Hugginface-hub 0.35.3. 

To reproduce the error, run with train_first = True to allow a model to be trained via QAT, saved and evaluated. Then, set train_first = False to allow code to attempt to load the saved quantized model for inference only. Let me know if you need any more information. 

```python
import os
import math
import transformers
from transformers import (
    PatchTSTConfig,
    PatchTSTForPrediction,
    Trainer,
    TrainingArguments,
)
from fms_mo import qmodel_prep
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np

from fms_mo import (
    qconfig_init,
    qmodel_prep
)
from fms_mo.utils.utils import patch_torch_bmm

contx_len = 512
pred_len = 96
train_epochs = 20
batch_size = 32

class TimeSeriesDataset(Dataset):
    """Your dataset class"""
    def __init__(self, data, context_length, prediction_length):
        self.data = data
        self.context_length = context_length
        self.prediction_length = prediction_length

    def __len__(self):
        return len(self.data) - self.context_length - self.prediction_length

    def __getitem__(self, idx):
        past = self.data[idx:idx + self.context_length]
        future = self.data[idx + self.context_length:idx + self.context_length + self.prediction_length]
        return {
            'past_values': torch.FloatTensor(past),
            'future_values': torch.FloatTensor(future)
        }

    def create_model():
        # Configuration
        config = PatchTSTConfig(
            context_length=contx_len,
            prediction_length=pred_len,
            num_input_channels=7,
            patch_length=16,
            stride=8,
            d_model=128,
            d_ff=256,
            num_attention_heads=4,
            num_hidden_layers=3,
            loss="mse",
            positional_dropout=0.2,
            path_dropout=0.2,
            ff_dropout=0.2,
        )
        transformers.set_seed(42)

        # Create model
        model = PatchTSTForPrediction(config)

        return model

## Create datasets
train_data = np.random.randn(7000, 7).astype(np.float32)
eval_data = np.random.randn(1000, 7).astype(np.float32)
test_data = np.random.randn(2000, 7).astype(np.float32)

train_dataset = TimeSeriesDataset(train_data, contx_len, pred_len)
eval_dataset = TimeSeriesDataset(eval_data, contx_len, pred_len)
test_dataset = TimeSeriesDataset(test_data, contx_len, pred_len)

train_steps = math.ceil(len(train_dataset) / batch_size)

# Training arguments
training_args = TrainingArguments(
    output_dir="./checkpoints/output",
    overwrite_output_dir=True,
    num_train_epochs=train_epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    learning_rate=1e-4,
    do_eval=True,
    eval_strategy="epoch",
    save_strategy="epoch",
    logging_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    logging_dir="./checkpoints/logs",
    label_names=["future_values"],
)

def train_and_test(train_first):

    best_model_path = f"./checkpoints/model"
    os.makedirs(best_model_path, exist_ok=True)

    # quantization parameters
    qcfg = qconfig_init()
    qcfg["qmodel_calibration"] = 5
    qcfg["nbits_w"] = 8
    qcfg["nbits_a"] = 8

    # create model
    model = create_model().to(torch.device('cuda'))

    # Create initial trainer to allow us to get dataloader
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
    )
    train_dataloader = trainer.get_train_dataloader()

    if train_first:
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=1e-4,
        )

        calib_data = [next(iter(train_dataloader)) for _ in range(qcfg["qmodel_calibration"])]

        # quantize model for QAT
        # qmodel_prep is performed inplace, so model will now contain quantization params
        qmodel_prep(
            model,
            calib_data,
            qcfg,
            optimizer=optimizer,
            save_fname=f"qgraph.pt",
         )

        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer=optimizer,
            T_max=train_epochs * train_steps,
            eta_min=0
        )

        # create new trainer to account for qmodel and optimizer with new params
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            optimizers=[optimizer, scheduler]
        )

        # train
        trainer.train()

        # Save
        model.save_pretrained(best_model_path)

        # evaluate
        with patch_torch_bmm(qcfg):
            results = trainer.evaluate(eval_dataset=test_dataset)
        print(f"Test results: {results}")
     
    else: # perform only testing
        training_args.do_train=False
        training_args.report_to="none"

        # run 1 round of calibration so that new model contains q parameters
        # not necessary for testing, but this is just to cover all bases
        qcfg["qmodel_calibration"] = 1
        calib_data = [next(iter(train_dataloader)) for _ in range(qcfg["qmodel_calibration"])]

        qmodel_prep(
            model,
            calib_data,
            qcfg,
            save_fname=f"qgraph.pt",
        )
        # load model from path
        model.from_pretrained(best_model_path)

        # Create trainer with qmodel
        trainer = Trainer(
            model=model,
            args=training_args,
            eval_dataset=test_dataset,
        )

        # Test
        with patch_torch_bmm(qcfg):
            results = trainer.evaluate(test_dataset)

        print(f"Results: {results}")

if __name__ == "__main__":

    # Set train_first = True to allow for a QAT run and a qmodel to be saved.
    # qmodel is also tested immediately after training. Subsequently,
    # set train_first = False to allow for loading of the save model.
    # This reproduces the error.
    train_first = True

    train_and_test(train_first)`

Hi @lorsonblair, I’m not sure I fully understand how QAT works in your setup. Normally, when saving the model at the end of QAT, all the fake quantization layers should be removed, resulting in a model identical to the original, non-prepared version. That way, you can simply quantize it later when loading it from the checkpoint.

Thanks for your reply, @MekkCyber. In this instance, the quantize features are being saved, which I would expect as the model is being used for inference. A sample of qlayers added during prep for QAT is shown below. qmodel_prep performs operations inplace: 

`model.encoder.layers.0.self_attn.k_proj.quantize_feature`
`model.encoder.layers.0.self_attn.k_proj.quantize_calib_feature`
`model.encoder.layers.0.self_attn.k_proj.quantize_weight`
`model.encoder.layers.0.self_attn.k_proj.quantize_calib_weight`

This is expected behavior as I would want to load the quantized model later for inference without having to retrain (perform QAT again as the larger models take a long time to train). Is there a way to force `from_pretrained` to load all model layers (features and parameters)? It seems like `from_pretrained` is creating a new model every time it's called, so the quantized parameters get discarded.  

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.