执行和延迟作业

当您运行通常的脚本时，指令按顺序执行。使用 Accelerate 在多个 GPU 上同时部署您的脚本会引入一个复杂情况：虽然每个进程都按顺序执行所有指令，但有些进程可能比其他进程更快。

您可能需要等待所有进程都到达某个点才能执行给定的指令。例如，在确保每个进程都完成了训练之前，您不应该保存模型，并且在所有模型权重都加载之前，您不希望继续训练。要做到这一点，只需在您的代码中写入以下行

accelerator.wait_for_everyone()

此指令将阻塞最先到达的所有进程，直到所有其他进程都到达该点（如果您只在一个 GPU 或 CPU 上运行脚本，则此指令不会有任何作用）。

下面列出了一些使用此实用程序的示例情况

其中一些与 main_process_first() 上下文管理器一起使用，该管理器利用 wait_for_everyone() 在触发并启动其他进程之前，先在主进程上运行特定的代码集

下载数据集

下载数据集时，应首先在主进程上下载它，然后加载缓存的数据集。

load_dataset 将在后台执行锁定以阻止同时发生多个下载，但如果您正在下载未使用此库的内容，则应使用此方法。

with accelerator.main_process_first():
    datasets = load_dataset("glue", "mrpc")

在后台，这与调用以下内容相同：

# First do something on the main process
if accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

# And then send it to the rest of them
if not accelerator.is_main_process:
    datasets = load_dataset("glue", "mrpc")
else:
    accelerator.wait_for_everyone()

保存 state_dict

保存模型的 state_dict 时，因为您通常只在主进程上保存一个文件，所以您应该指定这一点

if accelerator.is_main_process:
    model = accelerator.unwrap_model(model)
    torch.save(model.state_dict(), "weights.pth")

加载 state_dict

将 state_dict 加载到模型、优化器或调度器时，应等待所有工作进程加载权重，然后再继续训练。

with accelerator.main_process_first():
    state = torch.load("weights.pth")
    model.load_state_dict(state)

应用多工作进程 CPU 操作

在多个工作进程上应用 map() 操作（例如标记化）应首先在主进程上完成，然后传播到每个工作进程。

datasets = load_dataset("glue", "mrpc")

with accelerator.main_process_first():
    tokenized_datasets = datasets.map(
        tokenize_function,
        batched=True,
        remove_columns=["idx", "sentence1", "sentence2"],
    )

应用检查（例如提前停止）

要进行由特定进程设置的标志生效的检查，应使用 set_trigger 和 check_trigger API。执行此操作的有用示例可能包括使用提前停止和监控损失（因为每个进程的损失略有不同）等情况。

当您的条件满足时，调用 Accelerator.set_trigger()，当检查任何进程中是否满足该条件时，调用 Accelerator.check_trigger()

for (x,y) in data_loader:
    logits = model(x)
    loss = loss_func(logits, y)
    # Assume `should_do_early_stopping` is a custom defined function that returns a conditional
    if should_do_early_stopping(loss):
        accelerator.set_trigger()

    # Later in the training script when we need to check for the breakpoint
    if accelerator.check_trigger():
        break

< > 在 GitHub 上更新