21importcopy2223importtorch24importtorch.nnasnn2526fromlabmlimportexperiment27fromlabml.configsimportoption28fromlabml\_helpers.moduleimportModule29fromlabml\_nn.activations.ftaimportFTA30fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs31fromlabml\_nn.transformersimportMultiHeadAttention,TransformerLayer32fromlabml\_nn.transformers.utilsimportsubsequent\_mask

35classFeedForwardFTA(nn.Module):

40def\_\_init\_\_(self,d\_model:int,d\_ff:int,41activation:FTA,42dropout:float=0.1):

49super().\_\_init\_\_()

51self.layer1=nn.Linear(d\_model,d\_ff)

53self.layer2=nn.Linear(d\_ff\*activation.expansion\_factor,d\_model)

55self.dropout=nn.Dropout(dropout)

57self.activation=activation

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

61x=self.activation(self.layer1(x))

63x=self.dropout(x)

65returnself.layer2(x)

68classAutoregressiveTransformer(Module):

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

83super().\_\_init\_\_()

85self.transformer\_layers=nn.ModuleList([copy.deepcopy(layer)for\_inrange(n\_layers)])

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

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

93self.mask=None

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

100ifself.maskisNoneorself.mask.size(0)!=len(x):

102self.mask=subsequent\_mask(len(x)).to(x.device)

105x=self.emb(x)

107forlayerinself.transformer\_layers:108x=layer(x=x,mask=self.mask)

110x=self.readout(x)

113returnx,None

116classConfigs(NLPAutoRegressionConfigs):

125model:AutoregressiveTransformer

128n\_layers:int=4

131deep\_norm\_alpha:float132deep\_norm\_beta:float

135n\_heads:int=4

137d\_model:int=256

139d\_k:int=16

141d\_ff:int=256

144fta\_lower\_limit:float=-1.145fta\_upper\_limit:float=+1.146fta\_delta:float=0.2147fta\_eta:float=0.05

150@option(Configs.model)151def\_model(c:Configs):

157fta=FTA(c.fta\_lower\_limit,c.fta\_upper\_limit,c.fta\_delta,c.fta\_eta)

161m=AutoregressiveTransformer(c.n\_tokens,c.d\_model,c.n\_layers,162TransformerLayer(d\_model=c.d\_model,163feed\_forward=FeedForwardFTA(d\_model=c.d\_model,164d\_ff=c.d\_ff,165activation=fta,166dropout=0.1),167self\_attn=MultiHeadAttention(c.n\_heads,c.d\_model,168dropout\_prob=0.0),169dropout\_prob=0.0))

172returnm.to(c.device)

175defmain():

180experiment.create(name="fta",writers={'screen','labml'})

182conf=Configs()

184experiment.configs(conf,{

186'tokenizer':'character',

188'prompt\_separator':'',

190'prompt':'It is ',

192'text':'tiny\_shakespeare',

195'seq\_len':256,

197'epochs':32,

199'batch\_size':16,

201'inner\_iterations':10,

204'optimizer.optimizer':'Adam',205'optimizer.learning\_rate':3e-4,206})

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

212withexperiment.start():

214conf.run()

218if\_\_name\_\_=='\_\_main\_\_':219main()