Using topology-aware scheduling in Amazon SageMaker HyperPod

Data transfer efficiency is a critical factor in high-performance computing (HPC) and machine learning workloads. When using UltraServers with Amazon SageMaker HyperPod, SageMaker HyperPod automatically applies topology labels to your resources. Topology-aware scheduling helps allocate resources to minimize data transfer overheads by considering both instance topology (how resources are connected within an instance) and network topology (how instances are connected with each other). For more information about instance topology, see Amazon EC2 instance topology.

Topology-aware scheduling works with both clusters on Slurm and Amazon EKS. For general information about how topology works with Slurm, see the Topology guide in the Slurm documentation.

In Amazon SageMaker HyperPod, data transfer overheads typically come from three main sources:

UltraServer architecture

SageMaker HyperPod supports UltraServer architecture with p6e-gb200.36xlarge instances. An UltraServer contains up to 18 p6e-gb200.36xlarge instances, with 4 GPUs on each instance. All GPUs across all nodes are interconnected through NVLink switches, enabling data transfer between any two GPUs without using network interfaces.

This architecture provides a significant performance boost compared to individual instances. To leverage this architecture effectively, jobs should be submitted to compute nodes from a single UltraServer.

EKS topology label

In accordance with EC2 instance topology, HyperPod automatically labels your nodes with the following labels:

Using these labels, you can use topology-aware scheduling in HyperPod task governance to apply topology labels and annotations to optimize training efficiency of your workloads. For more information, see Using topology-aware scheduling in Amazon SageMaker HyperPod task governance.

Slurm network topology plugins

Slurm provides built-in plugins for network topology awareness. SageMaker HyperPod automatically selects and configures the appropriate topology plugin based on the instance types in your cluster.

Automatic topology selection

When you create a HyperPod Slurm cluster, the system inspects all instance groups and their associated instance types, identifies the GPU communication characteristics of each instance type, and configures Slurm with the appropriate topology plugin. This process runs automatically and does not require any configuration.

HyperPod manages topology through dynamically generated configuration files. On Slurm 25.11 and later, topology is defined in a topology.yaml file, which is the source of truth and supports multiple topology definitions and per-partition assignment. On Slurm 24.x, topology is defined in a topology.conf file with a single cluster-wide topology. As the cluster evolves through scaling operations or node replacements, HyperPod continuously reconciles the topology configuration to reflect the current cluster state. For more information, see Dynamic topology updates.

Instance types that support network topology

HyperPod configures tree or block topology for instance types that support Amazon EC2 instance topology. For the authoritative, up-to-date list of supported instance types, see Prerequisites for Amazon EC2 topology in the Amazon EC2 User Guide.

Supported instance types include accelerated computing families such as G6e, G7e, P4d, P4de, P5, P5e, P5en, and P6e-GB200, as well as AWS Trainium families such as Trn1, Trn1n, and Trn2. UltraServer instance types (for example, ml.p6e-gb200.36xlarge) use block topology, and other topology-capable instance types use tree topology.

Using the topology/tree plugin

The topology/tree plugin models hierarchical communication structures with multiple bandwidth tiers. Tree topology enables Slurm to place jobs in a way that minimizes cross-tier communication and maximizes locality.

Tree topology is used for instance types with hierarchical interconnects, where distributed training workloads benefit from locality-aware placement. This includes instance types such as ml.p5.48xlarge, ml.p5e.48xlarge, and ml.p5en.48xlarge.

SageMaker HyperPod automatically configures the topology/tree plugin when your cluster uses these instance types. The generated topology configuration maps nodes into a switch hierarchy that reflects the communication tiers of your hardware.

Ensure your slurm.conf includes:

TopologyPlugin=topology/tree

Configuration

SageMaker HyperPod automatically configures tree topology based on information provided by Amazon EC2. For more details about Amazon EC2 topology, see Amazon EC2 instance topology.

On Slurm 25.11 and later, HyperPod defines topology in topology.yaml, which is the source of truth. A tree topology entry maps nodes into a switch hierarchy that reflects the communication tiers of your hardware:

- topology: tree
  cluster_default: true
  tree:
    switches:
      - switch: root
        children: leaf-0,leaf-1
      - switch: leaf-0
        nodes: compute-1,compute-2
      - switch: leaf-1
        nodes: compute-3,compute-4

On Slurm 24.x, tree topology is defined in topology.conf instead, which uses the following format:

SwitchName=nn-6fe9d8a965d34d181 Switches=nn-0b53107754517bf0e

SwitchName=nn-0b53107754517bf0e Switches=nn-424c855d4ad825aa4,nn-95acd7c656329fc30

SwitchName=nn-424c855d4ad825aa4 Nodes=ip-10-1-111-198
SwitchName=nn-95acd7c656329fc30 Nodes=ip-10-1-53-231

Usage

When the topology/tree plugin is configured, Slurm tries to allocate machines that are close to each other. You can force Slurm to allocate machines on a single switch by passing the --switch command line parameter to sbatch or srun:

sbatch --switch=1 ....

Using the topology/block plugin

NVIDIA developed a topology/block plugin that provides hierarchical scheduling across blocks of nodes with the following characteristics:

This plugin allocates nodes based on the defined network topology.

Block topology models uniform, high-bandwidth communication domains where all GPUs participate in a single high-speed domain with near-uniform latency. Block topology treats all nodes as part of a single cohesive communication unit. UltraServer architecture in SageMaker HyperPod supports the block plugin.

Block topology is used for UltraServer instance types such as ml.p6e-gb200.36xlarge.

Ensure your slurm.conf includes:

TopologyPlugin=topology/block

Configuration

SageMaker HyperPod automatically configures block topology. On Slurm 25.11 and later, block topology is defined in topology.yaml, which is the source of truth. Block and tree topologies for a cluster are defined together in the same file. The following example shows a cluster with a P5 partition (tree) and an UltraServer partition (block):

- topology: tree
  cluster_default: true
  tree:
    switches:
      - switch: root
        children: leaf-0,leaf-1
      - switch: leaf-0
        nodes: compute-p5-1,compute-p5-2
      - switch: leaf-1
        nodes: compute-p5-3,compute-p5-4
- topology: block
  cluster_default: false
  block:
    block_sizes:
      - 18
    blocks:
      - block: cb-001
        nodes: ultraserver-1-[1-18]

On Slurm 24.x, block topology is defined in topology.conf instead, which uses the following format:

BlockName=us1 Nodes=ultraserver1-[0-17]
BlockName=us2 Nodes=ultraserver2-[0-17]
BlockSizes=18

Usage

When submitting jobs, you can use the following additional arguments with sbatch and srun commands:

The following are sample scenarios you might consider when thinking about allocating blocks.

Allocate a whole block of nodes on an empty system

sbatch -N18

Allocate two blocks of nodes on an empty system

sbatch -N36

Allocate 18 nodes on one block + 6 nodes on another block

sbatch -N24

Allocate 12 nodes on one block and 12 nodes on another block

sbatch -N24 --segment=12

With --exclusive=topo, job must be placed on block with no other jobs

sbatch -N12 --exclusive=topo

Partition-level topology selection

Starting with Slurm 25.11, HyperPod supports topology configuration at the partition level. Each partition is assigned a topology based on the instance types of its compute instance groups, so a single cluster can run tree topology in one partition and block topology in another. Partitions that contain instance types without network topology support inherit a cluster-wide flat default, which keeps their nodes schedulable.

HyperPod resolves the topology for each partition as follows:

The cluster-wide default topology is selected based on the topologies present across the cluster: if only one topology type is present, that type is the default; if both block and tree are present with no non-topology groups, tree is the default; and if any non-topology compute group is present, flat is the default so that those nodes remain schedulable.

The following table shows an example of how HyperPod resolves topology for a cluster with mixed instance types.

In this example, the P5 partition uses tree topology and the UltraServer partition uses block topology. With partition-level topology, each partition uses its optimal topology within the same cluster, so you no longer need to create separate clusters to give each instance type its ideal topology model.

Flat topology default

When a cluster contains a mix of topology-capable and non-topology instance types, HyperPod applies topology/flat as the cluster default. Topology-aware partitions are pointed at their tree or block topology through per-partition Topology= directives in slurm.conf (Slurm 25.11 and later), while partitions with non-topology nodes have no Topology= directive and inherit the flat default. This ensures every node remains schedulable.

Disable or change topology plugin

When a Slurm cluster is created, HyperPod automatically selects the optimal topology plugin. To manually change the topology plugin, update the TopologyPlugin value in slurm.conf on the controller node.

To disable topology-aware placement, set the plugin to topology/flat (or topology/default):

# Set this value to disable topology-aware placement
TopologyPlugin=topology/flat

Dynamic topology updates

Topology-aware scheduling continuously maintains topology correctness as your cluster changes. The topology is automatically recalculated and the topology configuration file is regenerated when any of the following events occur:

When the topology is updated, new nodes are incorporated into the correct topology structure, removed nodes are pruned, and the Slurm configuration is updated without requiring manual intervention. This ensures that the topology always reflects the actual cluster state.

Best practices for UltraServer topology

For optimal performance with UltraServer architecture in SageMaker HyperPod:

When thinking about segmenting, consider the following:

Limitations in SageMaker HyperPod topology aware scheduling

With Slurm 25.11 and later, heterogeneous clusters (clusters with different instance types) are supported through partition-level topology and the flat cluster default. Each partition receives the topology best suited to its instance types, and non-topology and UltraServer nodes remain schedulable within the same cluster. For more information, see Partition-level topology selection.

On clusters running Slurm 24.x, a single cluster-wide topology applies, and the topology/block plugin has the following limitations with heterogeneous clusters:

For heterogeneous clusters on Slurm 24.x, consider these alternatives: