fix dist config keys

Author: grimoire

lmdeploy/pytorch/engine/model_agent.py

@@ -729,7 +729,7 @@ def __update_inputs(next_token_ids, model_metas, extra_inputs):
                 inputs,
                 return_logits=return_logits,
                 sync_long_context=sync_long_context,
-                return_routed_experts=return_routed_experts and need_output,
+                return_routed_experts=return_routed_experts and self.need_output,
             )
             logits = output['logits'][0]  # [bs, seq, prob] -> [seq, prob]
             seq_length = output.get('seq_length', inputs.seq_length)
@@ -959,7 +959,7 @@ def _build_model(self):
             update_custom_module_map(custom_module_map)
         logger.debug(msg_with_rank(rank, 'build model.'))
         # for router replay
-        need_output = self.dist_ctx.dp > 1 or self.dist_ctx.rank % self.dist_ctx.tp == 0
+        need_output = self.dist_config.dp > 1 or self.need_output
         enable_return_routed_experts = self.misc_config.enable_return_routed_experts and need_output
 
         build_model_ctx = BuildModelContext(