【HuggingFace LLM】经典NLP微调任务之问答

正文

数据集准备

python 复制代码
inputs = tokenizer(
    question,
    context,
    max_length=100,
    truncation="only_second",
    stride=50,
    return_overflowing_tokens=True,
)

for ids in inputs["input_ids"]:
    print(tokenizer.decode(ids))
    
'[CLS] To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France? [SEP] Architecturally, the school has a Catholic character. Atop the Main Building\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend " Venite Ad Me Omnes ". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basi [SEP]'
'[CLS] To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France? [SEP] the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend " Venite Ad Me Omnes ". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin [SEP]'
'[CLS] To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France? [SEP] Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive ( and in a direct line that connects through 3 [SEP]'
'[CLS] To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France? [SEP]. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive ( and in a direct line that connects through 3 statues and the Gold Dome ), is a simple, modern stone statue of Mary. [SEP]'

由于answer不一定出现在context哪个位置,因此不可以像之前处理的,直接truncation='True'后,把超过max_length部分给切割掉。因此这里采用stride的策略滑窗式生成样本

!NOTE #重要

经过tokenizer后,会得到如下几个结果:

  1. input_ids:[CLS] question [SEP] context [SEP]形式token列表,batch输入为多列表;
  2. token_type_ids:[None] 0 0 0 [None] 1 1 1 [None]的形式,区分token是来自question/context/special tokens,在QA任务中多用来标记Q/A的区分界限
  3. attention_mask:用于标记哪些是[PAD]
  4. offset_mapping:返回类似(0, 3)的格式,分别对应 questioncontext原始字符串中的切片,特殊字符的offset(0, 0),在QA任务中用于确定answer是否在该样本中 ,并且用于tokens -> strings的切换
  5. overflow_to_sample_mapping:返回[0, 0, 0, 0, 1, 1]格式,用于标记返回的input_ids是属于哪个question/contextQA任务中多用于对应answers

数据集制作

python 复制代码
# HuggingFace官方案例
answers = raw_datasets["train"][2:6]["answers"]
start_positions = []
end_positions = []

for i, offset in enumerate(inputs["offset_mapping"]):
    # 找到指定答案
    sample_idx = inputs["overflow_to_sample_mapping"][i]
    answer = answers[sample_idx]

    # 定位答案在原context中的字符级始末位置
    start_char = answer["answer_start"][0]
    end_char = answer["answer_start"][0] + len(answer["text"][0])
    sequence_ids = inputs.sequence_ids(i) # 返回类似[None,0,0..,None,1,1]的列表,与token_type_ids作用类似

    # 寻找在input_ids中的[cls] question [sep] context [sep],context起始位
    idx = 0
    while sequence_ids[idx] != 1:
        idx += 1
    context_start = idx
    while sequence_ids[idx] == 1:
        idx += 1
    context_end = idx - 1

    # 只有当答案的始末都在当前样本的context中时,才给打标签
    if offset[context_start][0] > start_char or offset[context_end][1] < end_char:
        start_positions.append(0)
        end_positions.append(0)
    else:
        # 在context范围中,定位answer左侧的token级位置
        idx = context_start
        while idx <= context_end and offset[idx][0] <= start_char:
            idx += 1
        start_positions.append(idx - 1)

        # 在context范围中,定位answer右侧的token级位置
        idx = context_end
        while idx >= context_start and offset[idx][1] >= end_char:
            idx -= 1
        end_positions.append(idx + 1)

start_positions, end_positions
python 复制代码
# 根据理解自己编写
answers = raw_datasets["train"][2:6]["answers"]
start_positions = []
end_positions = []

for idx, item in enumerate(inputs["overflow_to_sample_mapping"]):
  tokens_length = len(inputs["input_ids"][idx])

  answer_for_idx = answers[item]
  answer_start = answer_for_idx["answer_start"][0]
  answer_end = answer_start + len(answer_for_idx["text"][0])

  ## find question start & end
  question_start = 0
  question_end = len([item for item in inputs["token_type_ids"][idx] if item == 0])

  context_start, _ = inputs["offset_mapping"][idx][question_end]
  _, context_end = inputs["offset_mapping"][idx][tokens_length-2]

  if (context_start <= answer_start) and (context_end >= answer_end):
    # 这里标签需要打的是input_ids[idx]中对应的start和end
    for i, (token_start, token_end) in enumerate(inputs["offset_mapping"][idx]):
      if token_start == answer_start: start_positions.append(i)
      if token_end == answer_end: end_positions.append(i)
    # start_positions.append(answer_start)
    # end_positions.append(answer_end)
  else:
    start_positions.append(0)
    end_positions.append(0)

print(list(zip(start_positions, end_positions)))

和官方案例中的区别在于:

  1. 官方案例使用while循环遍历式的去找到0/None/1的区分界限,我直接默认0在前面使用len(全0列表)去定位question/answer的分界;并且我默认遵循[CLS] question [SEP] context [SEP]的形式,因此我的代码不大具有普适性。
  2. 官方案例中,answer对应的token位置也是通过while进行遍历查找得到的,当遍历token,对应的offset右边界超过answer的字符级context下标(a, b)左侧a时,标注该下标+1为对应answertoken;反之从右往左进行遍历。
python 复制代码
max_length = 384
stride = 128


def preprocess_training_examples(examples):
    questions = [q.strip() for q in examples["question"]]
    inputs = tokenizer(
        questions,
        examples["context"],
        max_length=max_length,
        truncation="only_second",
        stride=stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    offset_mapping = inputs.pop("offset_mapping")
    sample_map = inputs.pop("overflow_to_sample_mapping")
    answers = examples["answers"]
    start_positions = []
    end_positions = []

    for i, offset in enumerate(offset_mapping):
        sample_idx = sample_map[i]
        answer = answers[sample_idx]
        start_char = answer["answer_start"][0]
        end_char = answer["answer_start"][0] + len(answer["text"][0])
        sequence_ids = inputs.sequence_ids(i)

        # Find the start and end of the context
        idx = 0
        while sequence_ids[idx] != 1:
            idx += 1
        context_start = idx
        while sequence_ids[idx] == 1:
            idx += 1
        context_end = idx - 1

        # If the answer is not fully inside the context, label is (0, 0)
        if offset[context_start][0] > start_char or offset[context_end][1] < end_char:
            start_positions.append(0)
            end_positions.append(0)
        else:
            # Otherwise it's the start and end token positions
            idx = context_start
            while idx <= context_end and offset[idx][0] <= start_char:
                idx += 1
            start_positions.append(idx - 1)

            idx = context_end
            while idx >= context_start and offset[idx][1] >= end_char:
                idx -= 1
            end_positions.append(idx + 1)

    inputs["start_positions"] = start_positions
    inputs["end_positions"] = end_positions
    return inputs

train_dataset = raw_datasets["train"].map(
    preprocess_training_examples,
    batched=True,
    remove_columns=raw_datasets["train"].column_names,
)
>>>len(raw_datasets["train"]), len(train_dataset)
>>>(87599, 88729)

一般在做数据集的时候,都可以按照上述流程,先取一个不大的小样本列表 ,经过一个流程后得到最终的start_postionend_position字段,然后再做一个函数,直接使用batch加速跑。

课后task

轮到你了! 使用 XLNet 架构时,会在左侧应用填充,并且问题和上下文会互换位置。请将我们刚才看到的所有代码适配到 XLNet 架构(并添加 padding=True )。请注意,应用填充后, [CLS] 标记可能不会位于第 0 个位置。

python 复制代码
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")

# 核心代码:将填充方向改为左侧
tokenizer.padding_side = "left"

inputs = tokenizer(
    raw_datasets["train"][2:6]["context"],
    raw_datasets["train"][2:6]["question"],
    max_length=100,
    truncation="only_first",
    stride=50,
    padding=True,
    return_overflowing_tokens=True,
    return_offsets_mapping=True,
)

print(f"The 4 examples gave {len(inputs['input_ids'])} features.")
print(f"Here is where each comes from: {inputs['overflow_to_sample_mapping']}.")

for ids in inputs["input_ids"]:
    print(tokenizer.decode(ids))

answers = raw_datasets["train"][2:6]["answers"]
for answer in answers:
  answer_start, answer_end = answer["answer_start"][0], answer["answer_start"][0] + len(answer['text'][0])
  print(answer_start, answer_end)

answers = raw_datasets["train"][2:6]["answers"]
start_list, end_list = [], []

for idx, offset in enumerate(inputs["offset_mapping"]):
  # 答案字符级定位
  answer_idx = inputs["overflow_to_sample_mapping"][idx]
  answer_for_idx = answers[answer_idx]
  answer_start, answer_end = answer_for_idx["answer_start"][0], answer_for_idx["answer_start"][0] + len(answer_for_idx['text'][0])

  # 寻找context与question的区分位置,[pad] context [sep] question [sep] [cls]
  mask_list = inputs.sequence_ids(idx)
  flag = 0
  while mask_list[flag] == None:
    flag += 1
  context_token_start = flag
  while mask_list[flag] != None:
    flag += 1
  context_token_end = flag - 1

  # 寻找[cls]在input_ids中的位置
  cls_index = inputs["input_ids"][idx].index(tokenizer.cls_token_id)

  # 寻找answer在context中的具体位置
  if (answer_start >= offset[context_token_start][0]) and (answer_end <= offset[context_token_end][1]):
    context_flag = context_token_start
    while context_flag <= context_token_end and offset[context_flag][1] <= answer_start:
      context_flag += 1

    answer_token_start = context_flag
    
    context_flag = context_token_end
    while context_flag >= context_token_start and offset[context_flag][0] >= answer_end:
      context_flag -= 1

    anaswer_token_end = context_flag

    start_list.append(answer_token_start)
    end_list.append(anaswer_token_end)
  else:
    start_list.append(cls_index)
    end_list.append(cls_index)

print(list(zip(start_list, end_list)))

验证集数据处理

由于需要把一个question/context对拆分为多个样本,并生成一个/多个answer,因此我们需要一个映射关系,好让我们知道这个answer是回答哪一个question的,因此会添加一个id列。

python 复制代码
def preprocess_validation_examples(examples):
	# 格式化question
    questions = [q.strip() for q in examples["question"]]
    # 这里padding是因为需要输入模型
    inputs = tokenizer(
        questions,
        examples["context"],
        max_length=max_length,
        truncation="only_second",
        stride=stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    sample_map = inputs.pop("overflow_to_sample_mapping")
    # 这里在定义question列表
    example_ids = []

    for i in range(len(inputs["input_ids"])):
	    # 例如一个question/context对生成了3个样本,那么sample_idx都是相同的例如8(表示第8个question/context对),那么对应的example["id"]也是3个相同的id号。
        sample_idx = sample_map[i]
        example_ids.append(examples["id"][sample_idx])

        sequence_ids = inputs.sequence_ids(i)
        offset = inputs["offset_mapping"][i]
        # 方便后续生成概率,这里直接把question和speicial tokens的token_type设为None,这样如果生成概率在其中的话直接丢弃即可
        inputs["offset_mapping"][i] = [
            o if sequence_ids[k] == 1 else None for k, o in enumerate(offset)
        ]

    inputs["example_id"] = example_ids
    return inputs

validation_dataset = raw_datasets["validation"].map(
     preprocess_validation_examples,
     batched=True,
     remove_columns=raw_datasets["validation"].column_names,
)
>>>len(raw_datasets["validation"]), len(validation_dataset)
>>>(10570, 10822)

使用API进行微调

我们已对上下文之外的标记对应的起始和结束 logits 进行了屏蔽,然后使用 softmax 将起始和结束 logits 转换为概率。通过计算对应两个概率的乘积,给每个 (start_token, end_token) 对赋予一个分数。最后寻找得分最高的且能得出有效答案的配对 (例如, start_token 小于 end_token )。

由于logits经过softmax后大小关系仍保持不变,因此可以省略softmax直接通过比较

log⁡(ab)=log⁡(a)+log⁡(b) \log(ab)=\log(a)+\log(b) log(ab)=log(a)+log(b)

两个logits相加对比大小即可。

python 复制代码
small_eval_set = raw_datasets["validation"].select(range(100))
trained_checkpoint = "distilbert-base-cased-distilled-squad"

tokenizer = AutoTokenizer.from_pretrained(trained_checkpoint)
eval_set = small_eval_set.map(
    preprocess_validation_examples,
    batched=True,
    remove_columns=raw_datasets["validation"].column_names,
)

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

import torch
from transformers import AutoModelForQuestionAnswering

eval_set_for_model = eval_set.remove_columns(["example_id", "offset_mapping"])
eval_set_for_model.set_format("torch")

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
batch = {k: eval_set_for_model[k].to(device) for k in eval_set_for_model.column_names}

trained_model = AutoModelForQuestionAnswering.from_pretrained(trained_checkpoint).to(
    device
)

with torch.no_grad():
    outputs = trained_model(**batch)
  1. 为编写后处理函数,验证可行性,这里用了一个已训练完成的distilbert-base-cased-distilled-squad模型进行结果生成;
  2. 此外,可以看到输入模型的数据格式要求是要一个dict字典形式
后处理
python 复制代码
# 获取两个列表,每个token的answer起点概率,每个token的answer终点概率
start_logits = outputs.start_logits.cpu().numpy()
end_logits = outputs.end_logits.cpu().numpy()

import collections

# eval_set中是多个inputs组成的列表
example_to_features = collections.defaultdict(list)
for idx, feature in enumerate(eval_set):
	# 这一步是在找id对应输出答案的映射关系,可以有一个id对应多个inputs
    example_to_features[feature["example_id"]].append(idx)

import numpy as np

n_best = 20
max_answer_length = 30
predicted_answers = []

# 这里的example就是id、question、context的字典
for example in small_eval_set:
    example_id = example["id"]
    context = example["context"]
    answers = []

	# 一个id可以对应多个inputs输入,对应模型生成多个结果
    for feature_index in example_to_features[example_id]:
        start_logit = start_logits[feature_index] # 1*100
        end_logit = end_logits[feature_index] # 1*100
        offsets = eval_set["offset_mapping"][feature_index] # 对应样本的(部分)context下标索引

        start_indexes = np.argsort(start_logit)[-1 : -n_best - 1 : -1].tolist()
        end_indexes = np.argsort(end_logit)[-1 : -n_best - 1 : -1].tolist()
        for start_index in start_indexes:
            for end_index in end_indexes:
                # Skip answers that are not fully in the context
                if offsets[start_index] is None or offsets[end_index] is None:
                    continue
                # Skip answers with a length that is either < 0 or > max_answer_length.
                if (
                    end_index < start_index
                    or end_index - start_index + 1 > max_answer_length
                ):
                    continue

                answers.append(
                    {
                        "text": context[offsets[start_index][0] : offsets[end_index][1]],
                        "logit_score": start_logit[start_index] + end_logit[end_index],
                    }
                )

    best_answer = max(answers, key=lambda x: x["logit_score"])
    predicted_answers.append({"id": example_id, "prediction_text": best_answer["text"]})
  1. np中使用[-1: -n_best-1: -1]是指,倒数n_best个数的意思 ,由此也可以知道np.argsort默认是由小到大进行排序,因此这里倒数取n_best选取最大概率的start_end对;
指标计算
python 复制代码
import evaluate

metric = evaluate.load("squad")
theoretical_answers = [
    {"id": ex["id"], "answers": ex["answers"]} for ex in small_eval_set
]

print(predicted_answers[0])
print(theoretical_answers[0])

{'id': '56be4db0acb8001400a502ec', 'prediction_text': 'Denver Broncos'}
{'id': '56be4db0acb8001400a502ec', 'answers': {'text': ['Denver Broncos', 'Denver Broncos', 'Denver Broncos'], 'answer_start': [177, 177, 177]}}

>>>metric.compute(predictions=predicted_answers, references=theoretical_answers)
>>>{'exact_match': 83.0, 'f1': 88.25}

!NOTE

SQuAD 评估指标(evaluate.load("squad")原生支持多参考答案场景(这是为了适配 SQuAD 数据集本身的标注特点:同一个问题可能有多个标注者给出的合法正确答案),无需你额外处理,metric 内部会自动按「最优匹配」原则计算,具体逻辑如下:

  1. Exact Match(EM,精确匹配)
    • 计算规则:对单个样本,只要模型的prediction_text与参考答案列表(answers["text"])中的任意一个 答案完全一致(大小写、空格、标点、特殊字符等严格匹配),该样本的 EM 值为1;否则为0
    • 最终 EM 分数:所有样本的 EM 值取算术平均(比如示例中 83.0 代表 83% 的样本至少匹配一个参考答案);
    • 你的示例中:参考答案是 3 个相同的 'Denver Broncos',预测文本恰好匹配,因此该样本 EM=1;若参考答案是['Denver Broncos', 'Broncos'],预测文本是Broncos,该样本 EM 也为 1;
  2. F1 Score(F1 值)
    F1 是基于「词级别的重叠度」计算的,处理多参考答案时更灵活:
    • 计算规则:对单个样本,先将预测文本和每个参考答案 分别拆分为单词列表;计算预测文本与「每个参考答案」的 F1 值(精确率 + 召回率的调和平均);取这些 F1 值中的最大值作为该样本的最终 F1 值;
    • 最终 F1 分数:所有样本的 F1 最大值取算术平均(示例中 88.25 是所有样本最优 F1 的平均值)。
python 复制代码
from tqdm.auto import tqdm

def compute_metrics(start_logits, end_logits, features, examples):
    example_to_features = collections.defaultdict(list)
    for idx, feature in enumerate(features):
        example_to_features[feature["example_id"]].append(idx)

    predicted_answers = []
    for example in tqdm(examples):
        example_id = example["id"]
        context = example["context"]
        answers = []

        # Loop through all features associated with that example
        for feature_index in example_to_features[example_id]:
            start_logit = start_logits[feature_index]
            end_logit = end_logits[feature_index]
            offsets = features[feature_index]["offset_mapping"]

            start_indexes = np.argsort(start_logit)[-1 : -n_best - 1 : -1].tolist()
            end_indexes = np.argsort(end_logit)[-1 : -n_best - 1 : -1].tolist()
            for start_index in start_indexes:
                for end_index in end_indexes:
                    # Skip answers that are not fully in the context
                    if offsets[start_index] is None or offsets[end_index] is None:
                        continue
                    # Skip answers with a length that is either < 0 or > max_answer_length
                    if (
                        end_index < start_index
                        or end_index - start_index + 1 > max_answer_length
                    ):
                        continue

                    answer = {
                        "text": context[offsets[start_index][0] : offsets[end_index][1]],
                        "logit_score": start_logit[start_index] + end_logit[end_index],
                    }
                    answers.append(answer)

        # Select the answer with the best score
        if len(answers) > 0:
            best_answer = max(answers, key=lambda x: x["logit_score"])
            predicted_answers.append(
                {"id": example_id, "prediction_text": best_answer["text"]}
            )
        else:
            predicted_answers.append({"id": example_id, "prediction_text": ""})

    theoretical_answers = [{"id": ex["id"], "answers": ex["answers"]} for ex in examples]
    return metric.compute(predictions=predicted_answers, references=theoretical_answers)
开始微调
python 复制代码
model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)
from huggingface_hub import notebook_login

notebook_login()

from transformers import TrainingArguments

args = TrainingArguments(
    "bert-finetuned-squad",
    evaluation_strategy="no",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    weight_decay=0.01,
    fp16=True,
    push_to_hub=True,
)

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=validation_dataset,
    tokenizer=tokenizer,
)
trainer.train()

predictions, _, _ = trainer.predict(validation_dataset)
start_logits, end_logits = predictions
compute_metrics(start_logits, end_logits, validation_dataset, raw_datasets["validation"])

使用Accelerate进行微调

python 复制代码
from torch.utils.data import DataLoader
from transformers import default_data_collator

train_dataset.set_format("torch")
validation_set = validation_dataset.remove_columns(["example_id", "offset_mapping"])
validation_set.set_format("torch")

train_dataloader = DataLoader(
    train_dataset,
    shuffle=True,
    collate_fn=default_data_collator,
    batch_size=8,
)
eval_dataloader = DataLoader(
    validation_set, collate_fn=default_data_collator, batch_size=8
)

model = AutoModelForQuestionAnswering.from_pretrained(model_checkpoint)
from torch.optim import AdamW

optimizer = AdamW(model.parameters(), lr=2e-5)

from accelerate import Accelerator

accelerator = Accelerator(fp16=True)
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

from transformers import get_scheduler

num_train_epochs = 3
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps,
)

from huggingface_hub import Repository, get_full_repo_name

model_name = "bert-finetuned-squad-accelerate"
repo_name = get_full_repo_name(model_name)
output_dir = "bert-finetuned-squad-accelerate"
repo = Repository(output_dir, clone_from=repo_name)

from tqdm.auto import tqdm
import torch

progress_bar = tqdm(range(num_training_steps))

for epoch in range(num_train_epochs):
    # Training
    model.train()
    for step, batch in enumerate(train_dataloader):
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

    # Evaluation
    model.eval()
    start_logits = []
    end_logits = []
    accelerator.print("Evaluation!")
    for batch in tqdm(eval_dataloader):
        with torch.no_grad():
            outputs = model(**batch)

        start_logits.append(accelerator.gather(outputs.start_logits).cpu().numpy())
        end_logits.append(accelerator.gather(outputs.end_logits).cpu().numpy())

    start_logits = np.concatenate(start_logits)
    end_logits = np.concatenate(end_logits)
    start_logits = start_logits[: len(validation_dataset)]
    end_logits = end_logits[: len(validation_dataset)]

    metrics = compute_metrics(
        start_logits, end_logits, validation_dataset, raw_datasets["validation"]
    )
    print(f"epoch {epoch}:", metrics)

    # Save and upload
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
    if accelerator.is_main_process:
        tokenizer.save_pretrained(output_dir)
        repo.push_to_hub(
            commit_message=f"Training in progress epoch {epoch}", blocking=False
        )
  1. collate_fn为 Transformers 提供的 default_data_collator
  2. 由于在Accelerate中,不同进程之前会存在补齐现象,因此需要使用截断防止输入pad;
  3. 使用gather把所有进程中的tensor进行聚合(放在不同的进程中跑);

!NOTE accelerate.gather

假设有一个长为18的列表 分5个给进程1;分5个给进程2;5个给进程3;剩3个给进程4(自动补齐至5) 后续计算出结果后,使用gather把结果进行聚合,但是需要排除自动补齐的4号进程后2个?

其中「自动补齐」不是gather的行为,而是数据分发阶段的操作,本质是多进程(多 GPU)分布式训练中,数据按进程数拆分时的「padding 补齐」:为了保证所有进程张量形状一致(避免分布式计算报错),最后一个进程会被自动补齐(padding)到其他进程相同长度(5);

accelerate.gather的核心是跨进程收集张量数据,把分散在不同进程(GPU)上的张量,收集到指定的主进程(默认是进程 0);按进程顺序拼接 成一个完整的张量返回;但不会识别 / 排除补齐的 padding 元素 ------ 它会把进程 4 的 5 个元素(3 个有效 + 2 个补齐)全部聚合,因此需要你手动排除补齐部分。

问题

总结