From 5da6ad27314cc1d3ebd43d7d9dcb373413c93048 Mon Sep 17 00:00:00 2001 From: bigmoyan Date: Sat, 26 Jul 2025 02:03:48 +0800 Subject: [PATCH] fix break for ckpt without _tp_plan (#39658) * fix break for ckpt without _tp_plan * Update src/transformers/modeling_utils.py * Update src/transformers/modeling_utils.py --------- Co-authored-by: wangzhengtao Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> --- src/transformers/modeling_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 374b308d19..d68f5f6dfe 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -5991,8 +5991,9 @@ def caching_allocator_warmup(model: PreTrainedModel, expanded_device_map: dict, if not accelerator_device_map: return + tp_plan = getattr(model, "_tp_plan", []) or [] tp_plan_regex = ( - re.compile("|".join([re.escape(plan) for plan in model._tp_plan])) + re.compile("|".join([re.escape(plan) for plan in tp_plan])) if _torch_distributed_available and torch.distributed.is_initialized() else None )