PyTorch分布式训练

数据并行训练

PyTorch 为数据并行训练提供了多种选项。一般来说，应用会从简单到复杂，从原型到量产。这些应用共同的发展轨迹是：

  1. 如果数据和模型可以放在一个 GPU 中，并且不关心训练速度，就使用单设备（single-device）训练。
  2. 如果服务器上有多个 GPU，并且您希望以最少的代码更改来加速训练，那么可以使用单机多 GPU DataParallel。
  3. 如果您想进一步加快训练速度并愿意编写更多代码来设置它，可以使用单机多 GPU DistributedDataParallel。
  4. 如果应用程序需要跨机器边界进行扩展，请使用多机 DistributedDataParallel 和 启动脚本。
  5. 如果预期会出现错误（例如，OOM）或者资源可以在训练期间动态加入和离开，则使用torchelastic启动分布式训练。

Untitled

Torch如何使用GPU

Untitled

  • _apply 方法
    • 遍历 _parameters：
      • 对参数调用fn进行处理，得到param_applied。
        • 用 param_applied 重新设置参数。
      • 如果参数有梯度，则：
        • 对参数的grad调用fn进行处理，得到grad_applied。
        • 用 grad_applied 重新设置参数的梯度。
    • 遍历 _buffers：
      • 对buf调用fn进行处理。

调用 cuda 或者 to 方法来移动模型到GPU，其实就是把模型的self._parameters和 self._buffers移动到 GPU，并没有对 self._modules进行移动。这个移动过程是递归调用的，是把模型每个叶子都移动到了 GPU 之上。

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
+
                                        |
+---------------------------------+     |     +----------------------------------+
| CPU                             |     |     | CPU                              |
|  +--------------+               |     |     |       +--------------------+     |
|  |Module        |               |     |     |       | Module             |     |
|  |              |               |     |     |       |                    |     |
|  | _parameters+----> Parameters |     |     |       |     _parameters ------+  |
|  |              |               |     |     |       |                    |  |  |
|  | _buffers +------> Buffers    |     |     |     +-----+ _buffers       |  |  |
|  |              |               |     |     |     | |                    |  |  |
|  | _modules     |               |     |     |     | |     _modules       |  |  |
|  |              |               |     |     |     | |                    |  |  |
|  +--------------+               |     |     |     | +--------------------+  |  |
|                                 |     |     |     |                         |  |
+---------------------------------+     |     +----------------------------------+
                                        |           |                         |
                                        +           |                         |
+------------------------------->  Module.cuda() +---------------------------------> Time
                                        +           |                         |
                                        |           |                         |
+---------------------------------+     |     +----------------------------------+
| GPU                             |     |     | GPU |                         |  |
|                                 |     |     |     |                         |  |
|                                 |     |     |     |       Parameters  <-----+  |
|                                 |     |     |     |                            |
|                                 |     |     |     |                            |
|                                 |     |     |     +---->  Buffers              |
|                                 |     |     |                                  |
|                                 |     |     |                                  |
+---------------------------------+     |     +----------------------------------+
                                        |
                                        +