14importdatetime1516importtorch17importtorch.distributed1819fromlabmlimportexperiment,monit,tracker20fromlabml.configsimportoption21fromlabml.loggerimportinspect22fromlabml\_nn.neox.samples.finetuneimportPipelineParallelTrainerConf

27classConfigs(PipelineParallelTrainerConf):28rank:int29world\_size:int

32@option(Configs.optimizer,'Zero3Adam')33def\_optimizer(c:Configs):

39fromlabml\_nn.optimizers.adam\_fp16importAdamFP1640returnAdamFP16(c.model.get\_trainable\_chunk(),lr=c.learning\_rate)

43@option(Configs.model,'Zero3')44def\_model(c:Configs):

48fromlabml\_nn.scaling.zero3importZero3Layer,Zero3Sequential

51\_=c.fine\_tuner

54modules=[]55forminmonit.iterate('Zero3',c.layers):56modules.append(Zero3Layer(m.to(c.device),57c.rank,c.world\_size,c.device,c.dtype))

60model=Zero3Sequential(modules)

63returnmodel

66defmain(rank:int,world\_size:int,init\_method:str='tcp://localhost:23456'):

71withmonit.section('Distributed'):72torch.distributed.init\_process\_group('nccl',73timeout=datetime.timedelta(seconds=30),74init\_method=init\_method,75rank=rank,76world\_size=world\_size)

79device=torch.device(f'cuda:{rank}')80torch.cuda.set\_device(device)

83experiment.create(name='zero3\_neox',writers={'screen','labml'})84experiment.distributed(rank,world\_size)

87conf=Configs()

90experiment.configs(conf,{91'model':'Zero3',92'optimizer':'Zero3Adam',9394'device':device,95'rank':rank,96'world\_size':world\_size,9798'learning\_rate':3e-4,99'max\_seq\_len':128,100'batch\_size':16,101})

104withexperiment.start():

106\_=conf.model

109forepochinmonit.loop(conf.epochs):110conf.train\_epoch()111tracker.new\_line()

115if\_\_name\_\_=='\_\_main\_\_':

117inspect([torch.cuda.get\_device\_name(i)foriinrange(torch.cuda.device\_count())])118inspect(119n\_gpus=torch.cuda.device\_count(),120mpi=torch.distributed.is\_mpi\_available(),121nccl=torch.distributed.is\_nccl\_available(),122)123124n\_gpu=torch.cuda.device\_count()

127torch.multiprocessing.spawn(main,args=(n\_gpu,),nprocs=n\_gpu,join=True)