| # tpu.yaml - TPU accelerator families available on this cluster. | |
| # | |
| # Read by job_submit/tpu, which validates --layout requests against it, and by | |
| # gres/tpu, which uses vendor_id and pci_device_id to detect chips on a node. | |
| # | |
| # A family is described by two chip geometries rather than by chip and node | |
| # counts. Counts cannot express a shape -- on tpu7x below, 256 chips is 64 | |
| # nodes, which could be 4x4x16 or 4x8x8 -- and the shapes are what the geometry | |
| # rules are stated in terms of. chips_per_node, the fabric's dimension count | |
| # and nodes_per_segment are all derived from them. | |
| # | |
| # Install as $SYSCONFDIR/tpu.yaml (typically /etc/slurm/tpu.yaml). | |
| # See tpu.yaml(5). | |
| tpu_types: | |
| # 3D torus. A node carries 4 chips arranged 2x2x1, so a 2x2x4 request is | |
| # 16 chips on 4 nodes. | |
| - type: tpu7x | |
| shape_per_node: "2x2x1" | |
| # One segment (the scale beyond which an optical circuit switch is used), | |
| # 64 chips on 16 nodes. Requests up to this size fit within a single | |
| # segment, each of their axes must divide 4, and their longest axis can | |
| # be at most twice their shortest. | |
| shape_per_segment: "4x4x4" | |
| # Pod limit. Requests larger than one segment and up to this are | |
| # satisfied by stitching whole segments together over OCS, so each of | |
| # their axes must be a multiple of 4. | |
| nodes_per_pod: 2304 | |
| vendor_id: "0x1ae0" | |
| pci_device_id: "0x0076" | |
| # 2D mesh. A node carries 4 chips arranged 2x2, so a 4x4 request is | |
| # 16 chips on 4 nodes. The whole pod is treated as a single segment, | |
| # so every request must fit inside 16x16, with its longest axis at most | |
| # twice its shortest. | |
| - type: v6e | |
| shape_per_node: "2x2" | |
| shape_per_segment: "16x16" | |
| nodes_per_pod: 64 | |
| vendor_id: "0x1ae0" | |
| pci_device_id: "0x006f" |