PyTorch TPU分布式预测结果收集及样本顺序对齐方法

现有实现存在的卡点 #

当前基于Kaggle提供的TPU v3-8环境、PyTorch栈微调BERT问答模型时，验证阶段采用ParallelLoader实现8个核心并行预测后，存在两个待解决的核心问题：

1. 缺乏跨核心回收预测结果的可靠通道，无法拿到全部8个核心上输出的start/end logits
2. 回收后的结果无法保证与原始验证集样本顺序一一匹配，无法支撑后续EM值、F1值的准确计算

现有参考实现代码 #

def _run():
    MAX_LEN = 192 # maximum text length in the batch (cannot have too high due to memory constraints)
    BATCH_SIZE = 16 # batch size (cannot have too high due to memory constraints)
    EPOCHS = 2 # number of epochs
    train_sampler = torch.utils.data.distributed.DistributedSampler(
          tokenized_datasets['train'],
          num_replicas=xm.xrt_world_size(), # tell PyTorch how many devices (TPU cores) we are using for training
          rank=xm.get_ordinal(), # tell PyTorch which device (core) we are on currently
          shuffle=True
    )
    
    train_data_loader = torch.utils.data.DataLoader(
        tokenized_datasets['train'],
        batch_size=BATCH_SIZE,
        sampler=train_sampler,
        drop_last=True,
        num_workers=0,
    )
    
    valid_sampler = torch.utils.data.distributed.DistributedSampler(
          tokenized_datasets['validation'],
          num_replicas=xm.xrt_world_size(),
          rank=xm.get_ordinal(),
          shuffle=False
    )
    
    valid_data_loader = torch.utils.data.DataLoader(
        tokenized_datasets['validation'],
        batch_size=BATCH_SIZE,
        sampler=valid_sampler,
        drop_last=False,
        num_workers=0
    )
    
    device = xm.xla_device() # device (single TPU core)
    model = model.to(device) # put model onto the TPU core
    xm.master_print('done loading model')
    xm.master_print(xm.xrt_world_size(),'as size')
    
    lr = 0.5e-5 * xm.xrt_world_size()
    optimizer = AdamW(model.parameters(), lr=lr) # define our optimizer
    
    for epoch in range(EPOCHS):
        gc.collect() 
        # use ParallelLoader (provided by PyTorch XLA) for TPU-core-specific dataloading:
        para_loader = pl.ParallelLoader(train_data_loader, [device]) 
        xm.master_print('parallel loader created... training now')
        gc.collect()
        
        call training loop:
        train_loop_fn(para_loader.per_device_loader(device), model, optimizer, device, scheduler=None)
        del para_loader
        model.eval()
        para_loader = pl.ParallelLoader(valid_data_loader, [device])
        gc.collect()

        model.eval()
        # call evaluation loop
        print("call evaluation loop")
        start_logits, end_logits = eval_loop_fn(para_loader.per_device_loader(device), model, device)

具体可落地的实现步骤 #

• 步骤1：显式记录每个核心负责的样本全局索引
  DistributedSampler在shuffle=False时，切分逻辑是固定的：rank为r的核心负责的样本索引为[r, r+8, r+16, ...]，只要在每个核心初始化完验证采样器后，显式导出当前核心拿到的全局索引序列，就能建立预测结果和原始样本的对应关系，不需要依赖隐式的切分规则。
  对应代码：

  # 初始化valid_sampler后立即调用，导出当前核心负责的样本全局索引
  valid_sampler.set_epoch(0)  # shuffle=False时epoch参数不影响切分结果，调用目的是初始化内部索引序列
  local_sample_indices = list(iter(valid_sampler))
  

• 步骤2：修改评估循环，本地绑定预测结果与对应样本索引
  每个核心在本地遍历验证loader时，逐batch收集对应的start logits、end logits，同时把和batch顺序严格对齐的本地样本索引切片一并存入本地变量，不要直接在跨核心通信前把索引和logits拆分传输。
  对应修改后的评估循环逻辑：

  def eval_loop_fn(loader, model, device, local_sample_indices):
      model.eval()
      all_local_start = []
      all_local_end = []
      all_local_indices = []
      # loader的遍历顺序和local_sample_indices的顺序严格对齐
      for batch_idx, batch in enumerate(loader):
          with torch.no_grad():
              inputs = {k: v.to(device) for k, v in batch.items()}
              outputs = model(**inputs)
              # 切片获取当前batch对应的全局样本索引
              batch_start = batch_idx * loader.batch_size
              batch_end = batch_start + len(batch["input_ids"])
              all_local_indices.extend(local_sample_indices[batch_start:batch_end])
              # 把logits搬到CPU存储，减少TPU内存占用
              all_local_start.append(outputs.start_logits.cpu())
              all_local_end.append(outputs.end_logits.cpu())
      local_start_logits = torch.cat(all_local_start, dim=0)
      local_end_logits = torch.cat(all_local_end, dim=0)
      return local_start_logits, local_end_logits, torch.tensor(all_local_indices)
  

• 步骤3：通过PyTorch XLA集合通信接口跨核心收集结果
  不要通过文件、本地缓存的方式跨核心传数据，直接调用xm.all_gather把所有核心上的张量按rank顺序拼接回收。如果验证集总样本数无法被8个核心整除，需要把每个核心的本地张量pad到统一长度（用-1填充空白位），避免通信时报shape不匹配的错误。
  对应代码：

  # 在每个核心上运行本地评估，拿到本地结果
  local_start, local_end, local_idx = eval_loop_fn(
      para_loader.per_device_loader(device),
      model,
      device,
      local_sample_indices
  )
  # 跨核心收集所有核心的张量
  gathered_start = xm.all_gather(local_start, dim=0)
  gathered_end = xm.all_gather(local_end, dim=0)
  gathered_indices = xm.all_gather(local_idx, dim=0)
  

• 步骤4：主核心上对齐顺序计算指标
  主核心拿到所有收集的结果后，按全局样本索引排序，就能得到和原始验证集顺序完全匹配的预测结果，直接传入EM、F1计算逻辑即可。
  对应代码：

  if xm.is_master_ordinal():
      # 按全局索引排序对齐顺序
      sorted_order = torch.argsort(gathered_indices)
      aligned_start = gathered_start[sorted_order].numpy()
      aligned_end = gathered_end[sorted_order].numpy()
      # 可直接将aligned_start、aligned_end与tokenized_datasets["validation"]的样本顺序匹配，计算EM与F1值
  

原有代码的冗余优化 #

• 删除连续两次调用model.eval()的重复逻辑
• 确保ParallelLoader的初始化、集合通信步骤在所有TPU核心上同步执行，不要只在主核心运行对应逻辑，否则会出现其他核心无法参与评估、结果不全的问题
• 保留验证采样器的drop_last=False配置，配合张量pad逻辑覆盖所有验证样本，避免漏算

内容的提问来源于stack exchange，提问作者佩特微