Requirements from Control and Management Viewpoint for Collective Communication Optimization

2.1. Memory Management

Scarce memory resources provided by network devices for collective communication MUST be scheduled and controlled, e.g. assigning a scheduling priority to collective communication offloading tasks. Compared to the amount of collective communication message in the applications such as AI, HPC, etc., it is severely mismatched and extremely scarce for memory resource provided by network devices for collective communication, such as network programmable switches.¶

Use Case[ESA]. The memory of programmable switch is scarce for the amount of gradient transmitted in distributed training. There is some existing work to solve this problem like pool-based streaming and dynamic sharing, which are not enough yet. A use case of fully utilizing the memory of programmable switch is that the control and management module of switch assigns a priority to the aggregation task, to dynamically and preemptively schedule the aggregation tasks in the data plane, thus making more full use of memory in the form of switch aggregators.¶

+----------+  +----------+           +----------+
|          |  |          |           |          |
| worker 1 |  | worker 2 |           | worker n |
|          |  |          |  ... ...  |          |
+----+-----+  +----+-----+           +-----+----+
     |             |                       |
     |             |                       |
     +------+------+-----------------------+
            |GB/TB-level gradients
+-----------+-----------+           +-----------+
|           |           |           |           |
|  +--------+--------+  |           |           |
|  |     Switch      |  |           |  Control  |
|  |   Aggregators   |  |  Manage   |     &     |
|  +-----------------<--+-----------+ Management|
|  |Memory for others|  | Schedule  |           |
|  +-----------------+  |           |           |
| Switch Memory = 10MB  |           |           |
+-----------------------+           +-----------+

2.2. Topology Management

Topology awareness and mapping work are REQUIRED to be done to put some of the end-host computing on the network nodes for collective communication optimization. In many collective operation tasks, the logical relationship between nodes is usually described in the form of graph, and then mapping to the physical network. Therefore, collective communication offloading requires awareness of the network topology and making efficient mappings.¶

Use Case. In the parameter server architecture commonly used in distributed training, the parameter server can be reasonably mapped to spine switches in the fat tree physical network with being aware of network topology. Under this mapping mechanism, the traffic path is more simplified and the traffic volume of the whole network is greatly compressed. Compared to the traditional collective communication mode, the optimized end-to-network or end-to-network-to-end one with topology awareness and mapping makes the physical topology and the logical topology closer, more friendly and unified.¶

                 Logical Topology
                +----------------+
                |Parameter Server|
                +--------+-------+
                         |
          +----------+---+------------+
          |          |                |
     +----+---+ +----+---+        +---+----+
     |Worker 1| |Worker 2| ...... |Worker n|
     +--------+ +--------+        +--------+
                         |Mapping
                         |
              +----------+---------+
              |Management & Control|
              |                    |
              | Topology Awareness |
              | Paths Planning     |
              +----------+---------+
                         |
                         |Mapping
                         v
                 Physical Topology
             +-----+         +-----+
             |Spine|         |Spine|
             +--+--+         +--+--+
                |               |
     +----------+-+------------++-----------+
     |            |            |            |
  +--+--+      +--+--+      +--+--+      +--+--+
  |Leaf |      |Leaf |      |Leaf |      |Leaf |
  +--+--+      +--+--+      +--+--+      +--+--+
     |            |            |            |
  +--+--+      +--+--+      +--+--+      +--+--+
  |     |      |     |      |     |      |     |
+-+-+ +-+-+  +-+-+ +-+-+  +-+-+ +-+-+  +-+-+ +-+-+
|GPU| |GPU|  |GPU| |GPU|  |GPU| |GPU|  |GPU| |GPU|
+---+ +---+  +---+ +---+  +---+ +---+  +---+ +---+

2.3. Interfaces Management

Some collective communication interfaces MUST be defined and managed for application developers to shield tedious network engineering details, such as flow control, packet organization, chip-specific programming language, etc. If not, applications developers will need too much arcane knowledge and expertise, which is beyond their willingness and prevent from the evolution of the emerging applications.¶

Use case. The industry and academy have actually proposed some abstractions of collective communication operations, such as collective communication libraries MPI, NCCL, NetRPC[NetRPC], etc. In the control plane, these interfaces need to be configured and instantiated to complete the part of collective communication functionality.¶

2.4. Data Management

The semantic gap between application data unit and network data unit is REQUIRED to be bridged. This semantic mismatch poses a huge obstacle to the support of upper layer applications in the underlying network.¶

Use Case. In the distributed training of LLMs, AllReduce, a common kind of collective communication operation, involves the transportation of large amount of message-level gradients or token data, which is orders of magnitude of packets in network in size. Because packet in network is designed originally not considering large data transmission, there is natural mismatch property between network and collective communication.¶

2.5. Computional Resources Management

The fusing of communication and computation operators is now a popular and consensus-based way to optimize LLMs training or HPC. This practice can make the scheduling of communication and computing more efficient. A universal and direct way is to consider computing when designing communication operators in the communication libraries. In this process, unified management of computational resources, e.g., AI-cores of GPU or CPU-cores, is REQUIRED, to optimize the scalability of computing cluster.¶

Use Case. Training large language models (LLMs) efficiently in a distributed setting is indeed a challenging task, primarily due to communication bottlenecks that arise when multiple nodes need to synchronize their data. Overlapping communication with computation, or said kernel fusion, is a key strategy to improve training efficiency. This approach aims to reduce idle times by initiating communication operators while the CPU or GPU is still performing computations, e.g., fusing Allgather and Batch-Matmul into a single larger kernel to reduce kernel launch overhead and better overlap communication with computation.¶

5.1. Normative References

[RFC2119]

Bradner, S., "Key words for use in RFCs to Indicate Requirement Levels", BCP 14, RFC 2119, DOI 10.17487/RFC2119, March 1997, <https://www.rfc-editor.org/info/rfc2119>.

5.2. Informative References

[ESA]

Wang, H., "Efficient Data-Plane Memory Scheduling for In-Network Aggregation", 2022.

[NetRPC]

Zhao, B., "NetRPC: Enabling In-Network Computation in Remote Procedure Calls", 2023.