get Token/Score from Paddle inference engine
| 57 | |
| 58 | |
| 59 | class TokenProcessor: |
| 60 | """ |
| 61 | get Token/Score from Paddle inference engine |
| 62 | """ |
| 63 | |
| 64 | def __init__(self, cfg, cached_generated_tokens, engine_worker_queue, split_connector): |
| 65 | paddle.device.set_device("cpu") |
| 66 | self.cfg = cfg |
| 67 | self.cached_generated_tokens = cached_generated_tokens |
| 68 | self.resource_manager = None |
| 69 | self.scheduler_metrics_logger = None |
| 70 | self.engine_worker_queue = engine_worker_queue |
| 71 | self.tokens_counter = Counter() |
| 72 | self.split_connector = split_connector |
| 73 | |
| 74 | if envs.FD_USE_GET_SAVE_OUTPUT_V1: |
| 75 | port = self.cfg.parallel_config.local_engine_worker_queue_port |
| 76 | llm_logger.debug( |
| 77 | f"create zmq get_save_output_rank{self.cfg.parallel_config.local_data_parallel_id}_{port}" |
| 78 | ) |
| 79 | self.zmq_server = ZmqIpcServer( |
| 80 | name=f"get_save_output_rank{self.cfg.parallel_config.local_data_parallel_id}_{port}", mode=zmq.PULL |
| 81 | ) |
| 82 | |
| 83 | self.speculative_decoding = self.cfg.speculative_config.method is not None |
| 84 | self.use_logprobs = self.cfg.model_config.enable_logprob |
| 85 | self.use_sampling_mask = getattr(self.cfg.model_config, "enable_keep_sampling_mask", False) |
| 86 | if not envs.FD_USE_GET_SAVE_OUTPUT_V1 and self.use_sampling_mask: |
| 87 | rank_id = self.cfg.parallel_config.local_data_parallel_id |
| 88 | port = self.cfg.parallel_config.engine_worker_queue_port[rank_id] |
| 89 | self.sampling_mask_zmq_server = ZmqIpcServer( |
| 90 | name=f"sampling_mask_output_rank_{rank_id}_{port}", mode=zmq.PULL |
| 91 | ) |
| 92 | llm_logger.info(f"create zmq sampling_mask_output_rank_{rank_id}_{port}") |
| 93 | self.enable_draft_logprob = self.cfg.speculative_config.enable_draft_logprob |
| 94 | |
| 95 | if self.speculative_decoding: |
| 96 | if self.use_logprobs: |
| 97 | self.output_tokens = paddle.full( |
| 98 | shape=[MAX_BSZ * MAX_DRAFT_TOKENS * (K + 1) + MAX_BSZ + 3, 1], fill_value=2, dtype="int64" |
| 99 | ) |
| 100 | self.output_scores = paddle.full( |
| 101 | shape=[MAX_BSZ * MAX_DRAFT_TOKENS * (K + 1), 1], fill_value=0.0, dtype="float32" |
| 102 | ) |
| 103 | self.output_ranks = paddle.full(shape=[MAX_BSZ * MAX_DRAFT_TOKENS], fill_value=0, dtype="int64") |
| 104 | else: |
| 105 | self.output_tokens = paddle.full( |
| 106 | shape=[SPECULATE_MAX_BSZ * MAX_DRAFT_TOKENS + SPECULATE_MAX_BSZ + 2], |
| 107 | fill_value=2, |
| 108 | dtype="int64", |
| 109 | ) |
| 110 | elif self.use_logprobs: |
| 111 | self.output_tokens = paddle.full(shape=[MAX_BSZ * (K + 1) + 2, 1], fill_value=2, dtype="int64") |
| 112 | self.output_scores = paddle.full(shape=[MAX_BSZ * (K + 1), 1], fill_value=0.0, dtype="float32") |
| 113 | self.output_ranks = paddle.full(shape=[MAX_BSZ], fill_value=0, dtype="int64") |
| 114 | else: |
| 115 | self.output_tokens = paddle.full(shape=[MAX_BSZ + 2, 1], fill_value=2, dtype="int64") |
| 116 | self.worker = None |
no outgoing calls