13importcopy1415importtorch16importtorch.nnasnn1718fromlabmlimportexperiment19fromlabml.configsimportoption20fromlabml\_helpers.moduleimportModule21fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs22fromlabml\_nn.normalization.deep\_normimportDeepNormTransformerLayer23fromlabml\_nn.transformersimportMultiHeadAttention24fromlabml\_nn.transformers.feed\_forwardimportFeedForward

27classAutoregressiveTransformer(Module):

34def\_\_init\_\_(self,n\_tokens:int,d\_model:int,n\_layers:int,layer:DeepNormTransformerLayer):

41super().\_\_init\_\_()

43self.transformer=nn.Sequential(\*[copy.deepcopy(layer)for\_inrange(n\_layers)])

46self.emb=nn.Embedding(n\_tokens,d\_model)

48self.readout=nn.Linear(d\_model,n\_tokens)

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

55x=self.emb(x)

57x=self.transformer(x)

59x=self.readout(x)

62returnx,None

65classConfigs(NLPAutoRegressionConfigs):

74model:AutoregressiveTransformer

77n\_layers:int=32

80deep\_norm\_alpha:float81deep\_norm\_beta:float

84n\_heads:int=4

86d\_model:int=64

88d\_k:int=16

91@option(Configs.deep\_norm\_alpha)92def\_deep\_norm\_alpha(c:Configs):

98return(2.\*c.n\_layers)\*\*(1./4.)

101@option(Configs.deep\_norm\_beta)102def\_deep\_norm\_beta(c:Configs):

108return(8.\*c.n\_layers)\*\*-(1./4.)

111@option(Configs.model)112def\_model(c:Configs):

116m=AutoregressiveTransformer(c.n\_tokens,c.d\_model,c.n\_layers,117DeepNormTransformerLayer(d\_model=c.d\_model,118deep\_norm\_alpha=c.deep\_norm\_alpha,119deep\_norm\_beta=c.deep\_norm\_beta,120feed\_forward=FeedForward(d\_model=c.d\_model,121d\_ff=c.d\_model\*4),122self\_attn=MultiHeadAttention(c.n\_heads,c.d\_model,123dropout\_prob=0.0)))124125returnm.to(c.device)

128defmain():

133experiment.create(name="deep\_norm",writers={'screen','web\_api'})

135conf=Configs()

137experiment.configs(conf,{

139'tokenizer':'character',

141'prompt\_separator':'',

143'prompt':'It is ',

145'text':'tiny\_shakespeare',

148'seq\_len':256,

150'epochs':32,

152'batch\_size':16,

154'inner\_iterations':10,

157'n\_layers':50,

161'optimizer.optimizer':'Adam',162'optimizer.learning\_rate':1.25e-4,163})

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

169withexperiment.start():

171conf.run()

175if\_\_name\_\_=='\_\_main\_\_':176main()