Loss growing back after 3-5 epochs

[XinnuoXu]

Hi!
I'm trying to move my previously working model (based on OpenNMT) to pytorch-lightning, which is basically a Bart fine-tuning model.

Here are some key callbacks:

def training_step(self, batch, batch_idx):
        src = batch[0]
        tgt = batch[1].contiguous()
        mask_src = batch[2]
        mask_tgt = batch[3]
        labels = tgt[:, 1:].clone()
        labels[tgt[:, 1:] == self.pad_id] = -100
        loss, logits = self.model(src, tgt[:, :-1], mask_src, mask_tgt[:, :-1], labels)
        self.log('train_loss', loss)
        return loss

def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(),
                                    lr=self.args.lr,
                                    betas=(self.args.beta1, self.args.beta2),
                                    eps=self.args.adam_eps,
                                    weight_decay=self.args.weight_decay)
        scheduler = {'scheduler': NoamLR(optimizer, self.args.warmup_steps),
                     'monitor': 'metric_to_track',
                     'interval': 'step',
                     'frequency': 1,
                     'strict': True,
                    }
        return [optimizer], [scheduler]

Here is the initialization of pytorch-lightning trainer:

    train_obj = LightningObject(args, device)
    trainer = pl.Trainer(gpus=args.gpu_ranks,
                                num_nodes=args.num_nodes,
                                accelerator=args.lightning_accelerator,
                                max_epochs=args.train_epochs,
                                val_check_interval=args.val_check_interval,
                                accumulate_grad_batches=args.accum_count,
                                callbacks=[lr_monitor, checkpoint_callback],
                                log_every_n_steps=args.log_every_n_steps)
        trainer.fit(train_obj, train_loader)

and here are some of my hyper-parameters:

python train.py  \
        -mode train \
        -data_path ${DATA_PATH} \
        -model_path ${MODEL_PATH} \
        -lr 1e-04 \
        -pad_id 1 \
        -max_pos 800 \
        -adam_eps 1e-08 \
        -weight_decay 0.01 \
        -accum_count 5 \
        -batch_size 3 \
        -warmup_steps 1000 \
        -train_epochs 35 \
        -val_check_interval 660 \
        -log_every_n_steps 50 \
        -visible_gpus 2 \

I was training with one GPU on one node.

The loss on toy training dataset (1000 examples) went down from 5 to 1.7 -ish in the first couple of epochs and grow back to 2.5 later.

I also tried without using 'noam' schedular, but with the same learning rate, the loss followed 9+ -> 2+ -> 6+.

I also tried to decrease the learning rate but the similar situation happened going alone with longer time (more epochs).

Could you please point me to some potential directions?

Cheers,
Xinnuo