Function _embedding_bwd_kernel

tinygrad/nn/__init__.py:333–362

(grad_weight:UOp, grad_emb:UOp, idx:UOp)

  # this is the real atomic kernel
  def _embedding_bwd_kernel(grad_weight:UOp, grad_emb:UOp, idx:UOp) -> UOp:
    idx_flat, grad_emb_flat = idx.flatten(), grad_emb.reshape((idx.numel(), grad_weight.shape[-1]))

    embed_size = grad_weight.shape[-1]
    BLOCK_J = min(256, embed_size)
    assert embed_size % BLOCK_J == 0, f"embed_size {embed_size} must be divisible by {BLOCK_J}"

    n_j_blocks = embed_size // BLOCK_J
    i = UOp.range(grad_emb_flat.shape[0], 0)         # batch_size * sequence_length -> GLOBAL
    j_inner = UOp.range(BLOCK_J, 2, AxisType.LOOP if device in ("CPU", "NULL") else AxisType.LOCAL)  # BLOCK_J threads per workgroup
    j_outer = UOp.range(n_j_blocks, 1)
    j = j_outer * BLOCK_J + j_inner

    if is_vocab_sharded:
      # each device owns [offset, offset+local_vocab_size) of the global vocabulary
      dnum = UOp.variable("_device_num", 0, ndev-1)
      offset = dnum * local_vocab_size
      global_token_id = idx_flat[i].cast(dtypes.weakint)
      local_token_id = (global_token_id - offset).clip(0, grad_weight.shape[0]-1)
      in_range = (global_token_id >= offset) & (global_token_id < (offset + local_vocab_size))
      grad_val = in_range.where(grad_emb_flat[i, j].cast(dtypes.float), 0.0)
    else:
      local_token_id = idx_flat[i].clip(0, grad_weight.shape[0]-1).cast(dtypes.weakint)
      grad_val = grad_emb_flat[i, j].cast(dtypes.float)
    # atomic scatter-add: grad_weight[token_id, j] += grad_emb_flat[i, j]
    if device in ("CPU", "NULL"): atomic_arg = "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
    elif device == "AMD": atomic_arg = "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
    else: raise NotImplementedError(f"no atomics for device {device}")
    atomic = UOp(Ops.CUSTOM, dtypes.void, (grad_weight.index(local_token_id, j, ptr=True), grad_val), arg = atomic_arg)
    return atomic.end(i, j_outer, j_inner).sink(arg=KernelInfo(name="embedding_bwd", opts_to_apply=()))

  grad_weight_uop = grad_weight_uop.custom_kernel(grad_emb, idx, fxn=_embedding_bwd_kernel)[0]