15fromlabmlimportexperiment16fromlabml.configsimportoption17fromlabml\_nn.transformersimportTransformerConfigs18fromlabml\_nn.transformers.basic.autoregressive\_experimentimportConfigs19fromlabml\_nn.transformers.configsimportFeedForwardConfigs20fromlabml\_nn.transformers.primer\_ezimportSquaredReLU

23@option(FeedForwardConfigs.activation,'SquaredReLU')24def\_squared\_relu():

30returnSquaredReLU()

33@option(TransformerConfigs.encoder\_attn,'MultiDConvHeadAttention')34def\_d\_conv\_mha(c:TransformerConfigs):

40fromlabml\_nn.transformers.primer\_ezimportMultiDConvHeadAttention41returnMultiDConvHeadAttention(c.n\_heads,c.d\_model,dropout\_prob=c.dropout)

44@option(TransformerConfigs.encoder\_attn,'MultiDSharedConvHeadAttention')45def\_d\_shared\_conv\_mha(c:TransformerConfigs):

53fromlabml\_nn.transformers.primer\_ez.variationsimportMultiDSharedConvHeadAttention54returnMultiDSharedConvHeadAttention(c.n\_heads,c.d\_model,dropout\_prob=c.dropout)

57@option(TransformerConfigs.encoder\_attn,'MultiDPHConvHeadAttention')58def\_d\_per\_head\_conv\_mha(c:TransformerConfigs):

66fromlabml\_nn.transformers.primer\_ez.variationsimportMultiDPHConvHeadAttention67returnMultiDPHConvHeadAttention(c.n\_heads,c.d\_model,dropout\_prob=c.dropout)

70defmain():

72experiment.create(name="primer\_ez")

74conf=Configs()

76experiment.configs(conf,{

78'tokenizer':'character',

80'prompt\_separator':'',

82'prompt':'It is ',

84'text':'tiny\_shakespeare',

87'seq\_len':256,

89'epochs':128,

91'batch\_size':32,

94'inner\_iterations':10,

97'd\_model':512,98'transformer.ffn.d\_ff':2048,

101'optimizer.optimizer':'Adam',102'optimizer.learning\_rate':2.5e-4,

107'transformer.ffn.activation':'SquaredReLU',

112'transformer.encoder\_attn':'MultiDConvHeadAttention',113})

116experiment.add\_pytorch\_models({'model':conf.model})

119withexperiment.start():

121conf.run()

125if\_\_name\_\_=='\_\_main\_\_':126main()