diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
index 137c2e710..1e729332f 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
@@ -132,11 +132,14 @@ public class TransportMessage implements Serializable {
   }
 
   public ByteBuffer toByteBuffer() {
-    int totalBufferSize = payload.length + 4 + 4;
+    int payloadLength = payload != null ? payload.length : 0;
+    int totalBufferSize = payloadLength + 4 + 4;
     ByteBuffer buffer = ByteBuffer.allocate(totalBufferSize);
     buffer.putInt(messageTypeValue);
-    buffer.putInt(payload.length);
-    buffer.put(payload);
+    buffer.putInt(payloadLength);
+    if (payload != null) {
+      buffer.put(payload);
+    }
     buffer.flip();
     return buffer;
   }