fix rollout

Author: Jintao-Huang

swift/llm/infer/rollout.py

@@ -86,6 +86,7 @@ def get_rollout_engine_type(args: RolloutArguments, engine: GRPOVllmEngine):
 def llm_worker(args: RolloutArguments, data_parallel_rank: int, master_port: int, connection: Connection) -> None:
     # Set required environment variables for DP to work with vLLM
     args._import_external_plugins()
+    args._init_custom_register()
     os.environ['VLLM_DP_RANK'] = str(data_parallel_rank)
     os.environ['VLLM_DP_RANK_LOCAL'] = str(data_parallel_rank)
     os.environ['VLLM_DP_SIZE'] = str(args.vllm_data_parallel_size)
@@ -119,6 +120,7 @@ async def async_llm_worker(args: RolloutArguments, data_parallel_rank: int, mast
                            connection: Connection) -> None:
     # Set required environment variables for DP to work with vLLM
     args._import_external_plugins()
+    args._init_custom_register()
     engine = SwiftRolloutDeploy.get_infer_engine(args, template=args.get_template(None))
 
     rollout_engine = get_rollout_engine_type(args, engine)