Question about Run-ahead and Race Conditions in Pallas doc

[wenqinI]

Hi team,

I'm learning Pallas with this doc, that's excellent doc for beginner to explore in TPU/pallas world, thank you for it!

When I was going through this section, it said "If we have N devices, this means that the final device can be up to N iterations ahead of the first device.", but for the kernel it discussed, I suppose it 2 at most.

Because the kernel is a double-buffer kernel, and it use the below code snippet:

  # Signal to our left neighbor that we are ready to receive.
  # Without this signal, our left neighbor can be >=1 iteration ahead,
  # meaning it could write into our working slot.
  pl.semaphore_signal(
      capacity_sem,
      inc=1,
      device_id=(left_neighbor,),
      device_id_type=pl.DeviceIdType.MESH,
  )

.......

  # Block until our right neighbor is ready to receive.
  pl.semaphore_wait(capacity_sem, 1)

Which means there is 2 iteration lags at most, because there is only 2 buffers, is my understand correct?

[jburnim]

You are correct that the use of capacity_sem prevents one device from running more than one iteration ahead of its immediate neighbors.

Even with capacity_sem, one device can run ahead farther compared to devices that are farther away in the ring.

For example, suppose N=8 and device 4 pauses forever after signaling its left neighbor on capacity_sem in step 0. Then:

• Devices 0, 1, 2, 3, 6, and 7 complete step 0.
• Devices 0, 1, 2, and 7 complete step 1.
• Devices 0 and 1 complete step 2.
• No device is able to complete step 3.

(I do think we should update the text to say something like "If we have N devices, this means that a device can be up to N/2 iterations ahead of the device farthest from it in the ring.")

[wenqinI]

Got it, thank you for the clarification!