I have looked into how Fairseq does the warmup https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md 8
I am having a tricky time following the code, but it looks like the warmup is handled with PolynomialDecaySchedule (since there is a ‘@register_lr_scheduler(“polynomial_decay”)’, and ‘polynomial_decay’ is an arg for the pretraining command) , by passing an optimizer object, and specific args
I can deduce most of the args from the ‘Train Roberta Base’ command here
https://github.com/pytorch/fairseq/blob/master/examples/roberta/README.pretraining.md#2-train-roberta-base 1
Except for power, which I am guessing will just be left to the default of 1.
I am also wondering the decay is ‘and then linearly decayed’ as mentioned in the Roberta paper. I see that there is already a decay happening with AdamW, since its decay rate is set at 0.01, but I don’t think that’s a linear decrease, and the paper does mention that there is a linear decrease.
Also, the PolynomialDecaySchedule class mentions that its supposed to be for decaying the learning rate (“Decay the LR on a fixed schedule.”). So it sounds like the class is also handling this linear decay, based on the end_learning_rate arg.
It seems to be that no arg is passed in for that, so it uses the default of zero. I’m not so sure of this though, I feel like it may have been mentioned in the paper, and its a little tricky to see what exactly is being passed into PolynomialDecaySchedule , because it doesn’t seem like that class is being instantiated in the code. It seems to be instantiated a different way which I am not able to figure out.
Maybe the end learning rate is zero after all, since that’s how it’s done in Bert, which Roberta is based on.

  
      github.com
  
  
    google-research/bert/blob/master/optimization.py#L36

"""Creates an optimizer training op."""global_step = tf.train.get_or_create_global_step()learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32)# Implements linear decay of the learning rate.learning_rate = tf.train.polynomial_decay(    learning_rate,    global_step,    num_train_steps,    end_learning_rate=0.0,    power=1.0,    cycle=False)# Implements linear warmup. I.e., if global_step < num_warmup_steps, the# learning rate will be `global_step/num_warmup_steps * init_lr`.if num_warmup_steps:  global_steps_int = tf.cast(global_step, tf.int32)  warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32)  global_steps_float = tf.cast(global_steps_int, tf.float32)



  
  
    
    
  
  


Which also shows power =1.