Ring Attention simulation

ring_attention_simulation.py

# Ring Attention using threads to simulate workers
import time
import threading
import pickle

import torch
import torch.nn.functional as F

torch.manual_seed(1337)

# Input data
n_context = 1024
d_head = 2
n_workers = 16
chunk_size = n_context // n_workers
print(f'{n_context=}, {d_head=}, {n_workers=}, {chunk_size=}')

# These will be split (sharded) across the devices
q, k, v = torch.randn(n_context, d_head), torch.randn(n_context, d_head), torch.randn(n_context, d_head)

# For comparison with a correct implementation of Attention
# NOTE: Ignores the scaling factor of scaled dot-product attention
correct = F.softmax(q @ k.transpose(-2, -1), dim=-1) @ v

# Used by the inner loop of Ring Attention to accumulate current K, V against Q into O
# Adapted from: https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf
def flashattention(q, k, v, o, m, d):
  assert q.shape == k.shape == v.shape == o.shape # Q, K, V, O must all have same shape
  assert m.shape == d.shape
  assert q.shape[0] == m.shape[0] == d.shape[0] # mins, denoms must all have shape equal to context length
  for i in range(q.shape[0]): # Iterate over each token in context of K, V and accumulate into O
    x_i = q @ k[i:i+1, :].transpose(-2, -1)
    m_next = torch.maximum(m, x_i)
    d_next = d * torch.exp(m - m_next) + torch.exp(x_i - m_next)
    o_adjust = d * torch.exp(m - m_next) / d_next
    o_add = torch.exp(x_i - m_next) * v[i:i+1, :] / d_next
    o = o * o_adjust + o_add # Accumulate into output
    # Update state of maxes, partial denominators
    m = m_next
    d = d_next
  return o, m, d

# Ring Attention Simulation
class Device(threading.Thread):
  def __init__(self, device_id, q, k, v, barrier, n_workers):
    super().__init__()
    self.device_id = device_id
    self.next_id = (self.device_id + 1) % n_workers
    self.q = q
    self.k = k
    self.v = v
    # Output will be accumulated here using FlashAttention
    self.o = torch.zeros(q.shape)
    # Pad these with dummy dimension `(q.shape[0], 1)` to make broadcasting work for flash attention
    # There is one min/partial denominator per query/output row
    self.mins = torch.full((q.shape[0], 1), float('-inf'))
    self.denoms = torch.zeros((q.shape[0], 1))

    # For receiving messages
    self.barrier = barrier
    self.recv_buf = None

  def run(self):
    for step in range(n_workers):
      # Special case: First step: Process device's own k, v blocks
      k, v = self.k, self.v # k, v can be considered the buffer for received k, v blocks from the previous device
      if step != 0:
        k, v = self.recv()

      # Barrier to ensure all devices/workers have read/received K, V blocks from their buffers
      self.barrier.wait()

      if step != n_workers - 1: # No need to send on last step
        send(self.next_id, self.device_id, (k, v))

      # Barrier to ensure all devices/workers have sent K, V blocks to the next device's buffers
      self.barrier.wait()

      # Local flash attention to update output block
      self.o, self.mins, self.denoms = flashattention(self.q, k, v, self.o, self.mins, self.denoms)

  def recv(self):
    source, ser_data = self.recv_buf
    n_bytes = len(ser_data)
    # Simple simulation of network latency using fixed time plus bandwidth time
    latency = 0.01 + 0.00001 * n_bytes
    time.sleep(latency)
    print(f'Device {self.device_id} received {n_bytes} bytes from device {source} in {latency:.2f} s')
    return pickle.loads(ser_data)

  def __repr__(self):
    return f'Device({self.device_id=}, {self.q=}, {self.k=}, {self.v=}, {self.o=})'

barrier = threading.Barrier(n_workers) # To synchronize receives before sends
devices = [Device(
    i,
    # Each device gets a block/chunk of Q, K, V matrices
    q[i*chunk_size:(i+1)*chunk_size],
    k[i*chunk_size:(i+1)*chunk_size],
    v[i*chunk_size:(i+1)*chunk_size],
    barrier,
    n_workers) for i in range(n_workers)]

def send(dest, source, data):
  ser_data = pickle.dumps(data)
  devices[dest].recv_buf = source, ser_data

start_time = time.time()
for device in devices:
  device.start()
for device in devices:
  device.join()
end_time = time.time()

print(f'Ring Attention simulation time: {end_time - start_time:.2f} seconds')

output = torch.cat([d.o for d in devices])
# NOTE: There appear to be floating-point errors due to the changed order of operations introduced by this RingAttention implementation/simulation
# Hence, the increased tolerance from rtol, atol.
is_correct = torch.allclose(correct, output, rtol=1e-04, atol=1e-06)
print(f'Correct? {is_correct}')
if not is_correct:
  print(f'Correct: {correct}')
  print(f'Output: {output}')