Performs NCCL AllReduce to distribute blobs to all the GPUs.
(blob_names, devices, model, net, use_nccl)
| 1425 | |
| 1426 | |
| 1427 | def _AllReduceBlobsSingleHost(blob_names, devices, model, net, use_nccl): |
| 1428 | """Performs NCCL AllReduce to distribute blobs to all the GPUs.""" |
| 1429 | |
| 1430 | if len(devices) == 1: |
| 1431 | return |
| 1432 | |
| 1433 | # Now we need to Allreduce blobs on all the GPUs. |
| 1434 | # Pick GPU #0 as a master GPU. |
| 1435 | master_device_opt = core.DeviceOption(model._device_type, devices[0]) |
| 1436 | last_out = None |
| 1437 | concatenated_idx = set() |
| 1438 | |
| 1439 | for blob_name in blob_names: |
| 1440 | # Group by blob_name for reduce. |
| 1441 | blobs_group = list(model._device_grouped_blobs[blob_name].values()) |
| 1442 | if len(blobs_group) == 1: |
| 1443 | # Non-reducible |
| 1444 | continue |
| 1445 | assert len(blobs_group) == len(devices), \ |
| 1446 | "Each GPU from {}, should have a copy of {}.".format( |
| 1447 | devices, blob_name) |
| 1448 | |
| 1449 | if _IsGPUBlob(model, blob_name): |
| 1450 | with core.DeviceScope(master_device_opt): |
| 1451 | if not isinstance(blobs_group[0], core.GradientSlice): |
| 1452 | _AllReduce( |
| 1453 | devices, model, net, blob_name, use_nccl, last_out |
| 1454 | ) |
| 1455 | # last_out is used to serialize the execution of nccls |
| 1456 | last_out = blobs_group[0] |
| 1457 | |
| 1458 | else: |
| 1459 | # Sparse gradients: all-gather for indices and values |
| 1460 | master_ns = "{}_{}".format(model._device_prefix, devices[0]) |
| 1461 | ''' |
| 1462 | Skip if we have already copied concatenated indices |
| 1463 | to the indices of GradientSlice. This happens when two |
| 1464 | or more grad blobs are gathered with the same indices |
| 1465 | blob |
| 1466 | ''' |
| 1467 | skip_idx_concat = False |
| 1468 | for g in blobs_group: |
| 1469 | if g.indices in concatenated_idx: |
| 1470 | skip_idx_concat = True |
| 1471 | |
| 1472 | if not skip_idx_concat: |
| 1473 | grad_idx_concat, _ = net.Concat( |
| 1474 | [g.indices for g in blobs_group], |
| 1475 | ["{}/{}_index_concat".format(master_ns, blob_name), |
| 1476 | "{}/{}_index_splitinfo".format(master_ns, blob_name)], |
| 1477 | axis=0, |
| 1478 | name="note:data_parallel_model") |
| 1479 | |
| 1480 | for gpu, g in model._device_grouped_blobs[blob_name].items(): |
| 1481 | device_opt = core.DeviceOption(model._device_type, gpu) |
| 1482 | with core.DeviceScope(device_opt): |
| 1483 | model.Copy(grad_idx_concat, g.indices) |
| 1484 | concatenated_idx.add(g.indices) |
no test coverage detected
searching dependent graphs…