Quantized Model Output NaN / 0

[sophia1488]

❓ Questions

Hi, I want to apply DiffQ to my source separation model with PyTorch Lightning framework, and I added the quantizer following the Usage on ReadMe.

In my callback function during training, it works fine when evaluating the unquantized model with SDR.
But when I use the quantized model to run separation on MUSDB test set or other pop songs, the separation result is NaN or contains lots of 0s.

Do you have any comments or suggestions on this? Hope to get your reply!
Sincerely,

[adefossez]

During training, do you get any issue with the valid loss ? Do you log the model size from diffq ? How are you dumping and loading the model for evaluation ?

[sophia1488]

For MUSDB test set, its accompaniment SDR can reach 14.5 at step 20k (w/o adding quantized loss, it's normally 15.x), so I think it's acceptable. Later it degrades to 11 (even 8), but I saved the checkpoint with the best SDR.

Below is the minimal example of callback function during training:

import pytorch_lightning as pl
class EvaluationCallback(pl.Callback):
    def __init__(self, model, quantizer, ...):
        self.model = model
        self.quantizer = quantizer

    def on_batch_end():                                              
        # Evaluate on musdb test set
        for each mixture in musdb_testset:
            output = self.model.separate(mixture)
            sdr, _, _, _ = museval.evaluate(target, output)
        if median_sdr > best_sdr:
             # save checkpoints if it's the best SDR so far
             torch.save(self.model.state_dict(), "best_model.pth")
             torch.save(self.quantizer.get_quantized_state(), "best_quantized.pth")      # the quantized model

To load the model

model = MyModel()
diffq.restore_quantized_state(model, torch.load("best_quantized.pth"))
output = model.separate(mixture)

I didn't log the quantized model size at each step, but the model w/o using DiffQ is 198MB, whereas the quantized model (best_quantized.pth) is 22 MB and the best_model.pth is 222MB.

Thanks!

[adefossez]

you are probably getting a divergence at some point. you can try reducing the penalty level. it is a bit hard to debug without logs or more code that was used for training.

[sophia1488]

Hi @adefossez, sorry for the late reply. Below is the code used for training.
I set the penalty smaller (2e-5).
At epoch 1 (step 10000), True model size: 39.76 MB, Compressed model size: 28.84 MB,
and the SDR for MUSDB valid set at validation step is pretty good (1x).
The training loss & valid loss also looks fine.

class LitModel(LightningModule):
   def on_fit_start(self):
        opt = self.trainer.optimizers[0]  # this is a list, opt is Adam initialized in configure_optimizers()
        self.quantizer = DiffQuantizer(self.model, group_size=8)       # initialize the quantizer here
        self.quantizer.setup_optimizer(opt)

    def training_step(...):
        output_dict = self.model(input_dict)
        loss = l1_wav(...)
        loss += self.penalty * self.quantizer.model_size()   

    def validation_step(...):
        loss = l1_wav(...)
        sdr, _, _, _ = museval.evaluate([target[b].T], [outputs[b].T])
        ...

    def configure_optimizers(self):
        opt = torch.optim.Adam(self.model.parameters(), lr)
        ...

    def on_save_checkpoint(self, checkpoint):
        checkpoint['state_dict'] = quantizer.get_quantized_state()

When calculating SDR for validation, self.model is used, but the checkpoint saves quantizer.get_quantized_state().
I wonder if that'd make a difference.

Also, in Readme, there's a warning: you must always wrap your model with DistributedDataParallel after having created the quantizer. If I only use 1 gpu, should I also wrap the model?

model = DistributedDataParallel(
            model,
            device_ids=[torch.cuda.current_device()],
            output_device=torch.cuda.current_device())

Thank you and hope to get your feedback!

[adefossez]

Thanks for the extra details. You don't need to wrap with Distributed if you are using a single GPU.
All look good in what you are doing. If you load the checkpoint you dumped with best_model.pth, do you have the same result ? (you will probably need to filter the diffq params out: state_dict = {key: value for key, value in state_dict.items() if not key.endswith('_diffq')} ? Note that this will use the unquantized weights though, but this is for debugging.
Finally, you can try to dump the quantized weights with

with self.model.enter_quantize():
    torch.save(self.model.state_dict(), "best_model_quantized.pth")

and load them using the same snippet. This will contain the quantized weights values but stored as float 32. This is again for debugging. Finally, the get_quantized_state() returns the weights quantized and optimally bit packed. I'm trying to see at what stage do the NaN appear.

[sophia1488]

Hi @adefossez , thanks for your reply! I refactor my code and re-trained the model, so it took a while...
Now there's no NaN or 0s, but the performance of the quantized model drops a lot.

I've tried your suggestions:

1. Load the checkpoint and filter out diffq params. -> The result is fine!

# Load checkpoint
checkpoint = torch.load(checkpoint_path, map_location='cpu')
state_dict = {key: value for key, value in checkpoint['model'].items() if not key.endswith('_diffq')}
model.load_state_dict(state_dict)

# Result on test set (epoch 8)
(Valid) median accompaniment SDR: 15.168
 --- Median vocals SDR: 6.74 ---
 --- Median accompaniment SDR: 12.49 ---

2. Dump the quantized weights -> The result is fine!

# Dump
with self.quantizer.enter_quantize():
  ckpt = {'model': self.model.state_dict(), ...}
  torch.save(ckpt, "best_model_quantized.pth")
# Load checkpoint as 1.

# Result on test set (epoch 12)
Valid SDR: 15.819
--- Median vocals SDR: 7.75 ---
--- Median accompaniment SDR: 13.65 ---

3. Use quantized state -> not working though...

# Dump
ckpt = {'quantized': self.quantizer.get_quantized_state(), ...}
torch.save(ckpt, "best_model_quantized.pth")
# Load
checkpoint = torch.load(checkpoint_path, map_location='cpu')
diffq.restore_quantized_state(model, checkpoint['quantized'])

# Result on test set (epoch 8)
(Valid) median accompaniment SDR: 15.168
--- Median vocals SDR: -5.49 ---
--- Median accompaniment SDR: 0.13 ---

Do you have any ideas on this? Thank you a lot 🙏

[sophia1488]

I also applied LSQ on the pre-trained model, and the result is similar. 😥

Valid SDR: 15.043
--- Median vocals SDR: -6.57 ---
--- Median accompaniment SDR: -0.41 ---

[adefossez]

For LSQ what number of bits are you using ? Are you not just compressing too much here ?
Sorry for the very late reply...