accelerate 分布式技巧（一）

accelerate分布式技巧

简单使用

Accelerate是一个来自Hugging Face的库，它简化了将单个GPU的PyTorch代码转换为单个或多台机器上的多个GPU的代码。
Accelerate精确地抽象了与多GPU/TPU/fp16相关的模板代码，并保持Pytorch其余代码不变。

  import torch
  import torch.nn.functional as F
  from datasets import load_dataset
+ from accelerate import Accelerator

+ accelerator = Accelerator()
- device = 'cpu'
+ device = accelerator.device

  model = torch.nn.Transformer().to(device)
  optimizer = torch.optim.Adam(model.parameters())

  dataset = load_dataset('my_dataset')
  data = torch.utils.data.DataLoader(dataset, shuffle=True)

+ model, optimizer, data = accelerator.prepare(model, optimizer, data)

  model.train()
  for epoch in range(10):
      for source, targets in data:
          source = source.to(device)
          targets = targets.to(device)

          optimizer.zero_grad()

          output = model(source)
          loss = F.cross_entropy(output, targets)

-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()
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

通过添加上面5行代码，可以在任何类型的单节点或分布式节点设置（单CPU、单GPU、多GPU和TPU）上运行，也可以使用或不使用混合精度（fp16）如：Accelerator(cpu=args.cpu, mixed_precision=args.mixed_precision)。

+ eval_dataloader = accelerator.prepare(eval_dataloader)
  predictions, labels = [], []
  for source, targets in eval_dataloader:
      with torch.no_grad():
          output = model(source)

-     predictions.append(output.cpu().numpy())
-     labels.append(targets.cpu().numpy())
+     predictions.append(accelerator.gather(output).cpu().numpy())
+     labels.append(accelerator.gather(targets).cpu().numpy())

  predictions = np.concatenate(predictions)
  labels = np.concatenate(labels)

+ predictions = predictions[:len(eval_dataloader.dataset)]
+ labels = label[:len(eval_dataloader.dataset)]

  metric_compute(predictions, labels)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

config参数

您需要配置 Accelerate以了解当前系统是如何为训练设置的。为此，运行以下命令并回答提示给您的问题：
accelerate config --config_file xxx

## accelerate config --help
options:
  -h, --help            show this help message and exit
  --config_file CONFIG_FILE
                        The path to use to store the config file. Will default to a file named
                        default_config.yaml in the cache location, which is the content of the environment
                        `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment
                        variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`)
                        suffixed with 'huggingface'.

subcommands:
  {default,update}
    default             Create a default config file for Accelerate with only a few flags set.
    update              Update an existing config file with the latest defaults while maintaining the old
                        configuration.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

要编写不包括DeepSpeed配置或在tpu上运行等选项的裸机配置，您可以快速运行：
python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"

得到default_config.yaml文件：

{
  "compute_environment": "LOCAL_MACHINE",
  "debug": false,
  "distributed_type": "MULTI_GPU",
  "downcast_bf16": false,
  "machine_rank": 0,
  "main_training_function": "main",
  "mixed_precision": "fp16",
  "num_machines": 1,
  "num_processes": 4,
  "rdzv_backend": "static",
  "same_network": false,
  "tpu_use_cluster": false,
  "tpu_use_sudo": false,
  "use_cpu": false
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

要检查您的配置是否正常，请运行
accelerate env

- `Accelerate` version: 0.22.0
- Platform: Linux-3.10.0-1160.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.10.12
- Numpy version: 1.25.2
- PyTorch version (GPU?): 2.0.1 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 503.65 GB
- GPU type: NVIDIA A10
- `Accelerate` default config:
        Not found
1
2
3
4
5
6
7
8
9
10
11

一旦完成，您就可以通过运行来测试设置上的一切是否正常:

accelerate test --config_file path_to_config.yaml
1

accelerate launch 启动

accelerate launch --config_file default_config.yaml \
  path_to_script.py \
  --args_for_the_script
1
2
3

API 说明

导入Accelerator主类并在accelerate 对象中实例化一个, 该类的所有参数可以参考https://huggingface.co/docs/accelerate/v0.22.0/en/package_reference/accelerator#accelerate.Accelerator

from accelerate import Accelerator
accelerator = Accelerator()
1
2

将所有与训练相关的对象(优化器、模型、训练数据加载器、学习率调度器)传递给prepare()方法。这将确保一切为训练做好准备。

model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, lr_scheduler
)
1
2
3

特别是，您的训练数据加载器将在所有可用的 GPU/TPU 核心上进行分片，以便每个核心都能看到训练数据集的不同部分。此外，所有进程的随机状态将在每次迭代开始时通过数据加载器进行同步，以确保数据以相同的方式进行洗牌（如果您决定使用 shuffle=True 或任何类型的随机采样器）。

训练的实际批大小将是使用的设备数量乘以您在脚本中设置的批大小:例如，在4个gpu上进行训练，在创建训练数据加载器时，批大小为16，将以实际批大小为64进行训练。

或者，您可以在创建和初始化加速器时使用选项 split_batches=True，在这种情况下，无论您在 1、2、4 还是 64 个 GPU 上运行脚本，批处理大小将始终保持不变。

Replace the lineloss.backward()by accelerator.backward(loss).

分布式评估

validation_dataloader = accelerator.prepare(validation_dataloader)
1

至于您的训练数据加载器，这意味着(如果您在多个设备上运行脚本)每个设备将只能看到部分评估数据。这意味着您需要将您的预测组合在一起。使用gather_for_metrics()方法非常容易做到这一点。

for inputs, targets in validation_dataloader:
    predictions = model(inputs)
    # Gather all predictions and targets
    all_predictions, all_targets = accelerator.gather_for_metrics((predictions, targets))
    # Example of use with a *Datasets.Metric*
    metric.add_batch(all_predictions, all_targets)
1
2
3
4
5
6

只在一个进程上执行一条语句

if accelerator.is_local_main_process: # 本机多进程
    # Is executed once per server
1
2

另一个例子是进度条:为了避免在输出中出现多个进度条，您应该只在本地主进程中显示一个

from tqdm.auto import tqdm

progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
1
2
3

如果你在两台有多个gpu的服务器上运行你的训练，指令将在每台服务器上执行一次。如果您只需要对所有进程(而不是每台机器)执行一次，例如，将最终模型上载到模型中心，那么将其封装在这样的测试中

if accelerator.is_main_process:
    # Is executed once only
1
2

推迟执行

当您运行常规脚本时，指令是按顺序执行的。使用Accelerate同时在多个gpu上部署脚本会带来一些复杂性:虽然每个进程按顺序执行所有指令，但有些进程可能比其他进程更快。

在执行给定的指令之前，可能需要等待所有进程到达某个点。例如，在确定每个过程都通过训练完成之前，您不应该保存模型。为此，只需在代码中编写以下行代码

accelerator.wait_for_everyone()
1

该指令将阻塞所有先到达的进程，直到所有其他进程到达该点(如果您只在一个GPU或CPU上运行脚本，这将不会做任何事情)。

保存训练的模型可能需要一些调整：首先，您应该等待所有进程都达到脚本中的该点，如上所示，然后，您应该在保存模型之前解开模型。这是因为在执行prepare()方法时，您的模型可能已被放置在一个更大的模型中，该模型处理分布式训练。这反过来意味着，在不采取任何预防措施的情况下保存模型状态字典将考虑到潜在的额外层，并且最终将得到无法加载回基本模型的权重。 save_model() 方法将帮助您实现这一目标。它将解开您的模型并保存模型状态字典。

Saving/loading a model

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)
1
2

save_model()方法还可以将模型保存到分片检查点或使用安全系数格式。如下例子：

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)
1
2

如果您的脚本包含加载检查点的逻辑，我们还建议您在展开的模型中加载权重（这仅在您使模型经过prepare()后使用加载函数时才有用）。这是一个例子：

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))
1
2
3

请注意，由于所有模型参数都是对张量的引用，因此这会将您的权重加载到模型中。

如果您想将分片检查点或安全张量格式的检查点加载到具有特定设备的模型中，我们建议您使用 load_checkpoint_in_model() 函数加载它。这是一个例子：

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})
1

Saving/loading entire states

训练模型时，您可能希望保存模型、优化器、随机生成器和潜在的 LR 调度器的当前状态，以便在同一脚本中恢复。您可以分别使用 save_state() 和 load_state() 来执行此操作。

要进一步自定义通过 save_state() 保存状态的位置和方式，可以使用 ProjectConfiguration 类。例如，如果启用automatic_checkpoint_naming，则每个保存的检查点将位于Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

from accelerate.utils import ProjectConfiguration
accelerator_project_config = ProjectConfiguration(
    total_limit=args.checkpoints_total_limit
)
accelerator = Accelerator(
    gradient_accumulation_steps=args.gradient_accumulation_steps,
    mixed_precision=args.mixed_precision,
    log_with=args.report_to,
    logging_dir=logging_dir,
    project_config=accelerator_project_config,
)

1
2
3
4
5
6
7
8
9
10
11
12
13

如果您已通过 register_for_checkpointing() 注册了要存储的任何其他有状态项目，它们也将被保存和/或加载。

from accelerate import Accelerator

accelerator = Accelerator()
# Assume `CustomObject` has a `state_dict` and `load_state_dict` function.
obj = CustomObject()
accelerator.register_for_checkpointing(obj)
accelerator.save_state("checkpoint.pt")
1
2
3
4
5
6
7

传递给 register_for_checkpointing() 的每个对象都必须有一个要存储的 load_state_dict 和 state_dict 函数

梯度裁剪
如果您在脚本中使用梯度裁剪，则应分别将对 torch.nn.utils.clip_grad_norm_ 或 torch.nn.utils.clip_grad_value_ 的调用替换为clipgrad_norm() 和clipgrad_value() 。

from accelerate import Accelerator

accelerator = Accelerator(gradient_accumulation_steps=2)
dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)

for input, target in dataloader:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_func(output, target)
    accelerator.backward(loss)
    if accelerator.sync_gradients: # 在每个训练步骤中同步梯度, 目前是否在所有过程中同步梯度
        accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
    optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12
13

混合精度训练
如果您使用 Accelerate 在混合精度中运行训练，您将在模型内部计算损失（例如在 Transformer 模型中），从而获得最佳结果。模型外部的每个计算都将以全精度执行（这通常是损失计算所需的，特别是当它涉及 softmax 时）。但是，您可能希望将损失计算放入 autocast() 上下文管理器中：

from accelerate import Accelerator

accelerator = Accelerator(mixed_precision="fp16")
with accelerator.autocast():
    train()
1
2
3
4
5

混合精度训练的另一个警告是，梯度在开始时和有时在训练过程中会跳过一些更新：由于动态损失缩放策略，训练期间存在梯度溢出的点，并且损失缩放因子减小为避免在下一步中再次发生这种情况。
这意味着您可以在没有更新时更新您的学习率调度程序，这通常很好，但当您的训练数据很少或调度程序的第一个学习率值非常重要时，可能会产生影响。在这种情况下，当优化器步骤未完成时，您可以跳过学习率调度程序更新，如下所示：

if not accelerator.optimizer_step_was_skipped:
    lr_scheduler.step()
1
2

梯度累积
要执行梯度累积，请使用accumulate()并指定gradient_accumulation_steps。这还将自动确保在多设备训练时梯度同步或不同步，检查是否应该实际执行该步骤，并自动缩放损失：

accelerator = Accelerator(gradient_accumulation_steps=2)
model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

for input, label in training_dataloader:
    with accelerator.accumulate(model):
        predictions = model(input)
        loss = loss_function(predictions, label)
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
1
2
3
4
5
6
7
8
9
10
11

参数、属性和方法说明

• 设备类型distributed_type

class DistributedType(str, enum.Enum):
    # Subclassing str as well as Enum allows the `DistributedType` to be JSON-serializable out of the box.
    NO = "NO"
    MULTI_CPU = "MULTI_CPU"
    MULTI_GPU = "MULTI_GPU"
    MULTI_NPU = "MULTI_NPU"
    MULTI_XPU = "MULTI_XPU"
    DEEPSPEED = "DEEPSPEED"
    FSDP = "FSDP"
    TPU = "TPU"
    MEGATRON_LM = "MEGATRON_LM"
1
2
3
4
5
6
7
8
9
10
11

• init_trackers初始化跟踪器，

from accelerate import Accelerator
accelerator = Accelerator(log_with="tensorboard")
accelerator.init_trackers(
	project_name="my_project",
	config={"learning_rate": 0.001, "batch_size": 32},
    init_kwargs={"tensorboard": {"flush_secs": 60}},)

tensorboard_tracker = accelerator.get_tracker("tensorboard")
1
2
3
4
5
6
7
8

• end_training结束训练

在这里插入代码片
1

其他接口参考： https://huggingface.co/docs/accelerate/v0.22.0/en/package_reference/accelerator

进阶

指定GPU和最大显存

当一张显卡容不下一个模型时，我们需要用多张显卡来推理。

假如我们现在模型是一个Llama33B，那么我们推理一般需要使用66G的显存，假如我们想要使用6号和7号卡，每张卡允许使用的显存是35G。那么我们代码可以这样写：

from transformers import LlamaConfig,LlamaForCausalLM,LlamaTokenizer
from accelerate import init_empty_weights,infer_auto_device_map,load_checkpoint_in_model,dispatch_model, load_checkpoint_and_dispatch
import torch

cuda_list = '1,2'.split(',')
memory = '8GiB'
model_path = './custom_model/hf_llama_7b'
no_split_module_classes = LlamaForCausalLM._no_split_modules
print(no_split_module_classes)
max_memory = {int(cuda):memory for cuda in cuda_list}
config = LlamaConfig.from_pretrained(model_path)
with init_empty_weights():
    model = LlamaForCausalLM._from_config(config, torch_dtype=torch.float16) #加载到meta设备中，不需要耗时，不需要消耗内存和显存

device_map = infer_auto_device_map(model, max_memory=max_memory,no_split_module_classes=no_split_module_classes) #自动划分每个层的设备
load_checkpoint_in_model(model, model_path,device_map=device_map) #加载权重
model = dispatch_model(model, device_map=device_map) #并分配到具体的设备上

tokenizer = LlamaTokenizer.from_pretrained(model_path)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
torch.set_grad_enabled(False)
model.eval()
sents=['你是谁']
ids = tokenizer(sents, 
                max_length=1800,  
                padding=True, 
                truncation=True, 
                return_tensors="pt")
ids = ids.to(model.device) 
outputs = model.generate(**ids, do_sample=False)
print(outputs)
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

但是里面写的不是很齐全，如果用了infer_auto_device_map，就不需要用load_checkpoint_and_dispatch了，否则其他显卡还是会被占用少量显存，而且也不需要手动指定如何分割，挺方便的。