Your LM head is too big for one device. How do you shard it by vocabulary or by hidden width, combine the logits, and compute an exact top-k across shards?

Vocabulary sharding: each device holds a slice of rows and computes logits for its slice only. Softmax never needs the full vector: all-reduce the max, then the sum of exponentials. Hidden sharding: each device holds a slice of the hidden width and produces partial logits over the whole vocabulary, which must be all-reduced, a vocabulary-sized vector per token.

import torch
import torch.distributed as dist

def dist_topk(local, k, offset):
    # local: [B, V_shard] logits on this rank
    v, i = local.topk(k, dim=-1)
    i = i + offset  # to global vocabulary ids
    n = dist.get_world_size()
    vs = [torch.empty_like(v) for _ in range(n)]
    ix = [torch.empty_like(i) for _ in range(n)]
    dist.all_gather(vs, v)
    dist.all_gather(ix, i)
    v, i = torch.cat(vs, -1), torch.cat(ix, -1)
    top, pos = v.topk(k, dim=-1)
    return top, i.gather(-1, pos)

It's exact: every global top-k entry is in its own shard's top k.