16importtorch17fromtorch.utils.dataimportDataLoader1819fromlabmlimportexperiment,tracker20fromlabml.configsimportoption,calculate21fromlabml\_helpers.datasets.textimportSequentialUnBatchedDataset22fromlabml\_nn.transformers.alibiimportAlibiMultiHeadAttention23fromlabml\_nn.experiments.nlp\_autoregressionimporttranspose\_batch24fromlabml\_nn.transformersimportTransformerConfigs25fromlabml\_nn.transformers.gptimportConfigsasGPTConfigs

28classConfigs(GPTConfigs):

36transformer:TransformerConfigs='GPT\_ALiBi'

38valid\_seq\_len:int=12839valid\_loader='shuffled\_longer\_valid\_loader'

41defother\_metrics(self,output:torch.Tensor,target:torch.Tensor):

46ifself.seq\_len\<output.shape[0]:

48tracker.add(f'loss.{self.seq\_len-1}.',self.loss\_func(output[self.seq\_len-1],target[self.seq\_len-1]))

50tracker.add(f'loss.0.',self.loss\_func(output[0],target[0]))

52tracker.add(f'loss.{int(output.shape[0])-1}.',self.loss\_func(output[-1],target[-1]))

55def\_alibi\_mha(c:TransformerConfigs):

59returnAlibiMultiHeadAttention(c.n\_heads,c.d\_model,dropout\_prob=c.dropout)

63calculate(TransformerConfigs.encoder\_attn,'alibi\_mha',\_alibi\_mha)64calculate(TransformerConfigs.decoder\_attn,'alibi\_mha',\_alibi\_mha)65calculate(TransformerConfigs.decoder\_mem\_attn,'alibi\_mha',\_alibi\_mha)

68@option(Configs.valid\_loader)69defshuffled\_longer\_valid\_loader(c:Configs):

73returnDataLoader(SequentialUnBatchedDataset(text=c.text.valid,74dataset=c.text,75seq\_len=c.valid\_seq\_len),76batch\_size=c.batch\_size,77collate\_fn=transpose\_batch,78shuffle=True)

81@option(Configs.transformer,'GPT\_ALiBi')82def\_transformer\_configs(c:Configs):

89conf=TransformerConfigs()

91conf.n\_src\_vocab=c.n\_tokens92conf.n\_tgt\_vocab=c.n\_tokens

94conf.ffn.activation='GELU'

97conf.src\_embed='no\_pos'98conf.tgt\_embed='no\_pos'

101conf.encoder\_attn='alibi\_mha'102conf.decoder\_attn='alibi\_mha'103conf.decoder\_mem\_attn='alibi\_mha'

106returnconf

109defmain():

111experiment.create(name="gpt\_alibi")

113conf=Configs()

115experiment.configs(conf,{

117'tokenizer':'character',

119'prompt\_separator':'',

121'prompt':'It is ',

123'text':'tiny\_shakespeare',

127'seq\_len':64,

129'valid\_seq\_len':80,

131'epochs':128,

133'batch\_size':128,

136'inner\_iterations':10,

139'transformer.d\_model':128,140'transformer.ffn.d\_ff':512,141'transformer.n\_heads':8,142'transformer.n\_layers':4,143'transformer.dropout':0.1,144})

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

150withexperiment.start():

152conf.run()

156if\_\_name\_\_=='\_\_main\_\_':157main()