17importtorch18fromtorchimportnn1920fromlabmlimportexperiment21fromlabml.configsimportoption22fromlabml.utils.pytorchimportget\_modules23fromlabml\_helpers.moduleimportModule24fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs

27classAutoregressiveModel(Module):

32def\_\_init\_\_(self,n\_vocab:int,d\_model:int,transformer:Module):33super().\_\_init\_\_()

35self.src\_embed=nn.Embedding(n\_vocab,d\_model)36self.transformer=transformer37self.generator=nn.Linear(d\_model,n\_vocab)

39defforward(self,x:torch.Tensor):

41x=self.src\_embed(x)

43res=self.transformer(x)

45returnself.generator(res),None

48classConfigs(NLPAutoRegressionConfigs):

55model:AutoregressiveModel5657d\_model:int=51258nu:int=159heads:int=860dropout:float=0.061d\_ff:int=204862n\_layers:int=6

65@option(Configs.model)66deffast\_weights\_transformer(c:Configs):

70fromlabml\_nn.transformers.fast\_weightsimportFastWeightsAttentionTransformer,\71FastWeightsAttentionTransformerLayer,FastWeightsAttention,FeedForward7273fromlabml\_nn.transformers.fast\_weightsimportDPFP74returnAutoregressiveModel(75c.n\_tokens,c.d\_model,76FastWeightsAttentionTransformer(77FastWeightsAttentionTransformerLayer(d\_model=c.d\_model,78attn=FastWeightsAttention(c.heads,c.d\_model,c.dropout,DPFP(nu=c.nu)),79feed\_forward=FeedForward(c.d\_model,c.d\_ff,c.dropout),80dropout\_prob=c.dropout),81c.n\_layers)).to(c.device)

84defmain():

86experiment.create(name="fast\_weights\_transformer")

88conf=Configs()

90experiment.configs(conf,

92{'tokenizer':'character',93'text':'tiny\_shakespeare',94'optimizer.learning\_rate':1.0,95'optimizer.optimizer':'Noam',96'prompt':'It is',97'prompt\_separator':'',9899'train\_loader':'shuffled\_train\_loader',100'valid\_loader':'shuffled\_valid\_loader',101102'seq\_len':128,103'epochs':128,104'batch\_size':16,105'inner\_iterations':25})

108experiment.add\_pytorch\_models(get\_modules(conf))

111withexperiment.start():

113conf.run()114115116if\_\_name\_\_=='\_\_main\_\_':117main()