加速器内部机制
在内部,Accelerate 通过首先分析脚本启动的环境来确定使用了哪种分布式设置、存在多少不同的进程以及当前脚本位于哪个进程中。所有这些信息都存储在~AcceleratorState
中。
此类在您第一次实例化 ~Accelerator 以及执行分布式设置所需的任何特定初始化时初始化。然后,它的状态通过所有 AcceleratorState 实例以唯一的方式共享。(同样的事情也可以使用它继承的更精简版本的 PartialState 来完成)
然后,当调用 prepare() 时,库将
- 将您的模型包装在适合分布式设置的容器中,
- 将您的优化器包装在 AcceleratedOptimizer 中,
- 将您的调度器包装在 AcceleratedScheduler 中
- 在 DataLoaderShard 或 DataLoaderDispatcher 中创建您的数据加载器的新版本
虽然模型、优化器和调度器只是放在简单的包装器中,但数据加载器会被重新创建。这主要是因为 PyTorch 不允许用户在数据加载器创建后更改其batch_sampler
,并且库通过更改batch_sampler
来处理数据在进程之间的分片,以便生成每隔num_processes
个批次(如果已启用)。
DataLoaderShard 对DataLoader
进行子类化,以添加以下功能
- 它在每次新的迭代中同步所有进程的适当随机数生成器,以确保任何随机化(如洗牌)在所有进程中都以完全相同的方式完成。
- 在生成批次之前,它会将批次放在适当的设备上(除非您已选择不使用
device_placement=True
)。
DataLoaderDispatcher 子类与 DataLoaderShard 不同,因为在遍历DataLoader
时,数据全部从进程 0 开始,然后再进行拆分并发送到每个进程,而不是在数据集级别进行。
随机数生成器同步默认情况下会同步
- 对于 PyTorch >= 1.6,给定采样器(如 PyTorch
RandomSampler
)的generator
属性 - 对于 PyTorch <=1.5.1,主随机数生成器
您可以使用主 Accelerator 的rng_types
参数选择要同步的随机数生成器。在 PyTorch >= 1.6 中,建议依赖于本地generator
以避免在所有进程的主随机数生成器中设置相同的种子。
同步主 torch(或 CUDA 或 XLA)随机数生成器会影响您可能在数据集中拥有的任何其他潜在的随机伪像(如随机数据增强),因为所有进程都将从 torch 随机模块获取相同的随机数(因此如果随机数据增强由 torch 控制,则会应用相同的随机数据增强)。
自定义采样器、批次采样器或可迭代数据集的随机化部分应使用本地torch.Generator
对象(在 PyTorch >= 1.6 中)完成,以传统的RandomSampler
为例。
如果您已安装 torchdata>=0.8.0
,并且您已将use_stateful_dataloader=True
传递到您的 DataLoaderConfiguration 中,这些类将直接继承自StatefulDataLoader
,并维护一个state_dict
。
有关内部机制的更多详细信息,请参阅 内部页面。
< > 在 GitHub 上更新