Fix omni grpo (#4469)

Co-authored-by: tastelikefeet <[email protected]>

Author: tastelikefeet

swift/trainers/rlhf_trainer/grpo_trainer.py

@@ -1199,7 +1199,10 @@ def _padding_free_output_hook(module, args, kwargs, result):
 
         if self.padding_free:
             llm_model = get_llm_model(model)
-            base_model = llm_model.model
+            if hasattr(llm_model, 'thinker'):
+                base_model = llm_model.thinker.model
+            else:
+                base_model = llm_model.model
             remove_handle1 = base_model.register_forward_pre_hook(
                 _padding_free_input_hook, with_kwargs=True, prepend=True)
             remove_handle2 = base_model.register_forward_hook(_padding_free_output_hook, with_kwargs=True, prepend=True)

swift/trainers/sequence_parallel/ulysses.py

@@ -667,7 +667,10 @@ def pre_forward_split_hook(_self, args, kwargs):
 
         llm_model = get_llm_model(model)
 
-        base_model = llm_model.model
+        if hasattr(llm_model, 'thinker'):
+            base_model = llm_model.thinker.model
+        else:
+            base_model = llm_model.model
         if hasattr(base_model, 'language_model'):
             self.causal_mask_func = base_model.language_model._update_causal_mask
         else: