Resume from checkpoint with quantization

[rotx-maxim]

I am using a workaround to allow resuming from checkpoint with active quantization. The requires_grad flags aren't set in the restored biases and weights (they seem to be present at checkpoint save time). So as a quick fix I use:

    def set_grad(m):
        """
        Force the `requires_grad` flag on all weights and biases
        """
        if isinstance(m, (nn.Linear, nn.Conv2d)):
            m.weight.requires_grad_()
            if hasattr(m, 'bias') and m.bias is not None:
                m.bias.requires_grad_()

    model.apply(set_grad)

Without this, I get the PyTorch error message element 0 of tensors does not require grad and does not have a grad_fn.

Looking for a proper way to fix this.

[guyjacob]

Unfortunately, resuming training with active quantization isn't really supported. The mechanism that's in there at the moment, with quantizer_metadata as you've noticed, is designed just for the case where you want to load a quantized model and perform evaluation (in this context, the command-line argument name resume is a bit misleading).
As it is now, what happens is that the model is quantized when the checkpoint is loaded. But then when the schedule YAML is loaded, it "quantizes" the already quantized module, the result of which is just wrong (it zeroes out the float_weight parameter).
So even if your workaround from above lets you bypass the error you're getting, I highly doubt you're getting correct results.

At this time we weren't planning to focus on this feature. If you want to have a go at it, I'll be happy to help with any questions.

[rotx-maxim]

@guyjacob, thanks for your reply. This explains why the hack seems to work for resuming quantized weight training -- it's the sequencing. My demo is different in this respect from the sample app in this repo. It does the following in the following order (it seemed to make sense that the model shouldn't be prepared twice, and not calling prepare_model at all causes the restore to miss the float_xxx).

1. Create and load compression scheduler (if present in checkpoint)

        if 'compression_sched' in checkpoint:
            compression_scheduler = distiller.CompressionScheduler(model, device=device)
            compression_scheduler.load_state_dict(checkpoint['compression_sched'])
   

2. Load quantizer metadata (if present in checkpoint), and prepare model

        if 'quantizer_metadata' in checkpoint:
            qmd = checkpoint['quantizer_metadata']
            quantizer = qmd['type'](model, **qmd['params'])
            quantizer.prepare_model()
   

3. Load model state dictionary:

       model.load_state_dict(.......)
   

4. Fix up requires_grad flags per comment above:

       model.apply(set_grad)
   

5. Load optimizer from checkpoint:

       if 'optimizer' in checkpoint:
           optimizer = Optimizer(model.parameters(), **optimizer_args)
           optimizer.load_state_dict(checkpoint['optimizer'])
   

6. Load YAML and create compression scheduler if we didn't restore one earlier and if we want to compress:

       if args.compress and not compression_scheduler:
           compression_scheduler = distiller.config.file_config(model, optimizer, args.compress, device)
   

I'm not clear why load_state_dict ignores the requires_grad flags but I haven't tried digging too deep.