[Chore] clean up debug info

vllm/distributed/afd_transfer/afd_connector/p2p_connector.py

@@ -60,7 +60,7 @@ class P2PAFDConnector(AFDConnectorBase):
             self.num_hidden_layers: int = (
                 self.config.model_config.hf_config.num_hidden_layers
             )
-            
+
         self.recv_attn_output_counter: int = 0
         self.recv_ffn_output_counter: int = 0
         self.dp_metadata_list: dict[int, DPMetadata] = {}
@@ -139,11 +139,9 @@ class P2PAFDConnector(AFDConnectorBase):
         for idx in range(num_of_stages):
             if idx == 0:
                 tensor_metadata_list[0] = tensor_metadata
-                logger.info(f"build tensor metadata: stage_{idx=}, size={tensor_metadata.size}")
             else:
                 new_size = list(tensor_metadata.size)
                 new_size[0] = connector_metadata.afd_tokens_lens[idx]
-                logger.info(f"build tensor metadata: stage_{idx=}, {new_size=}, {connector_metadata.afd_tokens_lens=}")
                 tensor_metadata_list[idx] = TensorMetadata(
                     tensor_metadata.device,
                     tensor_metadata.dtype,
@@ -167,7 +165,6 @@ class P2PAFDConnector(AFDConnectorBase):
         )
         metadata_tuple = (metadata, tensor_metadata)
         process_group.send_object(metadata_tuple, dst=dst)
-        logger.info(f"_send_metadata called build tensor metadata")
         self._tensor_metadata_list = self._build_tensor_metadata_list(
             tensor_metadata, metadata
         )
@@ -180,23 +177,32 @@ class P2PAFDConnector(AFDConnectorBase):
         (self._current_afd_connector_metadata, tensor_metadata) = (
             process_group.recv_object(src=src)
         )
-        logger.info(f"_recv_metadata called build tensor metadata")
         self._tensor_metadata_list = self._build_tensor_metadata_list(
             tensor_metadata, self._current_afd_connector_metadata
         )
-        logger.info(f"{self.config.parallel_config.data_parallel_size=}")
         if self.config.parallel_config.data_parallel_size > 1:
-            logger.info("jcz recv_metadata num_of_stages:{}".format(self._current_afd_connector_metadata.num_of_stages))
+            logger.info(
+                "jcz recv_metadata num_of_stages:{}".format(
+                    self._current_afd_connector_metadata.num_of_stages
+                )
+            )
             for stage_idx in range(self._current_afd_connector_metadata.num_of_stages):
                 num_tokens_per_ubatch = self._tensor_metadata_list[stage_idx].size[0]
-                logger.info(f"{stage_idx=}, {num_tokens_per_ubatch=}")
                 self.dp_metadata_list[stage_idx] = DPMetadata.make(
                     self.config.parallel_config,
                     num_tokens_per_ubatch,
-                    torch.tensor([num_tokens_per_ubatch] * self.config.parallel_config.data_parallel_size,
-                                device="cpu", dtype=torch.int32),
+                    torch.tensor(
+                        [num_tokens_per_ubatch]
+                        * self.config.parallel_config.data_parallel_size,
+                        device="cpu",
+                        dtype=torch.int32,
+                    ),
                 )
-            logger.info("jcz recv_metadata self.dp_metadata_list:{}".format(self.dp_metadata_list))
+            logger.info(
+                "jcz recv_metadata self.dp_metadata_list:{}".format(
+                    self.dp_metadata_list
+                )
+            )
 
     def _send_hidden_states(
         self,
@@ -229,7 +235,6 @@ class P2PAFDConnector(AFDConnectorBase):
             dtype=tensor_metadata.dtype,
             device=tensor_metadata.device,
         )
-        # logger.info(f"{__file__}: p2p recv hidden states: {hidden_states.shape=}, {tensor_metadata.size=}")
         torch.distributed.recv(
             hidden_states,
             src=process_group.ranks[src],
@@ -267,7 +272,6 @@ class P2PAFDConnector(AFDConnectorBase):
             self.recv_ffn_output_counter
             % self._current_afd_connector_metadata.num_of_stages
         )
-        logger.info(f"{stage_idx=}")
         hidden_states, work_list = self._recv_hidden_states(
             src,
             self.e2a_group,

vllm/model_executor/models/step3_text.py

@@ -303,10 +303,6 @@ class Step3TextDecoderLayer(nn.Module):
         else:
             hidden_states, residual = self.input_layernorm(hidden_states, residual)
 
-        # query, key and positions must have the same number of tokens
-        # /model_executor/layers/rotary_embedding/base.py
-        # positions.shape=torch.Size([8192]), hidden_states.shape=torch.Size([4096, 3712]）
-        logger.info(f"{positions.shape=}, {hidden_states.shape=}")
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
@@ -342,7 +338,6 @@ class Step3TextDecoderLayer(nn.Module):
             hidden_states = share_output + moe_output
         else:
             hidden_states = self.mlp(hidden_states)
-        logger.info(f"{type(hidden_states)=}")
         return hidden_states
 
 
@@ -353,7 +348,6 @@ class Step3TextModel(nn.Module):
         config = vllm_config.model_config.hf_config
         cache_config = vllm_config.cache_config
         quant_config = vllm_config.quant_config
-        logger.info(f"{quant_config=}")
         afd_config = vllm_config.afd_config
         self.vocab_size = config.vocab_size
         self.config = config
@@ -440,7 +434,6 @@ class Step3TextModel(nn.Module):
 
                 ubatch_hidden_states[stage_i] = hidden_states
                 ubatch_residual[stage_i] = residual
-                logger.info(f"create attn metadata:, {afd_metadata.afd_tokens_lens=}")
                 metadata = AFDConnectorMetadata.create_attention_metadata(
                     layer_idx=layer.layer_idx,
                     stage_idx=stage_i,
@@ -515,7 +508,6 @@ class Step3TextModel(nn.Module):
         hidden_states,
         layer_idx,
     ) -> torch.Tensor | IntermediateTensors:
-        logger.info(f"{type(self.layers)=}, {type(layer_idx)=}")
         hidden_states = self.layers[layer_idx].compute_ffn_output(hidden_states)
         return hidden_states
 
@@ -582,7 +574,6 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        logger.info(f"{__file__}: load_weights!")
         qkv_params_mapping = [
             # (param_name, shard_name, relative_start_idx, relative_end_idx)
             (
@@ -615,7 +606,6 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):
             (".gate_up_proj", ".up_proj", 1),
         ]
         params_dict = dict(self.named_parameters())
-        # logger.info(f"{params_dict.keys()=}")
         loaded_params: set[str] = set()
 
         expert_params_mapping = [
@@ -627,10 +617,6 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):
         disable_moe_stacked_params = [data[1] for data in expert_params_mapping]
 
         for name, loaded_weight in weights:
-            # logger.info(
-            #     f"{self.afd_role=}, {name=}, is_moe: {self.is_moe_weight(name)}, "
-            #     f"is_common: {self.is_common_weight(name)}"
-            # )
             if self.afd_role == "attention" and self.is_moe_weight(name):
                 continue
 
@@ -695,7 +681,6 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):
                         start_idx,
                         end_idx,
                     ) in qkv_params_mapping:
-                        # logger.info(f"{weight_name=}, {name=}")
                         if weight_name not in name:
                             continue
                         name = name.replace(weight_name, param_name)