SLURM批量运行🤗Transformer微调脚本的分布式资源分配问题
2026-7-15
问题描述
我有一个基于Hugging Face官方教程的🤗Transformer微调脚本，在远程SLURM服务器上交互式运行时可正常生成预期输出，但提交为批量作业时出现资源分配问题。错误源于TrainingArgs实例化阶段，经排查，问题出在accelerate.state第198行的代码：
elif get_int_from_env(["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"], 1) > 1:
交互式运行时所有环境变量值均为1，该行判断为False；批量作业时PMI_SIZE值为128，触发创建TCPStore对象，最终因登录匹配失败报错。我曾尝试将MASTER_ADDR和MASTER_PORT设置为登录服务器的地址与端口，却导致连接超时；尝试通过sagemaker和deepspeed运行TrainingArguments也出现相同问题。
报错堆栈信息：
╭───────────────────── Traceback (most recent call last) ──────────────────────╮ │ /users/home/korat/acode/finetune_classification.py:170 in <module> │ │ │ │ 167 │ │ df.to_csv(os.path.join(output_dir_base,f'results-{label_col}{a │ │ 168 │ │ 169 if __name__ == "__main__": │ │ ❱ 170 │ main() │ │ │ │ /users/home/korat/acode/finetune_classification.py:155 in main │ │ │ │ 152 │ │ i+=1 │ │ 153 │ │ if i>3: │ │ 154 │ │ │ break │ │ ❱ 155 │ │ loss = train_model(model_checkpoint, dataset, weight_decay=wei │ │ 156 │ │ │ │ │ │ │ │ │ adam_beta1=adam_beta1, │ │ 157 │ │ │ │ │ │ │ │ │ adam_beta2=adam_beta2, │ │ 158 │ │ │ │ │ │ │ │ │ adam_epsilon=adam_epsilon, │ │ │ │ /users/home/korat/acode/finetune_classification.py:86 in train_model │ │ │ │ 83 │ output_dir = output_dir_base + "/" + hyparam_comb_to_str(params) │ │ 84 │ │ │ 85 │ #get_int_from_env(["PMI_SIZE"], 1) = 128, and this is where the co │ │ ❱ 86 │ training_args = TrainingArguments( │ │ 87 │ │ │ │ 88 │ │ output_dir=output_dir, │ │ 89 │ │ evaluation_strategy=IntervalStrategy.STEPS, │ │ <string>:111 in __init__ │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/transformers/training_ │ │ args.py:1340 in __post_init__ │ │ │ │ 1337 │ │ if ( │ │ 1338 │ │ │ self.framework == "pt" │ │ 1339 │ │ │ and is_torch_available() │ │ ❱ 1340 │ │ │ and (self.device.type != "cuda") │ │ 1341 │ │ │ and (get_xla_device_type(self.device) != "GPU") │ │ 1342 │ │ │ and (self.fp16 or self.fp16_full_eval) │ │ 1343 │ │ ): │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/transformers/training_ │ │ args.py:1764 in device │ │ │ │ 1761 │ │ The device used by this process. │ │ 1762 │ │ """ │ │ 1763 │ │ requires_backends(self, ["torch"]) │ │ ❱ 1764 │ │ return self._setup_devices │ │ 1765 │ │ │ 1766 │ @property │ │ 1767 │ def n_gpu(self): │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/transformers/utils/gen │ │ eric.py:54 in __get__ │ │ │ │ 51 │ │ attr = "__cached_" + self.fget.__name__ │ │ 52 │ │ cached = getattr(obj, attr, None) │ │ 53 │ │ if cached is None: │ │ ❱ 54 │ │ │ cached = self.fget(obj) │ │ 55 │ │ │ setattr(obj, attr, cached) │ │ 56 │ │ return cached │ │ 57 │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/transformers/training_ │ │ args.py:1695 in _setup_devices │ │ │ │ 1692 │ │ │ del os.environ["ACCELERATE_USE_DEEPSPEED"] │ │ 1693 │ │ │ self._n_gpu = 1 │ │ 1694 │ │ else: │ │ ❱ 1695 │ │ │ self.distributed_state = PartialState(backend=self.ddp_ba │ │ 1696 │ │ │ self._n_gpu = 1 │ │ 1697 │ │ if not is_sagemaker_mp_enabled(): │ │ 1698 │ │ │ device = self.distributed_state.device │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/accelerate/state.py:23 │ │ 8 in __init__ │ │ │ │ 235 │ │ │ │ │ # Backend is not set by the user, we set it here │ │ 236 │ │ │ │ │ kwargs.pop("backend", None) │ │ 237 │ │ │ │ │ self.backend = backend │ │ ❱ 238 │ │ │ │ │ torch.distributed.init_process_group(self.backend, │ │ 239 │ │ │ │ self.num_processes = torch.distributed.get_world_size( │ │ 240 │ │ │ │ self.process_index = torch.distributed.get_rank() │ │ 241 │ │ │ │ self.local_process_index = local_rank │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/torch/distributed/dist │ │ ributed_c10d.py:900 in init_process_group │ │ │ │ 897 │ │ │ rendezvous_iterator = rendezvous( │ │ 898 │ │ │ │ init_method, rank, world_size, timeout=timeout │ │ 899 │ │ │ ) │ │ ❱ 900 │ │ │ store, rank, world_size = next(rendezvous_iterator) │ │ 901 │ │ │ store.set_timeout(timeout) │ │ 902 │ │ │ │ │ 903 │ │ │ # Use a PrefixStore to avoid accidental overrides of keys │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/torch/distributed/rend │ │ ezvous.py:245 in _env_rendezvous_handler │ │ │ │ 242 │ master_addr = _get_env_or_raise("MASTER_ADDR") │ │ 243 │ master_port = int(_get_env_or_raise("MASTER_PORT")) │ │ 244 │ │ │ ❱ 245 │ store = _create_c10d_store(master_addr, master_port, rank, world_s │ │ 246 │ │ │ 247 │ yield (store, rank, world_size) │ │ 248 │ │ │ │ /users/home/korat/.local/lib/python3.10/site-packages/torch/distributed/rend │ │ ezvous.py:176 in _create_c10d_store │ │ │ │ 173 │ │ return PrefixStore(f"/worker/attempt_{attempt}", tcp_store) │ │ 174 │ else: │ │ 175 │ │ start_daemon = rank == 0 │ │ ❱ 176 │ │ return TCPStore( │ │ 177 │ │ │ hostname, port, world_size, start_daemon, timeout, multi_t │ │ 178 │ │ ) │ │ 179 │ ╰──────────────────────────────────────────────────────────────────────────────╯ RuntimeError: Connection reset by peer
请问我应该尝试找到TCPStore的正确参数以成功登录服务器，还是寻找无需登录即可分配资源的解决方案？
解决方案
优先选择无需触发分布式初始化的方案，这比调试TCPStore参数更直接高效，具体操作如下：
1. 手动覆盖分布式环境变量
在实例化TrainingArguments前，强制设置相关环境变量为单进程模式：
import os # 覆盖SLURM注入的多进程环境变量 os.environ["PMI_SIZE"] = "1" os.environ["WORLD_SIZE"] = "1" os.environ["OMPI_COMM_WORLD_SIZE"] = "1"
2. 显式指定单进程参数
给TrainingArguments添加以下参数，直接禁用分布式逻辑：
training_args = TrainingArguments( # 你的原有参数... local_rank=-1, # 显式设置为-1关闭分布式检测 distributed_state=None, )
3. 调整SLURM作业脚本资源配置
确保批量作业只请求单节点单进程资源，避免SLURM自动注入多进程环境变量：
#!/bin/bash #SBATCH --nodes=1 #SBATCH --ntasks-per-node=1 #SBATCH --gres=gpu:1 # 根据需求调整GPU数量 #SBATCH --cpus-per-task=8 python finetune_classification.py
为什么不优先调试TCPStore参数？
- SLURM集群通常有严格的网络规则，登录节点和计算节点之间可能无法直接通过指定端口通信，导致
MASTER_ADDR/MASTER_PORT设置无效。 - 你的脚本在交互式模式下单进程运行正常，说明你不需要分布式训练，完全没必要额外适配多进程逻辑。
内容的提问来源于stack exchange，提问作者munchanr
相关产品推荐
相关产品推荐
