14fromlabmlimportexperiment15fromlabml.configsimportoption,calculate16fromlabml\_nn.transformersimportTransformerConfigs17fromlabml\_nn.transformers.basic.autoregressive\_experimentimportAutoregressiveTransformer,Configs

21def\_rotary\_pe\_mha(c:TransformerConfigs):22fromlabml\_nn.transformers.ropeimportRotaryPEMultiHeadAttention23returnRotaryPEMultiHeadAttention(c.n\_heads,c.d\_model,1.)

27calculate(TransformerConfigs.encoder\_attn,'rotary',\_rotary\_pe\_mha)28calculate(TransformerConfigs.decoder\_attn,'rotary',\_rotary\_pe\_mha)29calculate(TransformerConfigs.decoder\_mem\_attn,'rotary',\_rotary\_pe\_mha)

32@option(Configs.model,'rotary\_pe\_transformer')33def\_model(c:Configs):

37m=AutoregressiveTransformer(c.transformer.encoder,38c.transformer.src\_embed,39c.transformer.generator).to(c.device)4041returnm

44defmain():

46experiment.create(name="rotary\_pe\_transformer",writers={'screen'})

48conf=Configs()

50experiment.configs(conf,{

52'transformer.src\_embed':'no\_pos',53'transformer.tgt\_embed':'no\_pos',

56'transformer.encoder\_attn':'rotary',

59'model':'rotary\_pe\_transformer',

62'tokenizer':'character',

64'prompt\_separator':'',

66'prompt':'It is ',

68'text':'tiny\_shakespeare',

71'seq\_len':512,

73'epochs':32,

75'batch\_size':4,

78'inner\_iterations':10,

81'd\_model':128,82'transformer.ffn.d\_ff':512,83'transformer.n\_heads':16,84'transformer.dropout':0.0,

87'optimizer.optimizer':'Noam',88'optimizer.learning\_rate':1.,8990'dataloader\_shuffle\_with\_replacement':True91})

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

97withexperiment.start():

99conf.run()

103if\_\_name\_\_=='\_\_main\_\_':104main()