question: Scheduler.step() only called on the end of validation
context: What I did:
  def configure_optimizers(self):
  
      optimizer = torch.optim.AdamW(
          self.parameters(),
          lr=self.hparams.learning_rate,
          eps=1e-5,
      )
      scheduler = torch.optim.lr_scheduler.OneCycleLR(
          optimizer,
          max_lr=self.hparams.learning_rate,
          epochs=self.trainer.max_epochs,
          steps_per_epoch=len(self.datamodule.train_dataloader()),
      )
      set_trace()
      return {"optimizer": optimizer, "lr_scheduler": scheduler}
  I put an breakpoint on scheduler.step to see when it will be called.
  What I got:
  
  And I found
  
  I ran into breakpoint (scheduler.step is called) only on the end of validation.
  epochs=self.trainer.max_epochs=1 and steps_per_epoch=len(self.datamodule.train_dataloader())=144, which are correct.
  
  Code for your reference if needed
  import random
  from IPython.core.debugger import set_trace
  import torch
  import torch.nn.functional as F
  from transformers import AutoConfig, ElectraForTokenClassification
  import pytorch_lightning as pl
  from pytorch_lightning.callbacks.lr_monitor import LearningRateMonitor
  from pl_bolts.callbacks import PrintTableMetricsCallback
  from pytorch_lightning.loggers import WandbLogger
  
  class SortNumberDataset(torch.utils.data.Dataset):
      def __init__(self, dataset_size, vocab_size, sequence_length):
          super().__init__()
          self.dataset_size = dataset_size
          self.vocab_size = vocab_size
          self.sequence_length = sequence_length
  
      def __getitem__(self, i):
          x = [
              random.randint(0, self.vocab_size - 1) for _ in range(self.sequence_length)
          ]
          y = sorted(x)
          return {"x": torch.tensor(x), "y": torch.tensor(y)}
  
      def __len__(self):
          return self.dataset_size
  
  
  class NumberSorting(pl.LightningModule):
      def __init__(
          self,
          hf_config,
          learning_rate,
          trainset_size,
          valset_size,
          vocab_size,
          sequence_length,
          batch_size=128,
          num_workers=4,
      ):
          super().__init__()
          self.save_hyperparameters()
          self.datamodule = pl.LightningDataModule.from_datasets(
              SortNumberDataset(trainset_size, vocab_size, sequence_length),
              SortNumberDataset(valset_size, vocab_size, sequence_length),
              batch_size=batch_size,
              num_workers=num_workers,
          )
          # self.model = ElectraForTokenClassification(hf_config)
           # tie input/output embeddings
           delattr(self.model, "classifier")
           self.model.classifier = (
               lambda x: x @ self.model.electra.embeddings.word_embeddings.weight.t()
          )
          self.val_acc = pl.metrics.Accuracy()
  
      def forward(self, batch):
          result = self.model(input_ids=batch["x"], labels=batch["y"], return_dict=True)
          return result.logits.argmax(dim=-1), result.loss
  
      def training_step(self, batch, batch_idx):
          # self.log("lr", self.trainer.optimizers[0].param_groups[0]["lr"])
          return self(batch)[-1]
  
      def validation_step(self, batch, batch_idx):
          preds, loss = self(batch)
          self.val_acc(preds.view(-1), batch["y"].view(-1))
          self.log("valid_acc", self.val_acc, on_epoch=True)
  
      def configure_optimizers(self):
          timizer = torch.optim.AdamW(optimizer_grouped_parameters, eps=1e-5)
          optimizer = torch.optim.AdamW(
              self.parameters(),
              lr=self.hparams.learning_rate,
          )
          # return optimizer
          scheduler = torch.optim.lr_scheduler.OneCycleLR(
              optimizer,
              max_lr=self.hparams.learning_rate,
              epochs=self.trainer.max_epochs,
              steps_per_epoch=len(self.datamodule.train_dataloader()),
          )
          set_trace()
          return {"optimizer": optimizer, "lr_scheduler": scheduler}
  
  
  config = AutoConfig.from_pretrained(
      "google/electra-small-generator",
      pad_token_id=-1,
      max_position_embeddings=7,
      vocab_size=7,
      num_labels=7,
      embedding_size=10,
      hidden_size=10,
      intermediate_size=8,
      num_hidden_layers=2,
      num_attention_heads=2,
  )
  plmodule = NumberSorting(
      config,
      learning_rate=0.05,
      trainset_size=18333,
      valset_size=20000,
      vocab_size=7,
      sequence_length=7,
  )
  
  
  trainer = pl.Trainer(
      max_epochs=1,
      gpus="0",
      callbacks=[
          PrintTableMetricsCallback(),
          LearningRateMonitor(logging_interval="step", log_momentum=True),
      ],
      logger=WandbLogger(),
  )
  trainer.fit(plmodule)
answer: That is the default behaviour of learning rate schedulers, that they step at the end of the training epoch.
  Can I ask you what you are trying to achieve?
  If you want the learning rate scheduler to step after each batch, you can read more about what the output of configure_optimizers should look like here: https://pytorch-lightning.readthedocs.io/en/0.9.0/optimizers.html#learning-rate-scheduling