35importtorch36fromtorchimportnn3738fromlabmlimportexperiment39fromlabml.configsimportoption40fromlabml\_helpers.moduleimportModule41fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs42fromlabml\_nn.optimizers.configsimportOptimizerConfigs43fromlabml\_nn.transformersimportTransformerConfigs,Encoder44fromlabml\_nn.transformers.utilsimportsubsequent\_mask

47classGPT(Module):

55def\_\_init\_\_(self,encoder:Encoder,src\_embed:Module,generator:Module):

62super().\_\_init\_\_()63self.src\_embed=src\_embed64self.encoder=encoder65self.generator=generator

68self.mask=None

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

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

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

77x=self.src\_embed(x)

79x=self.encoder(x,self.mask)

81x=self.generator(x)

85returnx,None

88classConfigs(NLPAutoRegressionConfigs):

97model:GPT

99transformer:TransformerConfigs

101weight\_decay:float=0.1

103warmup\_steps:int=128\*128\*20

106optimizer='transformer\_optimizer'

109@option(Configs.transformer,'GPT')110def\_transformer\_configs(c:Configs):

117conf=TransformerConfigs()

119conf.n\_src\_vocab=c.n\_tokens120conf.n\_tgt\_vocab=c.n\_tokens

122conf.ffn.activation='GELU'

125returnconf

128def\_init\_weights(module):

137ifnotisinstance(module,(nn.Linear,nn.Embedding)):138return139140module.weight.data.normal\_(mean=0.0,std=0.02)

143ifisinstance(module,nn.Linear)andmodule.biasisnotNone:144module.bias.data.zero\_()

147@option(Configs.model)148def\_model(c:Configs):

152m=GPT(c.transformer.encoder,153c.transformer.src\_embed,154c.transformer.generator).to(c.device)

157m.apply(\_init\_weights)158159returnm

162@option(NLPAutoRegressionConfigs.optimizer)163deftransformer\_optimizer(c:NLPAutoRegressionConfigs):

171decay=set()172formn,minc.model.named\_modules():173forpn,pinm.named\_parameters():174fpn=f'{mn}.{pn}'ifmnelsepn# full param name175176iffpn.endswith('weight')andisinstance(m,nn.Linear):177decay.add(fpn)

180param\_dict={pn:pforpn,pinc.model.named\_parameters()}

182no\_decay=set(param\_dict.keys())-decay

185opt\_groups=[186{"params":[param\_dict[pn]forpninsorted(list(decay))],"weight\_decay":c.weight\_decay},187{"params":[param\_dict[pn]forpninsorted(list(no\_decay))],"weight\_decay":0.0},188]

193optimizer=OptimizerConfigs()

196optimizer.parameters=opt\_groups

199optimizer.optimizer='AdamWarmupCosineDecay'

202optimizer.d\_model=c.d\_model

205optimizer.weight\_decay=c.weight\_decay

207optimizer.learning\_rate=6e-4

209optimizer.betas=(0.9,0.95)

211optimizer.eps=1e-8

213optimizer.weight\_decouple=True

215optimizer.total\_steps=c.epochs\*len(c.text.train)//(c.batch\_size\*c.seq\_len)

217optimizer.warmup=c.warmup\_steps//(c.batch\_size\*c.seq\_len)218219returnoptimizer

222defmain():

224experiment.create(name="gpt")

226conf=Configs()

228experiment.configs(conf,{

230'tokenizer':'character',

232'prompt\_separator':'',

234'prompt':'It is ',

236'text':'tiny\_shakespeare',

239'seq\_len':128,

241'epochs':32,

243'batch\_size':128,

246'inner\_iterations':10,

249'transformer.d\_model':512,250'transformer.ffn.d\_ff':2048,251'transformer.n\_heads':8,252'transformer.n\_layers':6253})

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

259withexperiment.start():

261conf.run()

265if\_\_name\_\_=='\_\_main\_\_':266main()