clean

Author: hjh0119

swift/megatron/trainers/rollout_mixin.py

@@ -436,9 +436,9 @@ def _export_and_load_weights(self):
             llm_model = self.engine.inner_model
             patch_vllm_moe_model_weight_loader(llm_model)
             llm_model.load_weights(weight_iterator)
-            _model_config = getattr(getattr(self.engine, 'engine', None), 'model_config', None)
+            _model_config = self.engine.engine.model_config
             finish_vllm_weight_reload(
-                llm_model, model_config=_model_config, target_device=getattr(llm_model, 'device', None))
+                llm_model, model_config=_model_config, target_device=next(llm_model.parameters()).device)
         elif self.vllm_mode == 'server':
             self._load_weights_to_server_in_buckets(weight_iterator)
             self.vllm_client.process_weights_after_loading()

swift/rlhf_trainers/gkd_trainer.py

@@ -619,10 +619,7 @@ def _fetch_and_assemble_teacher_logprobs(self, chunks):
         all_raw = gather_object(local_raw)
 
         if self.accelerator.is_main_process:
-            non_thinking_prefix_ids = get_non_thinking_prefix_ids(self.template)
-            requests = [
-                build_teacher_infer_request(d, non_thinking_prefix_ids=non_thinking_prefix_ids) for d in all_raw
-            ]
+            requests = [build_teacher_infer_request(d) for d in all_raw]
             request_config = RequestConfig(prompt_logprobs=self.gkd_logits_topk, max_tokens=1, temperature=0.0)
             responses = self.teacher_client.infer(requests, request_config=request_config, use_tqdm=False)
             parsed_global = [parse_prompt_logprobs(r, topk=self.gkd_logits_topk) for r in responses]

swift/rlhf_trainers/rollout_mixin.py

@@ -813,11 +813,10 @@ def _move_full_model_to_vllm(self):
 
         # Re-run process_weights_after_loading once after ALL groups loaded
         if self.vllm_mode == 'colocate':
-            _model_config = getattr(getattr(self.engine, 'engine', None), 'model_config', None)
+            _model_config = self.engine.engine.model_config
+            llm_model = self.engine.inner_model
             finish_vllm_weight_reload(
-                self.engine.inner_model,
-                model_config=_model_config,
-                target_device=getattr(self.engine.inner_model, 'device', None))
+                llm_model, model_config=_model_config, target_device=next(llm_model.parameters()).device)
         elif self.vllm_mode == 'server' and self.accelerator.is_main_process:
             self.vllm_client.process_weights_after_loading()
 

swift/rlhf_trainers/utils.py

@@ -1141,8 +1141,8 @@ def finish_vllm_weight_reload(vllm_model, model_config=None, target_device=None)
     # Prefer vLLM's built-in
     if model_config is not None and target_device is not None:
         try:
-            from vllm.model_executor.model_loader.utils import process_weights_after_loading as _vllm_process
-            _vllm_process(vllm_model, model_config, target_device)
+            from vllm.model_executor.model_loader.utils import process_weights_after_loading
+            process_weights_after_loading(vllm_model, model_config, target_device)
             return
         except Exception as e:
             logger.warning(