15importtorch16importtorch.nnasnn1718fromlabmlimportexperiment,tracker19fromlabml.configsimportoption20fromlabml\_helpers.moduleimportModule21fromlabml\_helpers.train\_validimportBatchIndex22fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs

25classAutoregressiveModel(Module):

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

33self.src\_embed=nn.Embedding(n\_vocab,d\_model)

35self.transformer=transformer

37self.generator=nn.Linear(d\_model,n\_vocab)38self.mask=None

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

42ifself.maskisNoneorself.mask.size(0)!=len(x):43fromlabml\_nn.transformers.utilsimportsubsequent\_mask44self.mask=subsequent\_mask(len(x)).to(x.device)

46x=self.src\_embed(x)

48res,counts,route\_prob,n\_dropped,route\_prob\_max=self.transformer(x,self.mask)

50res=self.generator(res)

52returnres,counts,route\_prob,n\_dropped,route\_prob\_max

55classConfigs(NLPAutoRegressionConfigs):

64model:AutoregressiveModel65transformer:Module

68d\_model:int=128

70heads:int=4

72dropout:float=0.0

74d\_ff:int=256

76n\_layers:int=6

78n\_experts:int=4

80load\_balancing\_loss\_ceof=0.01

82is\_scale\_prob:bool=True

84drop\_tokens:bool=False

86capacity\_factor:float=1.0

88definit(self):89super().init()

91tracker.set\_scalar("lb\_loss.\*",False)92tracker.set\_scalar("route.\*",False)93tracker.set\_scalar("dropped.\*",False)

95defstep(self,batch:any,batch\_idx:BatchIndex):

101data,target=batch[0].to(self.device),batch[1].to(self.device)

104ifself.mode.is\_train:105tracker.add\_global\_step(data.shape[0]\*data.shape[1])

108withself.mode.update(is\_log\_activations=batch\_idx.is\_last):

110output,counts,route\_prob,n\_dropped,route\_prob\_max=self.model(data)

113cross\_entropy\_loss=self.loss\_func(output,target)

115total=counts.sum(dim=-1,keepdims=True)

119route\_frac=counts/total

122route\_prob=route\_prob/total

127load\_balancing\_loss=self.n\_experts\*(route\_frac\*route\_prob).sum()

130tracker.add('dropped.',total.new\_tensor(n\_dropped)/total)131tracker.add('route.min.',route\_frac.min())132tracker.add('route.max.',route\_frac.max())133tracker.add('route.std.',route\_frac.std())134tracker.add('route.max\_prob.',route\_prob\_max)135tracker.add("loss.",cross\_entropy\_loss)136tracker.add("lb\_loss.",load\_balancing\_loss)

141loss=cross\_entropy\_loss+self.load\_balancing\_loss\_ceof\*load\_balancing\_loss

144self.accuracy(output,target)145self.accuracy.track()

148ifself.mode.is\_train:

150loss.backward()

152torch.nn.utils.clip\_grad\_norm\_(self.model.parameters(),max\_norm=self.grad\_norm\_clip)

154self.optimizer.step()

156ifbatch\_idx.is\_last:157tracker.add('model',self.model)

159self.optimizer.zero\_grad()

162tracker.save()

165@option(Configs.model)166defautoregressive\_model(c:Configs):

170m=AutoregressiveModel(c.n\_tokens,c.d\_model,c.transformer)171returnm.to(c.device)

174@option(Configs.transformer)175defswitch\_transformer(c:Configs):

179fromlabml\_nn.transformers.switchimportSwitchTransformer,SwitchTransformerLayer,SwitchFeedForward180fromlabml\_nn.transformersimportMultiHeadAttention181fromlabml\_nn.transformers.feed\_forwardimportFeedForward182183returnSwitchTransformer(184SwitchTransformerLayer(d\_model=c.d\_model,185attn=MultiHeadAttention(c.heads,c.d\_model,c.dropout),186feed\_forward=SwitchFeedForward(capacity\_factor=c.capacity\_factor,187drop\_tokens=c.drop\_tokens,188is\_scale\_prob=c.is\_scale\_prob,189n\_experts=c.n\_experts,190expert=FeedForward(c.d\_model,c.d\_ff,c.dropout),191d\_model=c.d\_model),192dropout\_prob=c.dropout),193c.n\_layers)

196defmain():

201experiment.create(name="switch\_transformer",comment='')

203conf=Configs()

205experiment.configs(conf,

207{'tokenizer':'character',208'text':'tiny\_shakespeare',209'optimizer.learning\_rate':1.,210'optimizer.optimizer':'Noam',211'prompt':'It is',212'prompt\_separator':'',213214'transformer':'switch\_transformer',215'n\_experts':4,216217'drop\_tokens':True,218'capacity\_factor':1.2,219220'train\_loader':'shuffled\_train\_loader',221'valid\_loader':'shuffled\_valid\_loader',222223'seq\_len':64,224'epochs':128,225'batch\_size':32,226'inner\_iterations':25,227})

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

233withexperiment.start():

235conf.run()

239if\_\_name\_\_=='\_\_main\_\_':240main()