How to load checkpoint for a sharded model?

[ver217]

I tried to load checkpoint for a sharded model, but I got unexpected keys error. It seems that the sharded model is not initialized when loading checkpoint.

Here is my code:

from data import DaliImagenetDataModule
from loss import MulticlassBCEWithLogitsLoss, MixupLoss, LabelSmoothLoss
from timm.models.vision_transformer import vit_base_patch16_224, vit_small_patch16_224
from pytorch_lightning import seed_everything
from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
from colossalai.nn.optimizer import HybridAdam
from torchmetrics import Accuracy
import pytorch_lightning as pl
import torch
import os
from argparse import ArgumentParser


class ViTModule(pl.LightningModule):
    def __init__(self, ) -> None:
        super().__init__()
        self.criterion = MixupLoss(MulticlassBCEWithLogitsLoss(0.1))
        self.acc = Accuracy()

    def configure_sharded_model(self) -> None:
        self.model = vit_small_patch16_224(drop_rate=0.1, weight_init='jax', num_classes=100)

    def configure_optimizers(self):
        opt = HybridAdam(self.model.parameters(), lr=1.8e-2, weight_decay=0.1)
        scheduler = CosineAnnealingWarmupLR(opt, 300, 32)
        scheduler = {'scheduler': scheduler, 'interval': 'epoch'}
        return [opt], [scheduler]

    def training_step(self, batch, batch_idx):
        inputs, targets = batch
        logits = self.model(inputs)
        loss = self.criterion(logits, **targets)
        return loss

    def validation_step(self, batch, batch_idx):
        inputs, targets = batch
        if targets['targets_a'].ndim == 0:
            targets['targets_a'] = targets['targets_a'].unsqueeze(0)
        if targets['targets_b'].ndim == 0:
            targets['targets_b'] = targets['targets_b'].unsqueeze(0)
        logits = self.model(inputs)
        loss = self.criterion(logits, **targets)
        preds = torch.argmax(logits, 1)
        self.acc(preds, targets['targets_a'])
        self.log_dict({'valid_loss': loss, 'valid_acc': self.acc}, prog_bar=True, on_epoch=True, sync_dist=True)


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--np', type=int, default=1)
    parser.add_argument('--batch_size', type=float, default=128)
    parser.add_argument('--ckpt', default=None)
    args = parser.parse_args()
    assert args.batch_size % args.np == 0
    local_batch_size = args.batch_size // args.np
    seed_everything(42)
    dm = DaliImagenetDataModule(os.environ['DATA'], local_batch_size,
                                mixup_alpha=0.2, randaug_magnitude=10, randaug_num_layers=2,
                                gpu_aug=True)
    model = ViTModule()
    trainer_cfg = {
        'accelerator': 'cuda',
        'strategy': 'ddp',
    }
    trainer = pl.Trainer(
        max_epochs=300,
        devices=args.np,
        gradient_clip_val=1.0,
        # resume_from_checkpoint=args.ckpt,
        **trainer_cfg
    )
    trainer.fit(model, datamodule=dm, ckpt_path=args.ckpt)

[ver217]

[rohitgr7]

trainer_cfg = {
        'accelerator': 'cuda',
        'strategy': 'ddp',
    }

is it really sharded?

[ver217]

trainer_cfg = {
        'accelerator': 'cuda',
        'strategy': 'ddp',
    }

is it really sharded?

No, I just want to test configure_sharded_model(). I found a solution:

class ViTModule:
    ...
    def on_load_checkpoint(self, checkpoint) -> None:
        if not hasattr(self, 'model'):
            self.configure_sharded_model()

Is this right?

[rohitgr7]

configure_sharded_model

this should only be used with sharded strategies.