-
Notifications
You must be signed in to change notification settings - Fork 43
Expand file tree
/
Copy pathdistributed.py
More file actions
110 lines (88 loc) · 3.61 KB
/
Copy pathdistributed.py
File metadata and controls
110 lines (88 loc) · 3.61 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import datetime
import torch.distributed
import torch.nn as nn
from labml import experiment, monit, logger, tracker
from labml.configs import option
from labml.logger import Text, inspect
from labml_helpers.train_valid import hook_model_outputs, BatchIndex
from python_autocomplete.train import Configs as Configs_
class Configs(Configs_):
ddp_model: nn.parallel.DistributedDataParallel
def init(self):
tracker.set_queue("loss.*", 20, True)
tracker.set_scalar("accuracy.*", True)
hook_model_outputs(self.mode, self.ddp_model, 'model')
self.state_modules = [self.accuracy]
def step(self, batch: any, batch_idx: BatchIndex):
data, target = batch[0].to(self.device), batch[1].to(self.device)
if self.mode.is_train:
tracker.add_global_step(len(data))
with self.mode.update(is_log_activations=batch_idx.is_last):
output, *_ = self.ddp_model(data)
loss = self.loss_func(output, target)
self.accuracy(output, target)
self.accuracy.track()
tracker.add("loss.", loss)
if self.mode.is_train:
loss.backward()
self.optimizer.step()
if batch_idx.is_last:
tracker.add('model', self.ddp_model)
self.optimizer.zero_grad()
tracker.save()
def sample(self):
prompt = 'def train('
log = [(prompt, Text.subtle)]
for i in monit.iterate('Sample', 25):
data = self.text.text_to_i(prompt).unsqueeze(-1)
data = data.to(self.device)
output, *_ = self.ddp_model(data)
output = output.argmax(dim=-1).squeeze()
prompt += '' + self.text.itos[output[-1]]
log += [('' + self.text.itos[output[-1]], Text.value)]
logger.log(log)
@option(Configs.ddp_model)
def ddp_model(c: Configs):
return nn.parallel.DistributedDataParallel(c.model, device_ids=[c.device])
def main(local_rank, rank, world_size, uuid, init_method: str = 'tcp://localhost:23456'):
with monit.section('Distributed'):
torch.distributed.init_process_group("gloo",
timeout=datetime.timedelta(seconds=30),
init_method=init_method,
rank=rank,
world_size=world_size)
conf = Configs()
experiment.create(uuid=uuid,
name="source_code_ddp",
comment='lstm model')
experiment.distributed(local_rank, world_size)
experiment.configs(conf, {
'model': 'transformer_model',
'n_layers': 6,
'batch_size': 12,
'epochs': 32,
'optimizer.optimizer': 'Noam',
'optimizer.learning_rate': 1.0,
'device.cuda_device': local_rank,
'seq_len': 512,
'train_loader': 'shuffled_train_loader',
'valid_loader': 'shuffled_valid_loader'
})
experiment.add_pytorch_models(model=conf.ddp_model)
with experiment.start():
conf.run()
def _launcher():
import os
world_size = int(os.environ['WORLD_SIZE'])
run_uuid = os.environ['RUN_UUID']
local_rank = int(os.environ['LOCAL_RANK'])
rank = int(os.environ['RANK'])
inspect(world_size=os.environ['WORLD_SIZE'],
run_uuid=os.environ['RUN_UUID'],
local_rank=os.environ['LOCAL_RANK'],
rank=os.environ['RANK'],
master_addr=os.environ['MASTER_ADDR'],
master_port=os.environ['MASTER_PORT'])
main(local_rank, rank, world_size, run_uuid, 'env://')
if __name__ == '__main__':
_launcher()