pipe.py 21.9 KB
Newer Older
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
1
2
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

Tom Birch's avatar
Tom Birch committed
3
import argparse
4
5
6
from collections import defaultdict
from functools import reduce
import gc
7
import logging
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
8
import math
9
import operator
Tom Birch's avatar
Tom Birch committed
10
import os
11
import pprint
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
12
13
import time

14
15
16
17
from datasets.wikitext2_data import get_real_dataloaders as get_real_wikitext2_dataloaders
from datasets.wikitext2_data import get_synthetic_dataloaders as get_synthetic_wikitext2_dataloaders
from golden_configs import lm_wikitext2
from models import transformer_lm
18
import numpy as np
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
19
import torch
Tom Birch's avatar
Tom Birch committed
20
21
from torch.distributed import rpc
import torch.multiprocessing as mp
22
from torch.nn.parallel import DistributedDataParallel as DDP
23
from torch.optim import Adam
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
24

25
from fairscale.nn import Pipe
Tom Birch's avatar
Tom Birch committed
26
from fairscale.nn.model_parallel import initialize_model_parallel
27
from fairscale.nn.model_parallel.initialize import get_data_parallel_group, get_pipeline_parallel_group
28
from fairscale.nn.pipe import LazyModule, MultiProcessPipe
29
from fairscale.optim.oss import OSS
30
from fairscale.utils.testing import dist_init, get_worker_map
31
32


Tom Birch's avatar
Tom Birch committed
33
34
35
36
def init_random_seed(seed: int):

    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
37
    np.random.seed(seed)
Tom Birch's avatar
Tom Birch committed
38
39


40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def get_model_and_optimizer(args, device, config):
    """Return instantiated model and optimizer function."""

    if args.model_name == "lm":
        model = get_lm_model(args, device, config)

    lr = config["lr"]

    def make_adam(params):
        if args.ddp_zero:
            return OSS(params=params, optim=Adam, group=get_data_parallel_group(), lr=lr)
        else:
            return Adam(params, lr=lr)

    optimizer = make_adam
    return model, optimizer


def get_lm_model(args, device, config):
    """Get language model(based on GPT-2) used for sequence prediction."""

61
62
63
64
65
66
    ninp = config["ninp"]
    nhead = config["nhead"]
    initrange = config["initrange"]
    dropout = config["dropout"]
    vocab_size = config["vocab_size"]
    nhid = config["nhid"]
Tom Birch's avatar
Tom Birch committed
67
    ndecoder = args.num_decoder_layers
68

Tom Birch's avatar
Tom Birch committed
69
70
    if args.lazy_construction:
        layers = [
71
72
            LazyModule(lambda: transformer_lm.EmbeddingLayer(vocab_size, ninp, initrange)),
            LazyModule(lambda: transformer_lm.PositionalEncodingLayer(ninp, dropout)),
Tom Birch's avatar
Tom Birch committed
73
74
        ]
        for _ in range(ndecoder):
75
            layers.append(LazyModule(lambda: transformer_lm.TransformerDecoderLayer(ninp, nhead, nhid, dropout)))
Tom Birch's avatar
Tom Birch committed
76

77
        layers.append(LazyModule(lambda: transformer_lm.LinearLayer(ninp, vocab_size, initrange)))
Tom Birch's avatar
Tom Birch committed
78
79
        model = layers
    else:
80
        model = transformer_lm.TransformerLM(vocab_size, ninp, nhead, nhid, dropout, initrange, ndecoder).to(device)
81

82
    return model
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
83
84


Tom Birch's avatar
Tom Birch committed
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def get_tensors_by_size_bucket():

    size_buckets = defaultdict(int)
    for obj in gc.get_objects():
        if not isinstance(obj, torch.Tensor):
            continue
        if obj.device.type == "cuda":
            size_buckets[(*obj.size(),) + (obj.element_size(),)] += 1

    return size_buckets


def dump_size_buckets(size_buckets, prefix=""):

    total = 0
    for key, value in size_buckets.items():
        this = reduce(operator.mul, key) * value
        total += this
        print(prefix + f"{key} : {value}, {this}")

    print(prefix + f"total = {total}")


last_size_buckets = None
once = True


def safe_rank():
    try:
        return torch.distributed.get_rank()
    except AssertionError:
        return 0


def check_size_buckets():
    global last_size_buckets
    global once
    size_buckets = get_tensors_by_size_bucket()
    if last_size_buckets is not None:
        if size_buckets != last_size_buckets:
            print(f"difference is oustanding tensors: {safe-rank()}")
            dump_size_buckets(last_size_buckets, "old: ")
            dump_size_buckets(size_buckets, "new: ")
        if once:
            print(f"dumping buckets for: {safe_rank()}")
            dump_size_buckets(last_size_buckets, "old: ")
            dump_size_buckets(size_buckets, "new: ")
            once = False
    else:
        print(f"size buckets none on {safe_rank()}")
    last_size_buckets = size_buckets


def dump_cuda_tensors():
    print(f"dumping cuda tensors...")

    for obj in gc.get_objects():
        if not isinstance(obj, torch.Tensor):
            continue
        if obj.device.type == "cuda":
            size_buckets[(*obj.size(),) + (obj.element_size(),)] += 1

    print(f"outstanding cuda tensors:")
    total = 0
    for key, value in size_buckets.items():
        this = reduce(operator.mul, key) * value
        total += this
        print(f"{key} : {value}, {this}")
    print(f"total size = {total}")
    pprint.pprint(torch.cuda.memory_stats())


157
def log_number_of_parameters(model):
Tom Birch's avatar
Tom Birch committed
158
159

    num_params = reduce(operator.add, (reduce(operator.mul, x.size()) for x in model.parameters()))
160
    if hasattr(model, "group"):
161
162
163
        total = torch.Tensor([num_params])
        if torch.cuda.is_available():
            total = total.cuda()
164
165
        torch.distributed.all_reduce(total, group=model.group)
        logging.info(
166
            f"training model, #params = {num_params}, group: {model.group.rank()}, grank:"
167
168
169
170
171
            f" {torch.distributed.get_rank()}, sizes {model.group.size()}"
        )
        torch.distributed.barrier()
        if model.group.rank() == 0:
            logging.info(f"total #prams = {total.item()}")
Tom Birch's avatar
Tom Birch committed
172
    else:
173
        logging.info(f"training model, #params = {num_params}")
Tom Birch's avatar
Tom Birch committed
174
175


176
177
178
def get_device(model, index):
    if isinstance(model, DDP):
        model = model.module
179

180
181
182
183
184
185
    if not torch.cuda.is_available():
        return torch.device("cpu")
    if model.devices:
        return model.devices[index]
    else:
        return torch.cuda.current_device()
Tom Birch's avatar
Tom Birch committed
186

187

188
189
190
191
192
193
194
195
196
197
198
199
200
def get_fake_dataloader(lm_dataloader_len):
    fake_input = {"input": torch.zeros(args.batch_size)}

    class FakeDataset:
        def __getitem__(self, index):
            return fake_input

        def __len__(self):
            return lm_dataloader_len

    return FakeDataset()


201
202
def train(model_config, model, benchmark_config, args):
    lm_dataloader, _, _ = model_config["data"]
203
204
    criterion = benchmark_config["criterion"]
    vocab_size = benchmark_config["vocab_size"]
205
    optimizer = model_config["optimizer"]
206
207
208
209
210
211
212
213

    model.train()
    log_number_of_parameters(model)

    total_loss = 0.0
    word_counter = 0

    optimizer = optimizer(model.parameters())
Tom Birch's avatar
Tom Birch committed
214

215
    pipe_group = model.group if hasattr(model, "group") else None
216
217
218
219
220
221
222
223
224

    if args.ddp_zero:
        model = DDP(
            model,
            device_ids=[torch.cuda.current_device()],
            process_group=get_data_parallel_group(),
            find_unused_parameters=False,
        )

225
    # TODO(anj-s): Avoid sending fake data to all replicas except the first and last one.
226
    if pipe_group and pipe_group.rank() != 0 and pipe_group.rank() != (pipe_group.size() - 1):
227
        lm_dataloader = get_fake_dataloader(len(lm_dataloader))
228

229
230
    total_tokens = 0
    total_tokens_per_log_interval = 0
231
232
233
234
235
236
237
238
239
240
    bptt = 2
    start_time = time.time()
    epoch_start_time = 0.0

    def get_batch(source):
        seq_len = len(source) - 1
        data = source[0:seq_len]
        target = source[1 : 1 + seq_len]
        return data, target

Tom Birch's avatar
Tom Birch committed
241
    for i, batch in enumerate(lm_dataloader):
242
243
244
245
        if i == 1:
            epoch_start_time = time.time()

        source, target = get_batch(batch)
Tom Birch's avatar
Tom Birch committed
246
247
        if args.max_batch and i > args.max_batch:
            break
248
249
250

        if i > 0:
            total_tokens += source.numel()
251

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
252
        optimizer.zero_grad()
253
254
        try:
            if (pipe_group is None or pipe_group.rank() == 0) and not args.ddp_zero:
255
                tmp = source.to(get_device(model, 0))
256
257
                output = model(tmp)
            else:
258
                output = model(source)
259
260
261
262
        except Exception as e:
            raise RuntimeError(f"training failed on {torch.distributed.get_rank()}") from e

        if pipe_group is None or pipe_group.rank() == pipe_group.size() - 1:
263
            target = target.to(get_device(model, -1))
Tom Birch's avatar
Tom Birch committed
264
            output = output.to(target.device)
265

Tom Birch's avatar
Tom Birch committed
266
            loss = criterion(output.view(-1, vocab_size), target.view(-1))
267
268
269
270
            if args.ddp_zero:
                ddp_group = get_data_parallel_group()
                torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.SUM, group=ddp_group)
                loss /= ddp_group.size()
Tom Birch's avatar
Tom Birch committed
271
            loss.backward()
272
            del target
Tom Birch's avatar
Tom Birch committed
273
        else:
274
275
276
277
            if args.ddp_zero:
                model.module.back_helper(output)
            else:
                model.back_helper(output)
Tom Birch's avatar
Tom Birch committed
278
279
280

        del output

281
        torch.nn.utils.clip_grad_value_(model.parameters(), benchmark_config["clip_value"])
Tom Birch's avatar
Tom Birch committed
282
283
        optimizer.step()

284
        if pipe_group is None or pipe_group.rank() == pipe_group.size() - 1:
Tom Birch's avatar
Tom Birch committed
285
286
            total_loss += loss.item()
            log_interval = 1
287
            total_tokens_per_log_interval += source.numel()
Tom Birch's avatar
Tom Birch committed
288
289
290
            if i % log_interval == 0 and i > 0:
                cur_loss = total_loss / log_interval
                elapsed = time.time() - start_time
Jun Ru Anderson's avatar
Jun Ru Anderson committed
291
                print(
Tom Birch's avatar
Tom Birch committed
292
                    "| batch {:5d} | wps {:5.2f} | loss {:5.2f} | ppl {:8.2f}".format(
293
                        i, total_tokens_per_log_interval / elapsed, cur_loss, math.exp(cur_loss)
Jun Ru Anderson's avatar
Jun Ru Anderson committed
294
                    )
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
295
                )
296
                total_tokens_per_log_interval = 0
Tom Birch's avatar
Tom Birch committed
297
298
                total_loss = 0
                start_time = time.time()
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
299

300
301
302
303
304
305
306
307
    if epoch_start_time != 0:
        wps = total_tokens / (time.time() - epoch_start_time)
    else:
        raise RuntimeError(
            "Unable to benchmark on a single batch. Increase the size " " of the dataset and rerun the benchmark."
        )

    return wps, loss.item()
308

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
309

310
311
# TODO(anj-s): Add an option for users to be able to benchmark evaluate.
def evaluate(eval_model, data_source, criterion, ntokens):
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
312
313
    eval_model.eval()
    total_loss = 0.0
314
315
    # TODO(anj-s): Move this to the benchmark config if we want to benchmark evaluation.
    bptt = 35
316
317
318
319
320
321
322

    def get_batch(source, i, bptt):
        seq_len = min(bptt, len(source) - 1 - i)
        data = source[i : i + seq_len]
        target = source[i + 1 : i + 1 + seq_len].view(-1)
        return data, target

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
323
324
325
326
    with torch.no_grad():
        for i in range(0, data_source.size(0) - 1, bptt):
            data, targets = get_batch(data_source, i, bptt)
            output = eval_model(data)
327
            output = output.to(targets.device)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
328
329
330
331
332
333
334
335
336
            output_flat = output.view(-1, ntokens)
            total_loss += len(data) * criterion(output_flat, targets).item()
    return total_loss / (len(data_source) - 1)


def get_number_of_words(data):
    return data.size()[0] * data.size()[1]


337
def verify_lm_run(wps, golden_config):
338
339
340
    """Verify that words per second for a given benchmark run matches the golden data."""

    # Assert that words per second is within 3 standard deviations of the average
341
342
343
344
345
346
347
348
349
    # of five golden runs
    print("Throughput(wps) is {:.2f}.".format(wps))
    if not wps > (golden_config["avg_wps"] - (3 * golden_config["std_dev_wps"])):
        raise RuntimeError(
            "Throughput(wps):{:.2f} is below the golden threshold of an "
            "average value of {:.2f} and standard dev of {:.2f}.".format(
                wps, golden_config["avg_wps"], golden_config["std_dev_wps"]
            )
        )
350
351
352
353
354
355

    for i in range(4):
        print("Peak allocated bytes on cuda:0: {:1d}".format(torch.cuda.memory_stats(i)["allocated_bytes.all.peak"]))

    # Assert that memory usage on each GPU is within 10% of golden run
    # Right-hand-side is golden run bytes * 110%
356
357
358
359
360
361
362
    for i, golden_ref in zip(range(4), golden_config["peak_mem_usage"]):
        current_device_usage = torch.cuda.memory_stats(i)["allocated_bytes.all.peak"]
        if not current_device_usage < golden_ref * 1.1:
            raise RuntimeError(
                "Peak memory usage for cuda device {:d} is {:d} which"
                "is less than golden reference value of {:d}".format(i, current_device_usage, golden_ref)
            )
363
364


365
def benchmark_language_model(model_config, model, benchmark_config, args):
366
367
    golden_config = get_golden_config(args.model_name)
    epoch = benchmark_config["epochs"]
Jun Ru Anderson's avatar
Jun Ru Anderson committed
368
    print("-" * 110)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
369
    print("| start of epoch {:1d}".format(epoch))
Jun Ru Anderson's avatar
Jun Ru Anderson committed
370
    print("-" * 110)
371
    start_time = time.time()
372
    wps, loss = train(model_config, model, benchmark_config, args)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
373
    elapsed_time = time.time() - start_time
374
375
376
    print("-" * 110)
    print("| end of epoch {:1d} | time: {:5.2f}s | train loss {:5.2f} ".format(epoch, elapsed_time, loss))
    print("-" * 110)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
377

378
379
    print("wps ", wps)
    if len(model.balance) == 4:
380

381
        if args.model_name == "lm":
382
            verify_lm_run(wps, golden_config)
383
384
        else:
            raise RuntimeError("Unrecognized args.model_name " % args.model_name)
385
386


387
388
389
390
391
392
393
394
395
396
397
def generate_balance_weighted(num_devices, num_layers, fraction=0.5):
    balance = []
    layers_assigned = 0
    average_count = num_layers / num_devices
    last_layers = int(average_count * fraction)

    balance = generate_balance(num_devices - 1, num_layers - last_layers)
    balance.append(last_layers)
    return balance


398
399
400
401
402
403
404
405
406
407
408
409
410
def generate_balance(num_devices, num_layers):
    balance = []
    layers_assigned = 0
    for i in range(num_devices):
        x = (num_layers - layers_assigned) / (num_devices - i)
        if x.is_integer():
            balance.append(int(x))
            layers_assigned += x
        else:
            balance.append(math.ceil(x))
            layers_assigned += math.ceil(x)
    return balance

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
411

412
def get_synthetic_dataloaders(args, benchmark_config):
413
    """Returns dataloader for synthetic data."""
414

415
    if args.model_name == "lm":
416
        return get_synthetic_wikitext2_dataloaders(args, benchmark_config)
Tom Birch's avatar
Tom Birch committed
417
    else:
418
419
420
        raise RuntimeError("Unrecognized args.model_mame " % args.model_name)


421
def get_real_dataloaders(args, device, benchmark_config):
422
423
424
    """Returns dataloaders for real data."""

    if args.model_name == "lm":
425
        data = get_real_wikitext2_dataloaders(args, benchmark_config)
426
        ntokens, train_dataloader, valid_dataloader, test_dataloader = data
427
        benchmark_config["vocab_size"] = ntokens
428
        return train_dataloader, valid_dataloader, test_dataloader
429
430
431
432
    else:
        raise RuntimeError("Unrecognized args.model_mame " % args.model_name)


433
def create_model_config(args, benchmark_config=None):
434
435
436
    """Return a dict with the given model, dataset and optimizer."""

    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
437

438
    if args.use_synthetic_data:
439
        dataloader_fn = get_synthetic_dataloaders
440
    else:
441
442
443
444
445
446
447
448
449
        dataloader_fn = get_real_dataloaders

    data = dataloader_fn(args, device, benchmark_config)
    model, optimizer = get_model_and_optimizer(args, device, benchmark_config)
    return {
        "model": model,
        "optimizer": optimizer,
        "data": data,
    }
Tom Birch's avatar
Tom Birch committed
450
451


452
453
454
def create_benchmark_config(model_name):
    """Return a dict with configurations required for benchmarking `model_name` model."""

455
    if model_name == "lm":
456
457
458
459
460
461
462
463
464
465
        return lm_wikitext2.get_benchmark_config()
    else:
        raise RuntimeError("Unrecognized args.model_mame " % args.model_name)


def get_golden_config(model_name):
    """Return a dict with the golden data for throughput and memory usage."""

    if model_name == "lm":
        return lm_wikitext2.get_golden_real_stats()
466
467
    else:
        raise RuntimeError("Unrecognized args.model_mame " % args.model_name)
468
469
470
471
472


def benchmark_single_process(args):
    """Benchmark a given model using a single process and multiple devices."""

473
    num_devices = torch.cuda.device_count() if torch.cuda.is_available() else 1
474
    assert num_devices > 0
Tom Birch's avatar
Tom Birch committed
475
476
    init_random_seed(0)

477
    benchmark_config = create_benchmark_config(args.model_name)
478
    model_config = create_model_config(args, benchmark_config=benchmark_config)
479
    model = model_config["model"]
Tom Birch's avatar
Tom Birch committed
480

481
    balance = generate_balance(min(num_devices, 4), len(model))
482
    pipe_model = Pipe(model, balance, chunks=args.chunks, checkpoint=args.checkpoint)
483
    del model
484
    del model_config["model"]
Tom Birch's avatar
Tom Birch committed
485

486
    if args.dry_run:
487
        train(model_config, pipe_model, benchmark_config, args)
Tom Birch's avatar
Tom Birch committed
488
    else:
489
        benchmark_language_model(model_config, pipe_model, benchmark_config, args)
Tom Birch's avatar
Tom Birch committed
490
491
492
493


def run_mp_worker(args, available_workers):

494
    benchmark_config = create_benchmark_config(args.model_name)
495
    model_config = create_model_config(args, config=benchmark_config)
496
    model = model_config["model"]
Tom Birch's avatar
Tom Birch committed
497

498
    balance = generate_balance_weighted(get_pipeline_parallel_group().size(), len(model), 0.8)
499
    pipe_model = MultiProcessPipe(
Tom Birch's avatar
Tom Birch committed
500
501
502
503
        model,
        balance,
        chunks=args.chunks,
        worker_map=get_worker_map(),
504
        input_device=torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"),
Tom Birch's avatar
Tom Birch committed
505
506
        pipelined_backward=args.pipelined_backward,
        checkpoint=args.checkpoint,
507
        # TODO(anj-s): Do we need to comment this out? loss_fn=benchmark_config["criterion"],
508
509
    )
    if torch.cuda.is_available():
510
511
        pipe_model = pipe_model.cuda()
    if args.all_at_once and pipe_model.pipeline:
Tom Birch's avatar
Tom Birch committed
512
        print(f"running all at once")
513
        pipe_model.pipeline.all_at_once = True
Tom Birch's avatar
Tom Birch committed
514

515
    if args.use_synthetic_data:
516
        train(model_config, pipe_model, benchmark_config, args)
Tom Birch's avatar
Tom Birch committed
517
    else:
518
        benchmark_language_model(model_config, pipe_model, benchmark_config, args)
Tom Birch's avatar
Tom Birch committed
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554


def run_worker(rank, world_size, args):
    if args.world_size != 0:
        world_size = args.world_size
    dist_init(rank + args.rank_base, world_size, hostname=args.host)
    initialize_model_parallel(1, world_size)
    init_random_seed(0)
    run_mp_worker(args, world_size)

    rpc.shutdown()
    torch.distributed.destroy_process_group()


def bench_multi_process(args, all_at_once=False):
    if args.local_world_size != 0:
        world_size = args.local_world_size
    else:
        world_size = min(torch.cuda.device_count(), 2)
    mp.spawn(run_worker, args=(world_size, args), nprocs=world_size, join=True)


best_device_map = {
    0: "mlx5_0:1",
    1: "mlx5_0:1",
    2: "mlx5_1:1",
    3: "mlx5_1:1",
    4: "mlx5_2:1",
    5: "mlx5_2:1",
    6: "mlx5_3:1",
    7: "mlx5_3:1",
}


def bench_mpi(args):
    guess_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
555
556
557
    world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
    local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
    os.environ["UCX_NET_DEVICES"] = best_device_map[local_rank]
Tom Birch's avatar
Tom Birch committed
558
559

    os.environ["MASTER_ADDR"] = args.host
560
    os.environ["MASTER_PORT"] = "10638"
Tom Birch's avatar
Tom Birch committed
561
562
563
    if args.socket_name:
        os.environ["GLOO_SOCKET_IFNAME"] = args.socket_name
        os.environ["TP_SOCKET_IFNAME"] = args.socket_name
564
565
566
567
568

    torch.distributed.init_process_group(backend="gloo", rank=guess_rank, world_size=world_size)

    os.environ["MASTER_ADDR"] = args.host
    os.environ["MASTER_PORT"] = "10639"
Tom Birch's avatar
Tom Birch committed
569
570
571
    init_method = f"tcp://{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}"
    rank = torch.distributed.get_rank()
    world_size = torch.distributed.get_world_size()
572
    torch.cuda.set_device(local_rank % torch.cuda.device_count())
Tom Birch's avatar
Tom Birch committed
573
574
575
576
577
578
579
580
581

    rpc.init_rpc(
        f"Test{rank}",
        rank=rank,
        world_size=world_size,
        backend=rpc.BackendType.PROCESS_GROUP,
        rpc_backend_options=rpc.ProcessGroupRpcBackendOptions(rpc_timeout=20, init_method=init_method),
    )

582
583
584
585
586
587
    backends = {"model_parallel_backend": "nccl", "pipeline_backend": "mpi", "ddp_backend": "nccl"}

    if args.ddp_zero:
        initialize_model_parallel(1, 4, **backends)
    else:
        initialize_model_parallel(1, world_size, **backends)
Tom Birch's avatar
Tom Birch committed
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
    init_random_seed(0)

    run_mp_worker(args, world_size)

    rpc.shutdown()
    torch.distributed.destroy_process_group()


parser = argparse.ArgumentParser(description="benchmark")
parser.add_argument("--local-world-size", "-l", type=int, default=0, help="local world size")
parser.add_argument("--world-size", "-w", type=int, default=0, help="world size")
parser.add_argument("--rank-base", "-r", type=int, help="rank base", default=0)
parser.add_argument("--host", "-o", type=str, default="localhost", help="hostname")
parser.add_argument("--no-mpi", action="store_true", default=False, help="disable mpi")
parser.add_argument("--chunks", type=int, default=1, help="number of microbatches per batch")
parser.add_argument("--batch-size", type=int, default=8, help="size of a batch")
parser.add_argument("--all-at-once", action="store_true", default=False, help="do backward pass on whole batch at once")
parser.add_argument("--max-batch", type=int, default=4, help="Max number of batches")
parser.add_argument("--socket-name", type=str, default=None, help="socket ifname for gloo/tp")
parser.add_argument("--num-decoder-layers", type=int, default=10, help="Number of decoder layers in the model")
608
parser.add_argument("--ddp-zero", action="store_true", default=False, help="enable ddp")
Tom Birch's avatar
Tom Birch committed
609
610
611
612
613
614
615
616
617
618
619
620
parser.add_argument(
    "--lazy-construction", action="store_true", default=False, help="Number of decoder layers in the model"
)
parser.add_argument(
    "--checkpoint", default="never", choices=["always", "except_last", "never"], help="Checkpointing strategy for pipe"
)
parser.add_argument(
    "--pipelined-backward", dest="pipelined_backward", action="store_true", help="Pipelined backward pass"
)
parser.add_argument(
    "--no-pipelined-backward", dest="pipelined_backward", action="store_false", help="Pipelined backward pass"
)
621
622
parser.add_argument("--use_synthetic_data", action="store_true", help="Uses synthetic data for running benchmarks.")
parser.add_argument("--dry_run", action="store_true", help="Run a sample training run without regression testing.")
623
parser.add_argument(
624
625
626
627
    # TODO(anj-s): In the process of adding more models and hence the requirement for a flag.
    "--model_name",
    default="lm",
    help="Language Model(LM) used to benchmark nn.pipe.",
628
)
Tom Birch's avatar
Tom Birch committed
629
630
631
632
parser.set_defaults(pipelined_backward=True)

if __name__ == "__main__":
    args = parser.parse_args()
633
    # TODO(anj-s): Add support for multiprocess benchmarking.
Tom Birch's avatar
Tom Birch committed
634
635
    if args.no_mpi or "OMPI_COMM_WORLD_RANK" not in os.environ:
        print(f"Running benchmark with args: {args}")
636
        benchmark_single_process(args)
Tom Birch's avatar
Tom Birch committed
637
638
639
640
    else:
        if os.environ["OMPI_COMM_WORLD_RANK"] == "0":
            print(f"Running benchmark with args: {args}")
        bench_mpi(args)