From 9cabb1cdde18ced3946ba6f2b6efeb797f3847ac Mon Sep 17 00:00:00 2001 From: Xiao-Chenguang <973990275@qq.com> Date: Thu, 27 Nov 2025 09:55:41 +0800 Subject: [PATCH 1/2] Fix mixed torch.Tensor and DTensor error by registering fsdp forward for trainer.mode.generate --- src/transformers/trainer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 65331ba8322c..86b7c8feb4cf 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2338,6 +2338,10 @@ def _inner_training_loop( if self.is_fsdp_enabled: self.model = self.model_wrapped = model + + # Fix `got mixed torch.Tensor and DTensor` error in model.generate() for FSDP2 with LoRA + if is_fsdp2: + dist.fsdp.register_fsdp_forward_method(self.model, "generate") # for the rest of this function `model` is the outside model, whether it was wrapped or not if model is not self.model: From ee9152e7b71144d13c9d08db7af8e57bce1ffbdc Mon Sep 17 00:00:00 2001 From: yiminzme Date: Wed, 3 Dec 2025 09:59:47 +0800 Subject: [PATCH 2/2] Apply fsdp forward register when fsdp is enabled --- src/transformers/trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index 86b7c8feb4cf..e5f95daeee70 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2338,9 +2338,7 @@ def _inner_training_loop( if self.is_fsdp_enabled: self.model = self.model_wrapped = model - - # Fix `got mixed torch.Tensor and DTensor` error in model.generate() for FSDP2 with LoRA - if is_fsdp2: + # Fix `got mixed torch.Tensor and DTensor` error in model.generate() for FSDP2 with LoRA dist.fsdp.register_fsdp_forward_method(self.model, "generate") # for the rest of this function `model` is the outside model, whether it was wrapped or not